🦾

Physical AIに向けて:Robotic World Model (RWM)

に公開

はじめに

Physical AIに挑んでいくための一つの大きなテーマとしてWorld Modelがある。今回はZürichの論文をテーマに最新の情報を追っていくこととした。Zürich工科大学はEnd2End強化学習で4脚実ロボットを制御し、現在の4脚や二足歩行ロボットの技術的基盤や方向性を示したパイオニアだと思っている。

2017年「Asymmetric Actor Critic for Image-Based Robot Learning」という論文が出てきた。これはUCバークレーやOpenAI、CMUなどに所属する研究者によって提唱された手法であるが、これを主軸に「特権情報」を利用した訓練によってEnd2End強化学習を実機で実現し、様々な成果を世に送り出したのがZürich工科大学である。

無論、今回の主題であるWorld Model自体もだいぶ前からある概念であるし、Zürich工科大学から出てきたものではない。しかし、重要なのは僕の好きな物理モデリングと最適制御理論に基づくような話はすっ飛ばされ、データドリブンでロボットを制御できるということを実機を用いて示してくれたことである。

深層強化学習と模倣学習は、大きくロボットのAI化を進めてきたが、現在はWorld Modelの実応用に向けて色々な分野で研究が進んでいる。(これはロボットに限ったことではない)

今回の論文もしっかり実機を動かしており、今後の潮流を作りうると思って選定した。

下記の論文の結論を読むだけでも良い。

技術的な興味があるなら「二重自己回帰」の項目だけでも抑えると良い。

今回の論文の本質的な貢献はここである。そして、この部分は、これまで時系列のフォアキャストタスクに、深層学習で真剣に取り組んだことがある人であれば一度は試したことがあるのではないか?というものになっている。

論文の結論

さすがNotebookLM、良い感じの一枚絵を作成してくれた。

研究が達成した成果、学術的な貢献、および残された課題を文章で書くと下記の通りだ。

成果

提案された「Robotic World Model (RWM)」は、二重自己回帰機構と自己教師あり学習を組み合わせることで、特定のロボットやタスクに関するドメイン知識を必要とせず、長期間にわたる正確なダイナミクス予測を実現した。自己教師あり学習というのは、要するに取得した時系列データで未来予測のデータセットで訓練するということだ。人間がラベルを張らなくて良いという意味で自己教師あり学習と呼んでいる。(この言葉は今も納得いってないが…まあ仕方ない)

これにより、誤差の蓄積を効果的に抑制し、100ステップを超える長期にわたる想像上の軌道生成を可能とした。

技術的貢献

RWMをシミュレータとして活用し、PPO(Proximal Policy Optimization)アルゴリズムでポリシーを最適化するフレームワークMBPO-PPOを提案した。これによって、ロボットのシミュレータは通常動作が遅いのだが、RWMを活用することで計算が高速化できる。すなわち強化学習が速く行えるということだ。

従来よりReduced Order Modelingとかサロゲートモデリングなどと呼ばれている概念である。流体解析とかではよくある話であるが、ここではセンサーやロボットの動作を含むマルチモーダルなデータの時間発展を二重自己回帰でサロゲート化したということだ。

このフレームワークを用いてシミュレーション内で学習したポリシーを、四足歩行ロボット(ANYmal D)やヒューマノイドロボット(Unitree G1)の実機にゼロショット(追加学習なし)で転移させることに成功し、実環境での頑健な動作を実証した。これは正直驚異的である。

「実データ」->「RWM」->「PPO」という流れで情報が伝わる。つまり推定・近似が二回入っているのである。それでも実機に戻してあげたらちゃんと動くというのは、数年前だとにわかには信じられない。(World Modelの長期予測性能の低さ、ないし二回の近似によるモデル化誤差は通常許容範囲を超える)

課題

高性能な物理シミュレータで膨大な時間をかけて学習させたモデルフリー強化学習の最高性能には、まだ及ばない点が挙げられる。やはり二回の近似が入ると性能が落ちるという点は否めないらしい。それでも様々なモデル・学習アルゴリズムを「高性能な物理シミュレータ」で個別に検討するには異常な時間が掛かると推察される。World Modelであれば一度作れば、その中での訓練結果を比較するという作業の効率は劇的に上がるだろう。World Modelでの訓練と物理シミュレータでの訓練の結果に相関さえあれば、事前検討として利用可能だと自分は思う。

ただ、学習初期の不安定さを避けるため、シミュレーションデータによる事前学習が必要である。これは特に接触ダイナミクスが不連続な脚式ロボットのタスクで顕著であり、マニピュレーションタスクでは不要であったらしい。これは結構痛い気がする。とはいっても、これはPPO等の強化学習アルゴリズムをまず妥当な状況に乗せるために必要なのだと思われるため、World Model側の課題であるか強化学習側の課題であるかはまあ微妙だと思う。

