2018年に、 強化学習 ニューラル ネットワークの世界モデルのコンテキストで初めて導入され、すぐにこの基本原理が世界モデルに適用されました。強化学習を実装する著名なモデルには、再帰状態空間モデルの潜在空間からの強化学習を導入した Dreamer フレームワークがあります。DreamerV2 は、離散潜在空間の使用により複合エラーが減少する可能性があることを実証し、DreamerV3 フレームワークは、固定ハイパーパラメータを使用して、さまざまなドメインにわたる一連のタスクで人間のようなパフォーマンスを実現できました。 

さらに、画像生成モデルと世界モデルの間には類似点があり、生成視覚モデルの進歩が世界モデルにも応用できることを示しています。 自然言語処理 フレームワークの人気が高まると、DALL-E フレームワークと VQGAN フレームワークが登場しました。これらのフレームワークは、画像を離散トークンに変換する離散オートエンコーダを実装し、自己回帰トランスフォーマーのシーケンスモデリング機能を活用して、非常に強力で効率的なテキストから画像への生成モデルを構築することができました。同時に、拡散モデルが注目を集め、今日では、拡散モデルは高解像度画像生成の主要なパラダイムとしての地位を確立しています。拡散モデルと強化学習が提供する機能により、2 つのアプローチを組み合わせる試みが行われており、その目的は、軌道モデル、報酬モデル、プランナー、およびオフライン強化学習におけるデータ拡張のポリシーとして拡散モデルの柔軟性を活用することです。 

ワールド モデルは、強化学習エージェントを安全かつ効率的にトレーニングするための有望な方法を提供します。従来、これらのモデルは、離散潜在変数のシーケンスを使用して環境ダイナミクスをシミュレートします。ただし、この圧縮では、強化学習に不可欠な視覚的な詳細が見落とされる可能性があります。同時に、拡散モデルが画像生成で人気が高まり、離散潜在変数を使用する従来の方法に挑戦しています。この変化に触発されて、この記事では、拡散ワールド モデル内でトレーニングされた強化学習エージェントである DIAMOND (DIffusion As a Model Of eNvironment Dreams) について説明します。拡散をワールド モデリングに適したものにするために必要な設計上の選択を検討し、視覚的な詳細を強化するとエージェントのパフォーマンスが向上することを示します。DIAMOND は、競争力のある Atari 100k テストで新しいベンチマークを設定し、平均人間正規化スコア 1.46 を達成しました。これは、ワールド モデル内で完全にトレーニングされたエージェントの中で最高のスコアです。 

ダイヤモンド：環境の夢のモデルとしての普及

環境の世界モデルまたは生成モデルは、生成エージェントが環境を計画し、推論するためのより重要なコンポーネントの 1 つとして浮上しています。強化学習の使用は近年かなりの成功を収めていますが、強化学習を実装するモデルはサンプル効率が悪いことで知られており、実際のアプリケーションが大幅に制限されています。一方、世界モデルは、サンプル効率が大幅に向上し、モデルが実際の経験から学習できるようにすることで、さまざまな環境で強化学習エージェントを効率的にトレーニングできることが実証されています。最近の世界モデリング フレームワークでは、通常、環境ダイナミクスを離散的な潜在変数のシーケンスとしてモデル化し、モデルは潜在空間を離散化して、複数ステップの時間範囲にわたるエラーの複合を回避します。このアプローチは大きな結果をもたらす可能性がありますが、情報の損失も伴い、再構築の品質の低下と一般性の損失につながります。情報の損失は、自律走行車のトレーニングなど、情報を明確に定義する必要がある実際のシナリオにとって大きな障害になる可能性があります。このようなタスクでは、信号機の色や前方の車両の方向指示器などの視覚入力の小さな変化や詳細によって、エージェントのポリシーが変わる可能性があります。離散潜在変数の数を増やすと情報損失を回避できますが、計算コストが大幅に増加します。 

さらに、近年、拡散モデルは高品質画像生成フレームワークの主要なアプローチとして浮上しています。拡散モデルに基づいて構築されたフレームワークはノイズ処理を逆転させることを学習し、離散トークンをモデル化するより確立されたアプローチのいくつかと直接競合するため、世界モデリングで離散化の必要性を排除する有望な代替手段を提供します。拡散モデルは、簡単に条件付けでき、モード崩壊なしで複雑なマルチモーダル分布を柔軟にモデル化できることで知られています。これらの属性は世界モデリングにとって重要です。条件付けにより世界モデルがエージェントの行動を正確に反映できるようになり、より信頼性の高いクレジット割り当てにつながるためです。さらに、マルチモーダル分布をモデル化すると、エージェントのトレーニングシナリオの多様性が向上し、全体的なパフォーマンスが向上します。 

