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

さらに、画像生成モデルと世界モデルとの類似性から、生成的なビジョンモデルでの進歩が世界モデルに利益をもたらすことが示唆されました。自然言語処理フレームワークでのトランスフォーマーの使用が人気を博したとき、DALL-EとVQGANフレームワークが登場しました。これらのフレームワークは、画像を離散的なトークンに変換するために離散的なオートエンコーダーを実装し、オートレグレッシブトランスフォーマーのシーケンスモデリング能力を利用して、高度に強力で効率的なテキストから画像への生成モデルを構築することができました。同時に、拡散モデルが注目を集め、現在、拡散モデルは高解像度画像生成の支配的なパラダイムとして確立されています。拡散モデルと強化学習の能力を考慮して、これら2つのアプローチを組み合わせる試みが行われており、拡散モデルの柔軟性を軌道モデル、報酬モデル、プランナー、およびオフライン強化学習のデータ増強のポリシーとして利用することを目的としています。

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

DIAMOND：拡散を環境の夢のモデルとして

世界モデルまたは環境の生成モデルは、生成エージェントが環境について計画し、推論するための重要なコンポーネントとして登場しています。強化学習の使用は近年大きな成功を収めてきましたが、強化学習を実装するモデルは、サンプル効率が低いことで知られており、これは実世界での応用において大きな制限となります。他方で、世界モデルは、さまざまな環境で強化学習エージェントを効率的にトレーニングする能力を示し、モデルが実世界の経験から学ぶことを可能にしました。最近の世界モデリングフレームワークは通常、環境のダイナミクスを離散的な潜在変数のシーケンスとしてモデル化します。モデルは潜在的な空間を離散化して、複合エラーを避けるため、離散的な潜在変数を使用することで大きな結果をもたらすことができます。ただし、このアプローチは情報の損失にも関連しており、再構成の品質と一般性の損失につながります。情報の損失は、情報が適切に定義されている必要がある実世界のシナリオ、たとえば自動運転車のトレーニングにおいて、大きな障害となる可能性があります。こうしたタスクでは、視覚的な入力の小さな変更や詳細、たとえば交通信号の色や前の車両のターンシグナルは、エージェントのポリシーを変更する可能性があります。離散的な潜在変数の数を増やすことで情報の損失を避けることができますが、計算コストが大幅に増加します。

さらに、近年、拡散モデルは、高品質の画像生成フレームワークの支配的なアプローチとして登場しました。拡散モデルは、ノイズ化プロセスを逆転させることでサンプルを生成し、非平衡熱力学から多大なインスピレーションを得ています。DIAMONDフレームワークは、連続的な時間変数で索引付けされた拡散プロセスを考慮し、対応する余分と境界条件とともに、扱いやすい非構造化事前分布を持ちます。さらに、生成モデルを取得するために、DIAMONDフレームワークはノイズからデータへのマッピングを行う必要があります。逆転プロセスも拡散プロセスであり、時間の逆方向に実行されます。さらに、任意の時点で、スコア関数を推定することは簡単ではありません。DIAMONDフレームワークは、実際のスコア関数へのアクセスがないため、スコアマッチングオブジェクトを実装して、スコアモデルをトレーニングするフレームワークを提供します。スコアベースの拡散モデルは、無条件の生成モデルを提供します。ただし、環境のダイナミクスを条件付けるために、世界モデルとして使用するには、条件付き生成モデルが必要です。DIAMONDフレームワークは、一般的なPOMDPアプローチを考慮し、過去の観測と行動を使用して未知のマルコフ状態を近似することができます。図1に示すように、DIAMONDフレームワークはこの履歴を使用して拡散モデルを条件付けて、次の観測を直接推定および生成します。理論的には、DIAMONDフレームワークは任意のSDEまたはODEソルバーを使用できますが、NFE（関数評価の数）とサンプルの品質のトレードオフがあり、拡散モデルの推論コストに大きな影響を与えます。

