DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

技術概要：DAWM – 行動推論による遷移を用いたオフライン強化学習のための拡散行動世界モデル

1. 問題定義

拡散ベースの世界モデルは、オフライン強化学習（RL）のための現実的な長期軌道の合成において強力な能力を示してきた。しかし、既存の手法は、1 ステップ時間的差分（TD）学習に依存する標準的な値ベースのオフライン RL アルゴリズム（例：TD3BC、IQL）との間に重大な互換性の問題を抱えている。

• 行動非依存モデルの限界: Diffusion World Model (DWM) などの手法は、将来の状態と報酬を生成するが、行動を明示的に生成しない。その結果、完全な遷移 $(s, a, r, s')$ を構築できず、状態、行動、報酬の一貫性を必要とする 1 ステップ TD 学習（効果的なクリティック正則化に不可欠）との互換性が欠如している。
• 結合モデルの限界: 状態、報酬、行動を同時にモデル化しようとするアプローチは、しばしば訓練の複雑性と不安定性の増大に陥り、実用的な性能の低下を招く。
• 計画ベースの手法の限界: 逆動力学モデル（IDM）や結合モデルを介して行動を復元する拡散ベースのプランナー（例：Diffuser、Decision Diffuser）は、通常、反復サンプリングと事後の報酬選択に依存する。これにより推論コストが高騰し、1 ステップ TD 学習の効率性が迂回されてしまう。

核心的な課題は、結合モデルの訓練不安定性や反復計画の推論コストを招くことなく、安定かつ効率的な 1 ステップ TD 学習を支援する合成データを生成することである。

2. 手法：DAWM

著者は、軌道生成と行動推論を分離し、オフライン RL に適した完全なラベル付き遷移を生成するモジュール型フレームワーク**DAWM（Diffusion Action World Model）**を提案する。このフレームワークは、主に 2 つのフェーズで構成される。

フェーズ 1：条件付き拡散軌道生成

• 拡散モデル $p_\theta$ は、現在の状態 $s_t$、行動 $a_t$、および目標とする**到達報酬（return-to-go, $R_{tg}$）**に条件付けられた、将来の状態と報酬のシーケンスを生成するように訓練される。
• このモデルは、条件付き分布 $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$ を定義する。
• 訓練目的は、クラスターフリーガイダンスを用いたスコアリングノイズ除去目的関数により、予測ノイズと真のノイズ間の期待誤差を最小化することである。
• 重要なのは、このモデルが行動非依存の軌道（状態と報酬のみ）を生成することであり、拡散プロセス内で行動を結合モデル化する複雑さを回避することで、生成タスクを単純化している点である。

フェーズ 2：逆動力学モデル（IDM）による行動補完

• 行動非依存の軌道を完全な遷移に変換するために、別途訓練された逆動力学モデル（IDM） $f_\phi$ が用いられる。
• IDM は、対角ガウス分布として分布 $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ をモデル化することで、欠落している行動を再構成する。
• IDM は、状態遷移が与えられた際の観測された行動の尤度を最大化するように、実世界のオフラインデータで訓練される。
• 推論時には、拡散モデルで生成された軌道に IDM を適用して行動を推論し、完全なタプル $(s, a, r, s')$ を得る。

オフライン方策訓練

• 合成された完全な遷移は、標準的な値ベースのアルゴリズム、具体的には TD3 に行動模倣正則化を追加したTD3BC、および優位性重み付き教師あり回帰を行うIQLを用いて、オフライン RL エージェントの訓練に使用される。
• Q 関数は、再構成された遷移から計算されたベルマンターゲットを用いて更新され、効率的な 1 ステップ TD 更新を可能にする。

3. 主要な貢献

1. モジュール型フレームワーク: 著者は、拡散ベースの状態・報酬生成と、行動推論のための別途訓練された IDM を組み合わせた効率的なフレームワークを提案する。この設計は、結合モデルの訓練不安定性を回避しつつ、完全な遷移を生成する。
2. 1 ステップ TD 学習との互換性: 完全な $(s, a, r, s')$ 遷移を生成することで、DAWM は 1 ステップ TD 学習に依存する保守的なオフライン RL アルゴリズム（TD3BC、IQL）の利用を可能にし、反復拡散プランナーに対する計算効率の高い代替手段を提供する。
3. 実証的パフォーマンス: 本論文は、保守的なオフライン RL アルゴリズムがこれらの拡張軌道での訓練から大幅に利益を得ており、複数の D4RL 移動タスクにおいて、既存の拡散ベースのベースライン（DWM、DD）を一貫して上回ることを示している。
4. 合成データの妥当性: 結果は、DAWM によって生成された合成データで訓練されたエージェントが、実世界のオフラインデータセットで訓練されたエージェントと同等、あるいはそれ以上の性能を達成しうることを示しており、合成軌道が実用的な代替手段となりうる可能性を浮き彫りにしている。

4. 実験結果

実験は、データ品質（Medium、Medium-Replay、Medium-Expert）が異なる 9 つの標準的な D4RL 移動ベンチマーク（Hopper、Walker2d、HalfCheetah）で実施された。

• DWM との比較: DAWM は、行動非依存の DWM ベースラインを一貫して上回った。TD3BC を用いた場合、DAWM は平均**9.3%の改善を達成し、IQL を用いた場合は9.5%**の改善を達成した。特に Walker2d などの複雑な制御タスクにおいて、その改善は顕著であった。
• 計画ベースの手法（DD）との比較: DAWM は、正規化された報酬の点で Decision Diffuser（DD）を上回った。さらに、推論中の反復サンプリングを回避するため、DAWM は DD に比べて約4.0 倍高速な推論を実現し、著しく高い効率性を示した。
• 実データとの比較: DAWM 生成データで訓練されたエージェントは、標準的な TD3BC/IQL を用いて実データで訓練されたエージェントと同等の正規化報酬を達成し、場合によっては（例：TD3BC を用いた Hopper-MR）、それらを上回った。SSORL（実データ上で IDM を用いたベースライン）は良好な性能を示したが、著者は SSORL が最長の軌道にフィルタリングする恩恵を受けていることを指摘しており、これは DAWM の固定長生成には存在しない要因である。
• アブレーション研究:
  • データ量（DAWM-T）: 条件付けペアの数を 8 分の 1 に削減（DAWM-T）しても、性能の低下は modest（軽微）であった。これは、行動補完メカニズムが単なるデータ量よりも重要な要因であることを示唆している。
  • 生成ホライズン: DAWM は、異なる生成ホライズン（$H=1, 3, 7$）において堅牢な性能を維持しており、ロールアウト長に関わらず 1 ステップ TD 学習フレームワークが安定していることを示している。

5. 意義と主張

本論文は、DAWM が効率的な生成モデルと TD ベースの方策学習の間のギャップを埋めると主張している。その主な意義は以下の点にある。

• 1 ステップ TD 学習の実現: 以前は行動非依存の世界モデルと互換性がなかった、標準的かつ安定したオフライン RL アルゴリズムに対して、拡散生成データを利用するメカニズムを提供する。
• 計算効率性: 反復計画手法に比べて推論コストを削減することで、大規模な生成オフライン RL に対する実用的な解決策を提供する。
• 実用的な代替手段: 結果は、このモジュール型アプローチで生成された合成軌道が、実世界のオフラインデータセットに対する実用的かつ高品質な代替手段となりうることを示唆しており、探索が高価または危険な実世界アプリケーションにおけるデータ不足の問題を緩和する可能性がある。

著者は、この手法が有望であることを示しつつも、将来的な研究として、視覚制御への拡張、高次元ロボティクスへの適用、および不完全なダイナミクス下での堅牢性の検討を行うと結論付けている。