これらの特性に基づいて、DIAMOND (DIffusion As a Model Of eNvironment Dreams) は、拡散世界モデル内でトレーニングされた強化学習エージェントです。DIAMOND フレームワークは、拡散世界モデルが長期間にわたって効率的かつ安定的であり続けるように、慎重な設計選択を行います。フレームワークは、これらの設計選択の重要性を示す定性分析を提供します。DIAMOND は、定評のある Atari 1.46k ベンチマークで平均人間正規化スコア 100 という新たな最先端技術を打ち立てました。これは、完全に世界モデル内でトレーニングされたエージェントの中で最高です。画像空間で動作することで、DIAMOND の拡散世界モデルは環境をシームレスに置き換えることができるため、世界モデルとエージェントの動作に関するより深い洞察が得られます。特に、特定のゲームでのパフォーマンスの向上は、重要な視覚的詳細のモデリングが改善されたことによるものです。DIAMOND フレームワークは、環境を、一連の状態、一連の離散アクション、一連の画像観察を含む標準 POMDP (部分観測マルコフ決定プロセス) としてモデル化します。遷移関数は環境のダイナミクスを記述し、報酬関数は遷移をスカラー報酬にマッピングします。観測関数は観測確率を記述し、画像観測を出力します。エージェントは状態に直接アクセスできないため、この画像観測を使用して環境を確認します。このアプローチの主な目的は、割引係数を使用して期待される割引リターンを最大化するように観測をアクションにマッピングするポリシーを取得することでした。世界モデルは環境の生成モデルであり、世界モデルを使用してシミュレートされた環境を作成し、実際の環境で強化学習エージェントをトレーニングしたり、世界モデル環境で強化学習エージェントをトレーニングしたりできます。図 1 は、DIAMOND フレームワークの想像力が時間の経過とともにどのように展開されるかを示しています。 

DIAMOND : 方法論とアーキテクチャ

本質的に、拡散モデルは、ノイズ処理を逆にすることでサンプルを生成する生成モデルの一種であり、非平衡熱力学から多大なインスピレーションを得ています。DIAMOND フレームワークは、扱いやすい非構造化事前分布を持つ、対応する周辺条件と境界条件を持つ連続時間変数でインデックス付けされた拡散プロセスを考慮します。さらに、ノイズからデータにマッピングする生成モデルを取得するには、DIAMOND フレームワークでプロセスを逆にする必要があり、その逆プロセスも拡散プロセスであり、時間を逆方向に実行します。さらに、DIAMOND フレームワークは真のスコア関数にアクセスしないため、任意の時点でスコア関数を推定することは簡単ではありません。このモデルは、スコア マッチング目標を実装することでこのハードルを克服します。このアプローチにより、フレームワークは、基礎となるスコア関数を知らなくてもスコア モデルをトレーニングできます。スコアベースの拡散モデルは、無条件生成モデルを提供します。ただし、環境ダイナミクスの条件付き生成モデルは、世界モデルとして機能する必要があり、この目的を果たすために、DIAMOND フレームワークは POMDP アプローチの一般的なケースを検討します。このアプローチでは、フレームワークは過去の観測とアクションを利用して、未知のマルコフ状態を近似できます。図 1 に示すように、DIAMOND フレームワークはこの履歴を利用して拡散モデルを調整し、次の観測を直接推定して生成します。DIAMOND フレームワークは、理論上は任意の SDE または ODE ソルバーを使用できますが、NFE または関数評価の数と、拡散モデルの推論コストに大きく影響するサンプル品質の間にはトレードオフがあります。 