上記の知識を基に、拡散ベースの世界モデルの実用的実現、特に拡散係数と対応する拡散アプローチの特定の選択について見てみましょう。DDPMは自然な候補ですが、DIAMONDフレームワークはEDM形式を使用し、拡散時間の実数値関数であるノイズスケジュールを持つ摂動カーネルを考慮します。フレームワークは、入力と出力の分散を維持するために、前処理器を選択します。ネットワークトレーニングは、劣化レベルに応じてシグナルとノイズを適応的に組み合わせます。ノイズが低い場合、ターゲットはクリーンシグナルと摂動シグナルとの差、つまり追加されたガウシアンノイズになります。直感的には、これにより、低ノイズレジームでのトレーニングオブジェクトが簡単になるのを防ぎます。実践的には、この目的はノイズスケジュールの端では高分散であるため、ノイズレベルを対数正規分布からサンプリングして、トレーニングを中ノイズ領域に集中させます。DIAMONDフレームワークは、ベクトル場のために標準的なU-Net 2Dコンポーネントを使用し、過去の観測と行動を保持するバッファーを保持します。DIAMONDフレームワークはこれらの過去の観測を次のノイズ観測に結合し、適応グループ正規化レイヤーを使用して、U-Netの残差ブロック内の入力アクションを結合します。

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

拡散モデルは、ノイズ化プロセスを逆転させることでサンプルを生成する、生成モデルのクラスです。非平衡熱力学から多大なインスピレーションを得ています。DIAMONDフレームワークは、連続的な時間変数で索引付けされた拡散プロセスを考慮し、対応する余分と境界条件とともに、扱いやすい非構造化事前分布を持ちます。さらに、生成モデルを取得するために、DIAMONDフレームワークはノイズからデータへのマッピングを行う必要があります。逆転プロセスも拡散プロセスであり、時間の逆方向に実行されます。さらに、任意の時点で、スコア関数を推定することは簡単ではありません。DIAMONDフレームワークは、実際のスコア関数へのアクセスがないため、スコアマッチングオブジェクトを実装して、スコアモデルをトレーニングするフレームワークを提供します。スコアベースの拡散モデルは、無条件の生成モデルを提供します。ただし、環境のダイナミクスを条件付けるために、世界モデルとして使用するには、条件付き生成モデルが必要です。DIAMONDフレームワークは、一般的なPOMDPアプローチを考慮し、過去の観測と行動を使用して未知のマルコフ状態を近似することができます。図1に示すように、DIAMONDフレームワークはこの履歴を使用して拡散モデルを条件付けて、次の観測を直接推定および生成します。理論的には、DIAMONDフレームワークは任意のSDEまたはODEソルバーを使用できますが、NFE（関数評価の数）とサンプルの品質のトレードオフがあり、拡散モデルの推論コストに大きな影響を与えます。

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

点推定の限界に従って、DIAMONDフレームワークは、人間の正規化スコアの平均とIQMの層化ブートストラップ信頼区間を提供します。さらに、パフォーマンスプロファイルと追加のメトリクスが次の図にまとめられています。

結果は、DIAMONDがベンチマーク全体で優れたパフォーマンスを示していることを示しています。11のゲームで人間のプレイヤーを上回り、平均人間の正規化スコア1.46を達成し、世界モデル内で完全にトレーニングされたエージェントの新しいベンチマークを設定しています。さらに、DIAMONDのIQMはSTORMと比較可能であり、他のすべてのベースラインを上回っています。DIAMONDは、 Asterix、Breakout、RoadRunnerなどの環境で、小さな詳細を捉えることが重要な場合に優れています。さらに、前述のように、DIAMONDフレームワークは、パイプラインで任意の拡散モデルを実装する柔軟性を持っています。EDMアプローチを選択していますが、DDPMモデルはすでに多くの画像生成アプリケーションで実装されているため、DDPMモデルを選択することも自然な選択です。EDMアプローチとDDPMの実装を比較するために、DIAMONDフレームワークは、同じネットワークアーキテクチャを使用して、同じ静的データセットで両方のバリアントをトレーニングします。世界モデルの推論コストに直接関係するため、ノイズ除去ステップの数を最小限に抑えることを目指しています。理想的には10未満のノイズ除去ステップを使用することを目指しています。ただし、ノイズ除去ステップの数を低すぎる設定すると、視覚的な品質が低下し、複合エラーが発生する可能性があります。さまざまな拡散バリアントの安定性を評価するために、n≤10の異なるノイズ除去ステップ数を使用して、t=1000のタイムステップまでオートレグレッシブに生成された想像トラジェクトを表示します。

