yoneken

SuperMimicは、ヒューマノイドの全身動作を学習するために、従来のPPOを中心とするモデルフリー強化学習ではなく、学習した微分可能な世界モデルを通じた誤差逆伝播法（BPTT）によって方策を直接最適化する。その結果、従来128基のGPUで7日間かけて学習したモデルの性能を、わずか1基のGPU、3時間の学習で上回った。 問題設定は、あらかじめ正解となる人間の参照動作データが存在し、それをロボットに模倣させるタスクである。正解となる動作の位置、姿勢、並進速度、角速度がロボットの実際の動作に近ければ近いほど、高い報酬を与える。 学習は以下の3段階を繰り返す。 ・現在の方策をMuJoCo Warp上で実行し、状態・行動・次状態をリプレイバッファに保存する。 ・8ステップの自己回帰予測を使い、ロボットのダイナミクスをニューラルネットワークで近似する世界モデルを学習する。 ・学習した世界モデル上で32ステップの仮想ロールアウトを実行し、報酬から行動、方策へと勾配を逆伝播（BPTT）して方策を最適化する。 これらを同一の学習ループで交互に繰り返し、世界モデルと方策を並行して学習する。 世界モデル自体は3層・各2048ユニットのMLPであり、巨大なモデルではない。これだけ小さなモデルでも機能するのは、ロボットの運動が視覚や言語とは異なり、比較的低次元の状態空間と共通の物理法則によって記述できるためではないかと考えられる。 高速に学習できる理由は、大きく2つある。1つ目は世界モデルによるシミュレーションが高速であることと、2つ目は誤差逆伝播法によって方策を効率的に最適化できることである。 世界モデルはMuJoCo Warpによる物理シミュレーションより高速に状態遷移を計算でき、実際に使われる並列環境数（2K〜16K）では10倍以上、最大64Kの並列環境では62倍高速に実行できる。 さらに、従来のPPOのようにサンプリングした行動から方策勾配を推定するのではなく、微分可能な世界モデルを通じて報酬の勾配を直接計算するため、効率的に方策を改善できる。 実験では、同じ世界モデル上でPPOを学習させた場合に比べ、約64分の1の仮想ステップ数で同等の報酬を達成している。 今回、世界モデルは8ステップの自己回帰予測で学習するが、方策の最適化では32ステップ先まで展開して勾配を計算する。なお、32ステップは50Hzの制御周期で0.64秒に相当する短い時間である。 また、32ステップだけでは長期的な安定性や将来の報酬を十分に評価できないため、Criticによる価値推定を組み合わせる。これにより、短いロールアウトで勾配を計算しながら、それ以降の長期的な報酬も考慮できる。 世界モデルは摩擦係数、重心位置、キャリブレーション誤差などのシミュレーションコンテキストを入力として使う。一方、実機に配置する方策はこれらを直接利用せず、関節位置・速度やIMUなどの観測情報を使う。いわゆるPrivileged Informationを学習時に利用する構成になっている。 コメント === 微分可能な世界モデルを使って効率的に学習するアイデアは古くから数多く提案されてきたが、今回の成功には、長期のダイナミクスを直接ロールアウトするのではなく、短い時間の予測と勾配計算に限定したことが大きく寄与していると考えられる。 具体的には、世界モデルを8ステップ（0.16秒）の自己回帰予測で学習し、方策のBPTTは32ステップ（0.64秒）に限定する。それ以降の長期的な価値についてはCriticを利用することで、予測誤差の蓄積を抑えながら方策を改善できる。 また、並列化可能で高速なMuJoCo Warpをさらに、学習された世界モデルで高速化し、1秒あたり合計数万秒分のロボット動作を仮想的に計算できることになる。 さらに、世界モデルを現在の方策が訪れる状態分布に合わせて継続的に学習することも重要だと考えられる。方策の改善によって訪れる状態や選択する行動が変化しても、新しいシミュレーションデータを収集して世界モデルを更新するため、方策が世界モデルの学習範囲から大きく逸脱する問題を抑えられる。 Appendix D.1に書かれている、方策が失敗しやすい区間を重点的にサンプリングする工夫も興味深い。動作データを1秒単位に分割し、Returnの低い区間を優先して学習するようにしている。例えば、宙返りの着地が難しい場合、動作の最初から繰り返すのではなく、着地直前の状態から重点的に学習できる。 また、今回のタスクは、あらかじめ正解となる動作が存在し、位置や姿勢の誤差から各ステップで滑らかな報酬を計算できるという特徴がある。こうした密な報酬も、微分可能な世界モデルを通じた方策最適化が成功した重要な要因だと考えられる。 そのため、正解動作が存在しないタスクや、物体との複雑な接触を伴う操作タスクにおいても同じような学習効率を達成できるかは、今後の重要な研究課題となる。