上記の学習に基づいて、特定の拡散アプローチの選択に対応するドリフト係数と拡散係数を含む拡散ベースの世界モデルの DIAMOND フレームワークの実用的な実現について見てみましょう。タスクに自然に適した候補である DDPM を選択する代わりに、DIAMOND フレームワークは EDM 定式化に基づいて構築され、ノイズ スケジュールと呼ばれる拡散時間の実数値関数を持つ摂動カーネルを考慮します。フレームワークは、あらゆる音声レベルの入力と出力の分散を維持するために、前処理を選択します。ネットワーク トレーニングでは、劣化レベルに応じて信号とノイズを適応的に混合し、ノイズが低い場合は、ターゲットはクリーンな信号と摂動された信号の差、つまり追加されたガウス ノイズになります。直感的に、これにより、低ノイズ領域でトレーニングの目的が自明になることがなくなります。実際には、この目標はノイズ スケジュールの両端での高い分散であるため、モデルは、中程度のノイズ領域を中心にトレーニングを連結するために、経験的に選択された対数正規分布からノイズ レベルをサンプリングします。DIAMOND フレームワークは、ベクトル フィールドに標準の U-Net 2D コンポーネントを使用し、フレームワークが自身を調整するために使用する過去の観測とアクションのバッファーを保持します。次に、DIAMOND フレームワークは、これらの過去の観測を次のノイズの多い観測に連結し、U-Net の残差ブロックの適応型グループ正規化レイヤーを通じてアクションを入力します。 

ダイヤモンド：実験と結果

包括的な評価のために、DIAMOND フレームワークは Atari 100k ベンチマークを選択します。Atari 100k ベンチマークは、幅広いエージェント機能をテストするように設計された 26 のゲームで構成されています。各ゲームでは、エージェントは環境内で 100 万回のアクションに制限されています。これは、評価前にゲームを学習するため、人間のゲームプレイの 2 時間にほぼ相当します。比較すると、制約のない Atari エージェントは通常 50 万ステップのトレーニングを行い、経験が 500 倍増加します。ゲームごとに 5 つのランダム シードを使用して、DIAMOND を最初からトレーニングしました。各トレーニング実行には約 12 GB の VRAM が必要で、単一の Nvidia RTX 2.9 で約 4090 日かかり、合計で 1.03 GPU 年に相当します。次の表は、すべてのゲームのスコア、平均、および人間で正規化されたスコアの IQM (四分位平均) を示しています。 

点推定の限界に従い、DIAMOND フレームワークは、次の図にまとめられているように、パフォーマンス プロファイルと追加のメトリックとともに、平均に対する層別ブートストラップ信頼度、および人間標準化スコアの IQM または四分位平均を提供します。 

結果は、DIAMOND がベンチマーク全体で非常に優れたパフォーマンスを発揮し、11 ゲームで人間のプレイヤーを上回り、超人的な平均 HNS 1.46 を達成し、完全にワールド モデル内でトレーニングされたエージェントの新記録を樹立したことを示しています。さらに、DIAMOND の IQM は STORM に匹敵し、他のすべてのベースラインを上回っています。DIAMOND は、Asterix、Breakout、RoadRunner など、細部のキャプチャが重要な環境で優れています。さらに、前述したように、DIAMOND フレームワークにはパイプラインに任意の拡散モデルを実装できる柔軟性があり、EDM アプローチを選択していますが、DDPM モデルを選択するのは自然な選択だったでしょう。なぜなら、DDPM モデルはすでに多数の画像生成アプリケーションで実装されているからです。EDM アプローチと DDPM 実装を比較するために、DIAMOND フレームワークは、エキスパート ポリシーを使用して収集された 100 万を超えるフレームを含む同じ共有静的データセットで、同じネットワーク アーキテクチャを使用して両方のバリアントをトレーニングします。ノイズ除去ステップの数は、世界モデルの推論コストに直接関係するため、ステップ数が少ないほど、想像上の軌道でエージェントをトレーニングするコストが削減されます。世界モデルが、タイムステップごとに 16 NFE を必要とする IRIS などの他のベースラインと計算上同等であることを保証するため、ノイズ除去ステップは数十以下、できればそれ以下を使用することを目指しています。ただし、ノイズ除去ステップの数を低く設定しすぎると、視覚的な品質が低下し、エラーが重なることがあります。さまざまな拡散バリアントの安定性を評価するために、次の図では、さまざまなノイズ除去ステップ数 n ≤ 1000 を使用して、t = 10 タイムステップまで自己回帰的に生成された想像上の軌道を示しています。 

