【DL輪読会】Octo: An Open-Source Generalist Robot Policy　　ーロボット基盤モデルの構築に向けてー

Octo: An Open-Source Generalist Robot Policy ーロボット基盤モデルの構築に向けてー 2024.2.2 Presenter: Tatsuya Matsushima @__tmats__ , Matsuo Lab 1

• https://twitter.com/__tmats__

概要 大規模データセットで学習したロボット制御モデル （ロボット基盤モデル）Octoの開発に関する報告 • transformerベースの拡散モデルの制御方策モデル • 800kエピソードのデータセットで学習 • モデル・学習コードを公開 2

書誌情報 Octo: An Open-Source Generalist Robot Policy • Dibya Ghosh*1, Homer Walke*1, Karl Pertsch*1,2, Kevin Black*1, Oier Mees*1, Sudeep Dasari3, Joey Hejna2, Tobias Kreiman1, Charles Xu1, Jianlan Luo1, You Liang Tan1, Dorsa Sadigh2, Chelsea Finn2, Sergey Levine1 • 1UC Berkeley, 2Stanford, 3Carnegie Mellon University • https://octo-models.github.io/paper.pdf • Website: https://octo-models.github.io/ • Github: https://github.com/octo-models/octo • 論文というよりテクニカルレポート • 試してうまくいったこと，いかなかったことが列挙されていて，同様なロボット基盤モデルを 開発する際に参考になる • どこかの学会に提出しているのかどうかは不明（arXivにもない） ※特に出典が明記されていない図は当論文から引用 3

• https://octo-models.github.io/paper.pdf
• https://octo-models.github.io/
• https://github.com/octo-models/octo

背景：ロボット制御での大規模モデル RT-1 / RT-2 大規模なデータセット・モデルを活用したEnd-to-Endな ロボットの制御方策の開発が加速 • 例：RT-1（35Mパラメータ，13台で130K軌道，744言語指示） • 言語指示に対応，transformerベース • 例：RT-2（12B/55Bパラメータ，同じデータ） • LLMのfinetuningで実現 https://robotics-transformer1.github.io/ https://robotics-transformer2.github.io/ 4

• https://robotics-transformer1.github.io/
• https://robotics-transformer2.github.io/

背景：ロボット制御での大規模モデル RT-X/Open X-Embodiment (OXE) dataset Google Deepmindと世界21研究機関がオフラインのロボット データセットを統一のフォーマットに変換し利用 • 22種類のロボット，527スキル (160,266タスク)，100万エピ ソード以上 • 日本からは東大3研究室が参加（松尾研・JSK・原田研） • RT-1/RT-2のモデルをこれらのデータで学習すると ロボットのmorphologyに（ある程度）汎化するモデルが得られた • 個別のデータで学習したRT-1/2よりも良い性能 5 https://robotics-transformer-x.github.io/

• https://robotics-transformer-x.github.io/

これまでのロボット基盤モデルの問題点 新しい入出力に拡張するのが困難 • 事前学習時とは異なる種類の観測に対応できない • 例）事前学習に使ってなかったproprioceptionをfinetuningで入れる • 事前学習時とは異なる種類の行動空間に対応できない • 例）事前学習では手先姿勢の差分を出力していたがfinetuningでは関節角を出力 言語指示以外のタスクの指示に対応できない • 例）言語指示とgoal-conditional両方に対応したい モデルや実装がオープンではない • 例）RT-1Xではモデルと重みは公開されているが学習コードが公開されていない • 例）RT-2Xは重みも公開されていない • APIのリクエストフォームだけがある 6

Octoの思想 オープンソースな大規模なモデルを作る • 800kエピソードのデータで学習されたtransformerベースの拡散 モデルの方策（27M・93M）のモデルとコードを公開 複数の入力形式に対応し，事前学習時とは異なる入出力形式の データでfinetuningできるようにする • Transformerのバックボーンのアーキテクチャをモジュラに なるように工夫（次で説明） 7

Octoのモデル Octoのモデルアーキテクチャ 入力・観測のtokenizer，transformerのバックボーン， readout head (action head)の3つで構成 • 系列の最初にタスク指示のトークンをtransformerのバックボーン に入力し，その後，観測のtoken列を入力 • Readoutのtokenを action headで行動に変換 • タスク指示のトークンは 言語指示and/orゴール画像 で作成 8