実機での安全かつ効果的なオンライン学習には依然として課題が残されている。これは、強化学習自体の問題でもある。

リサーチクエスチョン

本研究が解決しようとしている中心的な問いは、現代のロボティクスが直面する根源的な課題に基づいている。従来のロボット制御手法の多くは、一度デプロイされると、実世界で得られるデータから継続的に学習・適応する能力に欠けており、その頑健性(ロバスト性)は限られている。

既存のワールドモデル研究は、しばしば特定のタスクやロボットに特化した ドメイン知識(誘導バイアス) をモデル構造に組み込むことで性能を確保している。これは、新しいロボットや未知のタスクへのスケーラビリティや汎用性を損なう原因となる。

特に、脚式ロボットのように複雑で、一部の状態しか観測できず(部分観測性)、確率的なダイナミクスを持つ実世界のロボット制御において、予測誤差を蓄積させることなく長期間の未来を正確に予測することは非常に困難である。

これらの背景から、本研究は以下のリサーチクエスチョンに挑んだ内容となっている。

「ドメイン固有の知識に頼らず、多様なロボットタスクに適用可能で、長期間にわたる正確な予測ができる頑健なワールドモデルをいかにして学習し、実世界での効率的なポリシー最適化を実現するか?」

既に述べてある通り、上記のクエスチョンにはある程度答えているものの、本気のシミュレータ上での訓練には及ばないのと、デプロイ後に継続的にオンライン学習する部分には依然として課題がある。とはいえ、World Model上での長期予測が可能で、サロゲートを利用した強化学習がゼロショットで実機に適用できたのは驚異的なのは変わりない。

手法

問題設定

ロボットが動作する環境を 「部分観測マルコフ決定過程(POMDP)」 として定式化する。これは、ロボットが環境の全ての状態を直接観測できるわけではない、という実世界に近い設定である。この設定のもと、ワールドモデルを用いた学習は一般的に以下の3つのステップを繰り返す。

  1. 実環境でのデータ収集: 現在のポリシーに従ってロボットを実環境で動かし、観測と行動のデータを収集する。
  2. ワールドモデルの学習: 収集したデータを用いて、環境のダイナミクスを模倣するワールドモデルを学習・更新する。
  3. ワールドモデル内でのポリシー最適化: 学習したワールドモデルをシミュレータとして使い、その「想像」の中でポリシーを効率的に最適化する。

On-Policy形式の訓練であるからこうなってしまうか。ちょっと正直、スキームが想像と違ったぜ。
常識的に考えて、実機で適当に動作させたときのデータを大量に取得しておいて、そのあとはワールドモデルを訓練し、最後にサロゲート上での強化学習という流れが普通な気がするよね。ってか別にサロゲートに対しては普通にオンポリシーで動けるじゃん。なんでこのスキームなんだ?

二重自己回帰

RWMの最も重要な技術的貢献は、その独自の学習フレームワークにある。まず、モデルは 過去Mステップの観測・行動履歴を入力として、未来Nステップの観測を予測するように学習する。 このとき、単に1ステップ先を予測するだけでなく、予測した未来の観測を次のステップの入力として再度利用するという自己回帰的なプロセスで学習を進める。これは、モデルが自身の予測誤差も含めて学習することを意味する。つまり、訓練時に「Nステップのクローズドループ予測問題」を解かせているというわけだ。


https://arxiv.org/html/2501.10100v4 より

この学習方法は、長期予測における誤差の蓄積を抑制する上で極めて重要である。従来の「Teacher-Forcing」と呼ばれる手法では、学習時に常に正解データを次の入力として与える。このアプローチは、モデルを完璧な正解データのみで訓練するため、モデルは自らが犯した間違いから回復する方法を学ぶ機会がない。その結果、推論時(実用時)に自身の少し不完全な予測に頼らざるを得なくなると、小さな誤差が次の予測の誤差を呼び、雪だるま式に蓄積・増幅され、軌道予測が現実から破滅的に乖離してしまうのである。

https://arxiv.org/html/2501.10100v4 より
ちなみに時系列予測を生業としている人であれば、そうした方が良いと思って実験したことは一度はあると思う。正直、実装が面倒なんよね。面倒な割に通常のタスクだとパフォームしなかったりする。ただMPCの内部モデルのためのサロゲートモデルを作ろうとすると、確実にこちらのパターンの方が良い。World ModelからのPPOなんて、やっていることは代理モデルでMPCと本質的に解いている基礎方程式が変わらないので、妥当だと思われる。

RWMは、以下の目的関数を最小化することで、複数ステップ先の予測誤差を直接的に減らすように学習される。