DDPM（a）を使用すると、このレジームでは深刻な複合エラーが発生し、世界モデルがすぐに分布から外れてしまいます。一方、EDMベースの拡散世界モデル（b）は、長時間のホライズンで、1つのノイズ除去ステップでも、はるかに安定しています。想像トラジェクトは、DDPM（左）とEDM（右）に基づく拡散世界モデルで生成されます。初期観測はt=0で同じですが、各行はノイズ除去ステップの減少を表します。DDPMベースの生成は複合エラーに苦労し、ノイズ除去ステップの数が少ないとエラーが早く蓄積します。一方、DIAMONDのEDMベースの世界モデルは、n=1の場合でもはるかに安定しています。最適な1ステップ予測は、与えられたノイズ入力に対する可能な再構成の期待値であり、後部分布が多峰性の場合、分布外の場合になります。いくつかのゲーム、たとえばBreakoutは、1つのノイズ除去ステップで正確にモデル化できる決定的な遷移を持ちますが、他のゲームは部分的に観測可能であり、多峰性の観測分布を示します。このような場合、特定のモードに向けてサンプリング手順を導くための反復ソルバーが必要です。ボクシングゲームの図に示すように。したがって、DIAMONDフレームワークはすべての実験でn=3を設定します。

上記の図は、ボクシングでのシングルステップ（上行）とマルチステップ（下行）サンプリングを比較しています。黒いプレイヤーの動きは予測できないため、シングルステップの除去は可能な結果の間で補間し、ぼやけた予測を生成します。一方、複数ステップのサンプリングは、特定のモードに向けて生成を導くことで、明確な画像を生成します。興味深いことに、ポリシーが白いプレイヤーを制御するため、彼の行動は世界モデルにわかっているため、両方のシングルステップとマルティステップのサンプリングは、白いプレイヤーの位置を正しく予測します。

上記の図は、DIAMONDによって想像されたトラジェクトリとIRISによって想像されたトラジェクトリを比較しています。IRISによって生成されたトラジェクトリには、フレーム間の視覚的な不一致（白いボックスで強調表示）が含まれています。たとえば、敵が報酬として表示されたり、その逆が表示されたりします。こうした不一致はわずか数ピクセルに影響する可能性がありますが、強化学習の学習に大きな影響を与える可能性があります。たとえば、エージェントは通常、報酬をターゲットとし、敵を避けることを目指します。こうした小さな視覚的な不一致は、最適なポリシーを学習することをより困難にします。図は、IRIS（左）とDIAMOND（右）で想像された連続するフレームを示しています。白いボックスは、IRISによって生成されたトラジェクトリでのみ発生するフレーム間の不一致を強調しています。Asterix（上行）では、敵（オレンジ）が2番目のフレームで報酬（赤）になり、3番目のフレームで敵になり、4番目のフレームで再び報酬になります。Breakout（中行）では、ブロックとスコアがフレーム間で一貫性がありません。Road Runner（下行）では、報酬（道路上的小さな青い点）はフレーム間で一貫してレンダリングされていません。こうした不一致はDIAMONDでは発生しません。Breakoutでは、赤いブロックが破壊されたときにスコアが+7で信頼性の高い更新されます。

結論

この記事では、DIAMONDについて説明しました。DIAMONDは、拡散世界モデル内でトレーニングされた強化学習エージェントです。DIAMONDフレームワークは、拡散世界モデルが長時間のホライズンで効率的かつ安定したままであることを保証するための、慎重な設計選択を行います。DIAMONDは、Atari 100kベンチマークで新しいベンチマークを設定し、平均人間の正規化スコア1.46を達成し、世界モデル内で完全にトレーニングされたエージェントの最高スコアを達成します。画像空間で動作することで、DIAMONDの拡散世界モデルは環境をシームレスに置き換えることができ、世界モデルとエージェントの動作についての洞察を提供します。特に、特定のゲームでのパフォーマンスの向上は、重要な視覚的な詳細のモデリングの改善に帰因します。DIAMONDフレームワークは、環境を標準的なPOMDPまたは部分的に観測可能なマルコフ決定プロセスとしてモデル化し、状態のセット、離散的な行動のセット、画像観測のセットを持ちます。遷移関数は環境のダイナミクスを記述し、報酬関数は遷移をスカラー報酬にマップします。