Octoのモデルアーキテクチャ 入力・観測のtokenizer モダリティごとにtokenizerを用意 • 言語：Pretrainedされたtransformer（T5，111M）を利用 • 画像：浅いCNNを利用 • 16x16のパッチに切り特徴量 を1トークンにして並べる • 小さいパッチにすると性能 上がるが計算量増える トレードオフあり • transformerの計算量が token長の2乗に比例 9

Octoのモデルアーキテクチャ Transformerのバックボーン Block-wiseなcausal maskをかけてattention • その時刻とその前のステップの観測（青）とタスク指示のトークン（緑）のみにattention • 対応していないトークンはmask outする • これにより新しい入出力形式に対応可能 • 例）言語指示タスクではないときの言語に対応するトークン • Readoutトークンを各時刻の観測の 間に挿入 • Action headにより行動に変換される トークン • その時刻以前の観測トークンと タスクトークンのみにattend • 観測トークンやタスクトークンからは attentionがかからないように mask out 10

Octoのモデルアーキテクチャ Readout head (action head) Readoutトークンを行動に変換 • Diffusion policyを採用 • 3層のMLP • 隠れ層の次元256，residual connection有，layer normalization • ACTと同様に複数ステップ 行動を出力 • （参考）DL輪読会山根さん • https://www.docswell.com/s/DeepLearning2023/K7VL3W-dllearningfinegrained-bimanual-manipulation-with-lowcost-hardware 11

• https://www.docswell.com/s/DeepLearning2023/K7VL3W-dllearning-finegrained-bimanual-manipulation-with-lowcost-hardware

Octoのモデルアーキテクチャ Octoのモデルのデザインチョイス 画像のエンコーダのアーキテクチャの選択 • 過去の研究ではResNetベースの画像のエンコーダを使い，小さい transformerがが使われることが多かった • 本研究ではCNNを小さくしてtransformerをできるだけ大きくした • こちらの方が経験的に性能が高まった • ImageNetで学習ずみのResNetを利用しても性能は高まらなかった 複数の入力をくっつけるタイミング • Transformerは計算量がトークン長の2乗に比例 • それぞれの入力を別々にtransformerに入力すると遅くなる • 本研究では，ゴール画像がある場合は観測画像のチャネルに 追加してトークン長が長くなることを防ぐ 12

Octoの学習 Octoの事前学習に用いたデータ Open-X Embodiment (OXE)データセットから 25種類を利用（計800kエピソード） • 画像入力で行動空間が手先姿勢の差分のもの • ある程度画像の解像度が高くタスクの多様性が ありニッチすぎないタスクをやっているもの 各データセットの出現割合を調整 • 多様なタスクをやっているデータセットを2倍に • データの量が大きいタスクの割合を少し下げる 13

Octoの学習 Octoの事前学習に用いた目的関数 Action headにdiffusion policyを利用 • （参考）Diffusion policyに関するDL輪読会（山根さん） • https://www.docswell.com/s/DeepLearning2023/KP9XJXdldiffusion-policy-visuomotor-policy-learning-via-action-diffusion • Readoutトークンで条件づけて行動空間でのdenoisingとして行動 を出力 • Noise schedulerはDDPMを利用 • 学習時も推論時も20 diffusion steps 14

• https://www.docswell.com/s/DeepLearning2023/KP9XJX-dldiffusion-policy-visuomotor-policy-learning-via-action-diffusion

Octoの学習 Octoのfine-tuning 事前学習時と同じ目的関数でfinetuning • ターゲットドメイン（ロボットやタスク）のデータが 100エピソードぐらいあればできる • モデル全体をfinetuning • Action headだけ学習するより性能が良かった • 50kステップの学習 15

Octoの学習 モデルサイズ・データ拡張 2種類のモデルサイズで実験 • ViT-Sサイズ（27M），ViT-Bサイズ（93M） • ViT-Bサイズの事前学習はバッチサイズ2048で300kステップ • TPU v4-128 podで約14時間 • FinetuningはA5000x1で約5時間 • 推論時はどちらもRTX4090x1で10Hz以上出るらしい 16

実験 Zero-shotでの方策の利用 • 学習に用いたロボットと環境での検証 • RT-1Xより33パーセント高い成功率 • RT-2Xの55Bのモデルと同等レベルの成功率 17

実験 Finetuningで検証 • 一部新しい観測や行動空間を実験 • ResNet+transformerの組み合わせでscratchから学習する場合や VC-1の学習済みモデルを活用する場合よりも良い性能 18