\mathcal{L} = \frac{1}{N} \sum_{k=1}^{N} \alpha^k \left[ L_o(o'_{t+k}, o_{t+k}) + L_c(c'_{t+k}, c_{t+k}) \right]

ここで、o'は予測された観測、oは実際の観測、L_oは両者の誤差を表す。同様に、c'cは足の接触情報などの 「特権情報(privileged information)」 の予測と実際値を表し、L_cがその誤差である。

特権情報とは、実際のデプロイ時には直接的には取得できないが、実験環境やシミュレータ上では計測ができるような情報だ。これはリーク情報を使っていると言えなくもないが、例えばActor Criticの Critic 側にリーク情報を使う分にはActorを正しい情報に基づいて評価し、Actor自身は直接的に特権情報を参照しないし、デプロイ時にも必要ではないので上手く行く。これがAsymetric Actor Criticである。

ロボットの状態だけでなく、接触のような物理的に重要な情報も予測させることで、環境ダイナミクスを理解するための追加の学習シグナルを与えている点が特徴である。\alpha^kは減衰係数であり、遠い未来の予測よりも近い未来の予測精度を優先させるために使えるが、本研究では長期予測の安定性を重視し、\alpha=1.0に設定して全ての未来ステップを平等に扱っている。

さらに、RWMはGRU(Gated Recurrent Unit)ベースのネットワークアーキテクチャに、デュアルオートリグレッシブ機構を導入している。これは、(1) 隠れ状態を更新する 「内部自己回帰」と、(2) 予測結果を入力にフィードバックする「外部自己回帰」 の2つから構成され、長期的な依存関係を捉える能力を高め、予測の頑健性を確保している。

これも恐らくやったことある人はいると思う。同じくして実装が非常に面倒くさい。普通は内部自己回帰に頼って、例えばNeural State Spaceなどでは制御入力系列と状態の初期値のみを外から与え、N期先予測を行う実装が選ばれると思う。

学習済ワールドモデル上でのポリシー最適化

学習済みワールドモデル上でのポリシー最適化学習済みのRWMをシミュレータとして使い、ポリシーを最適化する手法として 「MBPO-PPO」 が提案されている。これは、論文のAlgorithm 1に示されるように、以下のプロセスで進行する。

  1. 実環境で現在のポリシーを使ってデータを収集し、リプレイバッファ \mathcal{D} に保存する。
  2. バッファのデータを用いて、RWM p_{\phi} を自己回帰学習で更新する。
  3. バッファからランダムに初期状態をサンプリングし、更新されたRWM p_{\phi} と現在のポリシー \pi_{\theta} を使って、「想像」の中で軌道を T ステップ分生成する。
  4. 「想像」の中で生成された軌道データ(状態、行動、報酬の系列)を使って、PPOアルゴリズムでポリシー \pi_{\theta} を更新する。

PPOのようなモデルフリー強化学習アルゴリズムを学習済みモデル上で使う際には、ポリシーがモデルの不正確な部分を意図的に悪用し、現実ではありえない高報酬を得ようとする問題が知られている。 しかし、RWMは非常に高い予測精度を持つため、100ステップ以上という従来の手法では考えられなかった長い想像上の軌道を安定して生成できる。これにより、長期的な視野でのポリシー最適化が可能となり、上記の問題を克服している。

結果

結果を見ると、どうもRWM-ARが今回の提案手法の二重自己回帰である。RWM-TFは一期先予測しか訓練していない(ので、制御時の長期予測が全然だめになっているはずである)。

また、モデルベース強化学習は他にも手法がある。とくにDreamerは有名だが、こちらに対しても良い結果が出ている。

所感

基本的な時系列予測に関する、これまで安直に扱われていた部分を丁寧に実装し評価したことで結果が得られたのだと思う。最適制御問題というのは理想的には無限時間先の状態を予測できることに(理論上は)なっているので、現在状態と制御入力列から(確率的だが法則としては)出力が正確に分かっていなければいけない。

自然言語のように、「全体として意味が通ればよく、返り読みをして全体を最後に直す」が許されないため、本質的に時系列のForecastingが必要なタスクである。これまでは自然言語の分野がAIの新規手法を開拓してきた歴史があるため、一部考慮されてこなかった、あるいは抜け落ちていた概念があるように思われる。

それらを適切に回収し、モデリング(インナーループとアウターループ)を基礎に立ち返って実直に行い、適切な評価(N期先のクローズドループ予測タスク)を行った結果だなぁと思った。

おそらく難しいことはやっていない。まだまだこれから追いつける。

ただ、なおのこと丁寧な実装と適切な評価を実直に実施していくことが大事なのだということを実感した。おそらく、この領域ではアイデア自体には価値が全くないこともしばしばある(思いついたことあるし、やったこともあるということが多数存在する)。ただ、それをやり続けることが大切だと改めて実感した。

Discussion