この状況でDDPM（a）を使用すると、深刻な複合誤差が生じ、世界モデルがすぐに分布から外れてしまうことがわかります。対照的に、EDMベースの拡散世界モデル（b）は、1回のノイズ除去ステップでも、長期間にわたってはるかに安定しています。 拡散世界モデル DDPM (左) と EDM (右) に基づくモデルが示されています。t = 0 での初期観測値はどちらも同じで、各行はノイズ除去ステップ数 n の減少に対応しています。DDPM ベースの生成では誤差が増大し、ノイズ除去ステップ数が少ないほど誤差の蓄積が速くなることがわかります。対照的に、DIAMOND の EDM ベースの世界モデルは、n = 1 の場合でもはるかに安定しています。最適な単一ステップ予測は、特定のノイズ入力に対する可能な再構成の期待値であり、事後分布がマルチモーダルである場合は分布から外れることがあります。Breakout などの一部のゲームでは、単一のノイズ除去ステップで正確にモデル化できる決定論的な遷移がありますが、他のゲームでは部分的な観測可能性を示し、マルチモーダルな観測分布になります。これらの場合、次の図の Boxing ゲームに示すように、サンプリング手順を特定のモードに導くために反復ソルバーが必要です。その結果、DIAMOND フレームワークでは、すべての実験で n = 3 に設定されました。

上の図は、ボクシングにおけるシングルステップ (上段) とマルチステップ (下段) のサンプリングを比較したものです。黒のプレイヤーの動きは予測不可能であるため、シングルステップのノイズ除去では可能な結果の間を補間し、予測がぼやけてしまいます。対照的に、マルチステップのサンプリングでは、生成を特定のモードに導くことで鮮明な画像を生成します。興味深いことに、ポリシーは白のプレイヤーを制御するため、その行動はワールド モデルに認識され、あいまいさが排除されます。したがって、シングルステップとマルチステップのサンプリングはどちらも白のプレイヤーの位置を正しく予測します。

上の図では、DIAMOND によって想像された軌道は、IRIS によって想像されたものと比較して、一般的に視覚的な品質が高く、実際の環境に忠実です。IRIS によって生成された軌道には、敵が報酬として表示されたり、その逆が行われたりするなど、フレーム間の視覚的な不一致 (白いボックスで強調表示) があります。これらの不一致は数ピクセルにしか影響しないかもしれませんが、強化学習に大きな影響を与える可能性があります。たとえば、エージェントは通常、報酬をターゲットにして敵を回避することを目指しているため、これらの小さな視覚的な不一致により、最適なポリシーを学習することがより困難になる可能性があります。この図は、IRIS (左) と DIAMOND (右) で想像された連続したフレームを示しています。白いボックスは、IRIS で生成された軌道にのみ現れるフレーム間の不一致を強調表示しています。Asterix (上段) では、敵 (オレンジ) が 7 番目のフレームで報酬 (赤) になり、XNUMX 番目のフレームで敵に戻り、XNUMX 番目のフレームで再び報酬になります。 Breakout (中央の列) では、レンガとスコアがフレーム間で一貫していません。Road Runner (下の列) では、報酬 (道路上の小さな青い点) がフレーム間で一貫してレンダリングされていません。これらの不一致は DIAMOND では発生しません。Breakout では、赤いレンガが壊れるとスコアが確実に +XNUMX 更新されます。 

まとめ：

この記事では、拡散世界モデル内でトレーニングされた強化学習エージェントである DIAMOND について説明しました。DIAMOND フレームワークは、拡散世界モデルが長期間にわたって効率的かつ安定的であり続けるように、慎重な設計選択を行います。フレームワークは、これらの設計選択の重要性を示す定性分析を提供します。DIAMOND は、定評のある Atari 1.46k ベンチマークで平均人間正規化スコア 100 という新たな最先端技術を打ち立てました。これは、完全に世界モデル内でトレーニングされたエージェントの中で最高です。画像空間で動作することで、DIAMOND の拡散世界モデルは環境をシームレスに置き換えることができるため、世界モデルとエージェントの動作に関するより深い洞察が得られます。特に、特定のゲームでのパフォーマンスの向上は、重要な視覚的詳細のモデリングが改善されたことによるものです。DIAMOND フレームワークは、環境を、一連の状態、一連の離散アクション、一連の画像観測を含む標準 POMDP (部分観測マルコフ決定プロセス) としてモデル化します。遷移関数は環境のダイナミクスを記述し、報酬関数は遷移をスカラー報酬にマッピングします。