うまくいったこと (Apendix F) 事前学習時に過去の履歴を入力する • コンテクストに1フレーム過去の履歴も入力するとzero-shot性能が向上 • 検証したタスクではこれ以上長い履歴を入力しても改善しなかった（他のタスクでは改善するかも） Action Chunkingを使う • ACTをと同様に将来の複数ステップの行動を予測 • 出力される動作に時間的な一貫性が生まれた 画像のパッチサイズを小さくする • 16x16のパッチと32x32のパッチで比較 • 16x16のパッチだとgraspingなど細かい動作が必要なタスクでの性能が向上 • （松嶋）CNNだから割と当たり前な気もするが・・・ • CNNのでも位置情報維持するようなものを使えば別に良い？(spatial softmaxとか？） • ただしtransformerのトークン長が長くなるので計算量増加するトレードオフあり • パッチの1辺を1/2にするとトークン長は4倍，計算量は16倍 ミニバッチのシャッフルのサイズを大きくする • データが多いのでシャッフルの仕方に工夫が必要 • 最大500kのシャッフルバッファを持てるようにした 19

（まだ）うまくいっていないこと① Action headにMSEを用いる • DIffusionではなくL2で行動のlossをとると，動きがとても遅くなった（hedging policy） • 例）Bridgeの環境では手先の回転に失敗 Action headの出力を離散化する（離散的な行動空間にする） • 行動空間を次元ごとに256分割して離散化して．cross-entropyでlossをとると， 決定的な方策になり（decisive policy），たまにgraspingが早すぎて失敗 ResNet Encoderを使う • トークン数が少なくなるので学習が早くなるが，zero-shot性能が下がった 事前学習されたEncoderを使う • ImageNetで事前学習されたResNetを使っても性能向上しなかった 20

（まだ）うまくいっていないこと② グリッパの行動を相対的に表現する • グリッパの状態（開いている：+1，閉じている：-1）ではなくグリッパの変化 （開ける:＋1，閉じる：-1，その他：0）で表現すると性能が下がった • 開けたり閉じたりする動作が減り，失敗時にリトライする動作も少なくなった • ほとんどのデータが変化なし（0）だからか Proprioceptionを入力に追加する • 性能が悪くなる傾向あり • ターゲットの行動とproprioceptionの入力の相関が高くcausal confusionが起こるため か 言語エンコーダのFinetuning • 言語エンコーダ（T5）のfinetuningをしても性能向上が見られなかった • 言語のアノテーションにそこまで多様性がないためか 21

考察 ハンドアイの設定での性能が低くなった • OXEのデータセットに少ないため，ハンドアイ用のエンコーダが undertrainingだからか？ • OXEの中に27％しかハンドアイのデータが含まれていない • 3rd parson viewのエンコーダと重み共有とかすると良くなるかも Goal-conditionalよりも言語条件付けの方が性能が低い • Goal-conditionalタスクはゴールのリラベリングでデータをたくさ ん作れるが，OXEに言語条件付けされたタスクは56%しかない • 言語とGoal画像のcross attentionをとるといいかもしれない 22

（松嶋の）感想・コメント （そもそもロボット基盤モデル全般の問題として） 「（寝室にいるときに）キッチンの冷蔵庫から水を取ってきて」 のような空間的推論が必要なモバイルマニピュレーションの タスクはそのままでは解けないのでは • 記憶必要なタスクにどうやってOctoのようなE2E方策を拡張す る？ オープンソースになったのはとてもありがたい • Finetuningして活用した研究がめちゃめちゃ増える気がする • というか当たり前のように使われる気がする 23

おしらせ 英文誌Advanced Roboticsの特集号 • 日本ロボット学会の英文誌 • Google・Metaなど海外で基盤モデルのロボット応用を研究する 第一人者（Andy Zeng, Chris Paxton）たちも本特集号のエディタ 陣を務める • サーベイ論文も歓迎 • 原稿〆切：2024/1/31 2024/2/29 （延長しました） • ただし，採録が決定次第，順次web上で公開予定（プレプリント公開 も規定のもとOK） 論文投稿をぜひご検討ください！！ 24

おしらせ ICRA2024＠横浜で料理ロボットのワークショップ • 東大・エジンバラ大の若手研究者・学生5名で1年間みっちり企画 • 料理タスクは，スケーリングの次に解くべき技術課題を 数多く含むという仮説に基づき企画 • CoRL2023でも料理系のタスクの研究が急増 • 複数の実機デモを用意・コンペとの共催も予定 • 投稿〆切：2024/3/31 • https://sites.google.com/view/icra2024cookingrobotics/ 投稿・参加をぜひご検討ください！！ 25

• https://sites.google.com/view/icra2024cookingrobotics/