arXiv	https://arxiv.org/abs/2412.12953
論文のライセンス	http://creativecommons.org/licenses/by/4.0/

Efficient Diffusion Transformer Policies with Mixture of Expert Denoisers for Multitask Learning

Moritz Reuss^1,∗, Jyothish Pari^2,∗, Pulkit Agrawal², Rudolf Lioutikov¹
¹Intuitive Robots Lab (IRL), Karlsruhe Institute of Technology, Germany
²Department of Electrical Engineering and Computer Science (EECS), MIT, USA

Abstract

拡散方策は模倣学習において広く使用されるようになり、多峰性や不連続な振る舞いの生成など、いくつかの魅力的な特性を提供している。モデルがより複雑な能力を捉えるために大規模化するにつれ、近年のスケーリング則が示すように、その計算需要は増大している。したがって、現在のアーキテクチャを継続することは計算上の障害となるであろう。この課題に対処するため、我々は模倣学習のための新しい方策として、Mixture-of-Denoising Experts (MoDE)を提案する。 MoDEは、現在の最先端のTransformerベースの拡散方策を凌駕しつつ、スパースな専門家とノイズ条件付きルーティングを通じてパラメータ効率の良いスケーリングを可能にし、専門家のキャッシングにより活性パラメータを40%、推論コストを90%削減する。我々のアーキテクチャは、この効率的なスケーリングとノイズ条件付き自己注意機構を組み合わせ、異なるノイズレベルにわたってより効果的なデノイジングを可能にする。 MoDEは、4つの確立された模倣学習ベンチマーク（CALVINとLIBERO）の134タスクで最先端の性能を達成する。特筆すべきは、多様なロボティクスデータでMoDEを事前学習することで、CALVIN ABCで4.01、LIBERO-90で0.95を達成したことである。これは、CNNベースおよびTransformer拡散方策の両方を、4つのベンチマークにわたって平均 $57\%$ 上回り、デフォルトの拡散Transformerアーキテクチャと比較して90%少ないFLOPsと少ない活性パラメータを使用している。さらに、我々はMoDEのコンポーネントに関する包括的なアブレーション実験を行い、拡散方策のための効率的でスケーラブルなTransformerアーキテクチャの設計に関する洞察を提供する。コードとデモはhttps://mbreuss.github.io/MoDE_Diffusion_Policy/で入手可能である。

1 Introduction

拡散モデルは、データサンプルにガウスノイズを加える反復プロセスを逆転させることを学習する(Ho et al., 2020; Song et al., 2020)。学習後、指示や画像などの目標に条件付けられた新しいサンプルを生成することができる。最近、拡散モデルは模倣学習（IL）のポリシーとして広く採用されるようになった(Octo Model Team et al., 2023; Reuss et al., 2023; Chi et al., 2023)。 ILは、エキスパートのデモンストレーションからエージェントを訓練し、多様なスキルを学習させるための強力なパラダイムである(Pomerleau, 1989; Nair et al., 2017; Pari et al., 2021; Fu et al., 2024)。

拡散ポリシーはILにとって魅力的な特性をいくつか提供する：多様なマルチモーダルな行動を生成でき(Jia et al., 2024)、より多くのデータでスケールし(Octo Model Team et al., 2023)、アクション空間の不連続性を扱うことができる(Chi et al., 2023)。しかし、主要な制限は高い計算コストであり、モデルが大きくなるにつれて訓練と推論の速度が遅くなる。標準的なアーキテクチャは数億のパラメータを含み(Chi et al., 2023)、アクションを生成するために多くのノイズ除去ステップを必要とする。画像やテキストのための大規模なエンコーダモジュールは、ILポリシーの計算要件をさらに増加させる。これにより、リアルタイムのロボティクスアプリケーション、特にモバイルロボットのような搭載コンピューティングリソースが限られたシナリオでの使用が制限される。

これらの課題に対処するため、我々は訓練と推論のためのFLOPsを少なくしながらモデル容量をスケールできる混合エキスパート（MoE）を探求する。スパースなMoEの背後にある核心的なアイデアは、各フォワードパスで総モデルパラメータの一部のみを利用することである。これは、複数のエキスパートサブネットワークとルーティングモデルを持ち、入力に基づいてエキスパートをスパースに活性化し、その出力を補間することで達成される。

本稿では、スケーラブルで効率的なMoE拡散ポリシーであるMixture-of-Denoising Experts Policy (MoDE)を紹介する。

Refer to caption — 図1: 提案されたMoDEアーキテクチャ（左）は、因果マスキングを用いたトランスフォーマーを使用しており、各ブロックにはノイズ条件付き自己注意機構とノイズレベルに基づいてトークンを専門モデルに割り当てるノイズ条件付きルーターが含まれる。この設計により、効率的でスケーラブルな行動生成が可能となる。右側では、ノイズ除去中にSwish-GLU活性化を持つ単純なMLP専門家のサブセットをルーターが活性化する様子が示されている。

我々の研究は、ノイズ除去プロセスの異なる段階間でほとんど転移が起こらないという、ノイズ除去プロセスのマルチタスク性を示す先行研究の結果に触発されている(Hang et al., 2024)。我々は、現在のノイズレベルに基づいてトークンを専門家に分配する新しいノイズ条件付きルーティングメカニズムを提案する。 MoDEは、ノイズ入力トークンと組み合わせたノイズ条件付き自己注意機構を活用し、ノイズ注入を強化している。我々が提案するポリシーは、より高い効率性で以前のDiffusion Policyを凌駕し、挑戦的な目標条件付き模倣学習ベンチマークであるCALVIN (Mees et al., 2022b)とLIBERO (Liu et al., 2023)の134の多様なタスクにわたって最先端の性能を示している。包括的なアブレーション研究を通じて、我々はトークンルーティング戦略、ノイズ注入技術、専門家分布、大規模ロボットデータセットCollaboration et al. (2023)での多様な事前学習など、様々な設計決定の影響を調査している。我々の貢献を以下に要約する：

•

我々はMoDEを導入する。これは新しい専門家混合Diffusion Policyであり、ノイズベースの専門家キャッシングとスパースなMoE設計により、密なトランスフォーマーベースラインと比較して90 $\%$ 少ないFLOPsと少ない活性パラメータを使用しながら、最先端の性能を達成する。
•

我々は4つのベンチマークにおける134のタスクにわたってMoDEの有効性を実証し、計算効率を向上させながら、以前のDiffusion Policyと比較して平均 $57\%$ の性能向上を示す。
•

我々は、ルーティング戦略とノイズ注入の重要性を調査する詳細なアブレーション研究を提示し、ノイズ除去ステップ全体にわたる専門家の利用を可視化してMoDEの主要コンポーネントを特定する。

2 Related Work

ロボティクスにおける拡散モデル。近年、拡散モデル(Song & Ermon, 2019; Ho et al., 2020; Karras et al., 2022)はロボティクスの文脈で広く採用されるようになった。これらは模倣学習のための方策表現として使用され(Chi et al., 2023; Reuss et al., 2024; Xian et al., 2023; Ke et al., 2024; Li et al., 2023b; Scheikl et al., 2023)、オフライン強化学習においても応用されている(Ajay et al., 2023a; Janner et al., 2022; Pari et al., 2022)。ロボティクスにおける拡散モデルの他の応用には、ロボット設計生成(Wang et al., 2023)、ビデオ生成(Du et al., 2023; Ko et al., 2023; Ajay et al., 2023b)、モーションプランニング(Carvalho et al., 2023; Urain et al., 2023)がある。ロボティクスにおいて拡散モデルを方策として使用する最も一般的なアーキテクチャは、畳み込みニューラルネットワーク（CNN）に追加のFiLM条件付け(Perez et al., 2018)を加えて、文脈情報に基づいて生成を誘導するものである。最近では、トランスフォーマーアーキテクチャが拡散方策のための強力な代替バックボーンとして採用されており、特にILにおいて顕著である。例としては、Octo(Octo Model Team et al., 2023)、BESO(Reuss et al., 2023)、3D-Diffusion-Actor(Ke et al., 2024)がある。 Sparse-DPはMoEを活用して異なるタスクに対する専門家を学習する(Wang et al., 2024)。しかし、汎用方策のための除ノイズ過程の計算効率と推論時間を改善するために専門家混合アーキテクチャを使用することを考慮した先行研究は存在しない。

Mixture-of-Experts（専門家の混合）。 MoEは、情報がモデル内を選択的に経路付けされるモデルの一種である。現代版のMoEは(Shazeer et al., 2017)で導入され、ルーティングまたはゲーティングネットワークが条件付きで入力を送信する専門家のサブセットを選択する。 Transformers (Vaswani et al., 2017)がデータとともに効果的にスケールするモデルであることが証明された後、(Fedus et al., 2022)でSwitch Transformersが提案され、モデルの各ブロックに専門家フィードフォワードネットワークを持つように修正された。 Switch Transformersは、現在も様々な大規模言語モデル（LLM） (Jiang et al., 2024; Du et al., 2022)で広く採用されている基礎を築いた。これにより、前方および後方のFLOPを密なカウンターパートよりも小さく保ちながら、より多くの総パラメータを可能にし、大幅な性能向上をもたらした。しかし、ルーターと専門家を並行して訓練することは自明でない最適化問題であり、専門家が特化する代わりに類似の機能を学習する専門家の崩壊など、最適でない解をもたらす可能性がある(Chi et al., 2022)。さらに、ルーターが専門家の小さなサブセットを選択し、すべての専門家を活用しない場合、ルーターの崩壊が発生する。これは負荷分散損失(Shazeer et al., 2017; Fedus et al., 2022)によって緩和され、ルーターがより均等に専門家間で入力を分散させることを促進する。複数の研究が、専門家選択ルーティング(Zhou et al., 2022)、微分k選択(Hazimeh et al., 2021)、凍結ハッシュ関数(Roller et al., 2021)、線形割り当て(Lewis et al., 2021)など、ルーティングを実行するための異なる方法を探求している。

拡散モデルにおけるマルチタスク学習ノイズ除去プロセスがマルチタスクであることが示されているHang et al. (2024)。この考えを活用し、マルチタスク学習に適したアーキテクチャを採用した研究が行われている。一部の研究では、ノイズ除去プロセスのどの段階にどのパラメータが特化しているかを明示的にスケジューリングしている(Park et al., 2023; Go et al., 2023)。これを拡張して、(Park et al., 2024)では、トレーニング中のガイダンスとしてスケジューリングを使用するだけでなく、ノイズ除去段階に基づいて表現を変調する方法も学習している。最後に、一部の研究では、ノイズ除去段階に応じて異なるアーキテクチャを採用している(Lee et al., 2024b)。

3 Method

本節では、我々の新しいMoE拡散政策であるMoDEを紹介する。まず、ILのための政策学習問題を定式化する。次に、拡散政策で使用されるフレームワークを要約し、その後、効率的な政策設計を可能にする我々の新しいノイズ条件付きルーティングとノイズ条件付き自己注意機構を備えたMoDEアーキテクチャを紹介する。最後に、効率的な推論のための我々のエキスパートキャッシング機構とMoDEの事前学習について説明する。

3.1 Problem Formulation

我々は、ロボットのデモンストレーションのデータセット $\mathcal{T}$ が与えられた場合の言語条件付き政策 $\pi_{\theta}(\bm{\bar{a}}|\bm{\bar{s}},\bm{g})$ の学習問題を考える。この政策は、長さ $h$ の状態埋め込みの履歴 $\bm{\bar{s}}=(\bm{s}_{i-h+1},\ldots,\bm{s}_{i})$ と望ましい目標 $\bm{g}$ を条件として、長さ $j$ の将来の行動の系列 $\bm{\bar{a}}=(\bm{a},\ldots,\bm{a}_{i+j-1})$ を予測する。データセットには $\tau\in\mathcal{T}$ 個の軌道が含まれており、各軌道は状態、行動、目標の三つ組の系列 $(\bm{\bar{s}}_{i},\bm{a}_{i},\bm{g})$ で構成される。 $\bm{g}$ は言語指示である。我々の政策は、状態履歴と目標の文脈が与えられた場合の行動系列の対数尤度を最大化するように訓練される：

\displaystyle\mathcal{L}_{\text{IL}}=\mathbb{E}\left[\sum_{(\bm{\bar{s}},\bm{% \bar{a}},\bm{g})\in\mathcal{T}}\log\pi_{\theta}\left(\bm{\bar{a}}|\bm{\bar{s}}% ,\bm{g}\right)\right].

(1)

3.2 Diffusion Policy

MoDEは、EDM (Karras et al., 2022)の連続時間拡散モデルを政策表現として使用する。拡散モデルは、初めにガウス摂動によってノイズを加え、その後このプロセスを逆転させることでデータを生成する一種の生成モデルである。 MoDEは、スコアベースの拡散モデルを適用して政策 $\pi_{\theta}(\bm{\bar{a}}|\bm{\bar{s}},\bm{g})$ を表現する。摂動と逆プロセスは、確率微分方程式を用いて以下のように記述できる：

\begin{split}\mathop{}\!\mathrm{d}\bm{\bar{a}}=\big{(}\beta_{t}\sigma_{t}-\dot% {\sigma}_{t}\big{)}\sigma_{t}\nabla_{a}\log p_{t}(\bm{\bar{a}}|\bm{\bar{s}},% \bm{g})dt+\sqrt{2\beta_{t}}\sigma_{t}d\omega_{t},\end{split}

(2)

ここで、 $\beta_{t}$ はノイズ注入を制御し、 $d\omega_{t}$ は無限小のガウスノイズを指し、 $p_{t}(\bm{\bar{a}}|\bm{\bar{s}},\bm{g})$ は拡散プロセスのスコア関数であり、順方向プロセスにおいてサンプルをデータ密度の高い領域から遠ざける。ノイズから新しいサンプルを生成するために、ニューラルネットワークはスコアマッチング（SM）(Vincent, 2011)を通じてスコア関数 $\nabla_{\bm{\bar{a}}}\log p_{t}(\bm{\bar{a}}|\bm{\bar{s}},\bm{g})$ を近似するよう訓練される。

\mathcal{L}_{\text{SM}}=\mathbb{E}_{\mathbf{\sigma},\bm{\bar{a}},\bm{\epsilon}% }\big{[}\alpha(\sigma_{t})\newline \|D_{\theta}(\bm{\bar{a}}+\bm{\epsilon},\bm{\bar{s}},\bm{g},\sigma_{t})-\bm{% \bar{a}}\|_{2}^{2}\big{]},

(3)

ここで、 $D_{\theta}(\bm{\bar{a}}+\bm{\epsilon},\bm{\bar{s}},\bm{g},\sigma_{t})$ は訓練可能なニューラルネットワークである。訓練中、我々は訓練分布からノイズをサンプリングし、それをアクション系列に加える。ネットワークはノイズ除去されたアクションを予測し、SMロスを計算する。

訓練後、我々はランダムノイズから始めて、数値ODE積分器を用いて逆SDEまたは関連するODEを離散的なステップで近似することで、新しいアクション系列を生成できる。したがって、我々は事前分布 $\bm{a}_{T}\sim\mathcal{N}(\mathbf{0},\sigma_{T}^{2}\mathbf{I})$ からノイズをサンプリングし、それを反復的にノイズ除去する。 MoDEは、拡散モデル用に設計された数値ODE解法であるDDIMソルバー(Song et al., 2021)を使用する。これにより、少ないステップでアクションの高速なノイズ除去が可能となる。 MoDEは、我々のすべての実験において $10$ ノイズ除去ステップを使用する。

3.3 Mixture-of-Experts Denoising

我々は今、MoDEを紹介する。これは、ノイズ条件付き専門家ルーティングを用いて拡散ベースのポリシーを強化する新しいアプローチである。この新規ルーティングメカニズムにより、必要な専門家を事前に計算し融合することで、より効率的な推論が可能となる。 MoDEのアーキテクチャとルーティングメカニズムの概要を図1に示す。

言語条件付けについて、MoDEは凍結されたCLIP言語エンコーダーを活用して潜在的な目標ベクトルを生成し、画像エンコーディングにはFiLM条件付きResNets-18/50を使用する。モデルは入力トークンの系列 $\bf{X}\in\mathbb{R}^{\text{tokens}\times D}$ とノイズレベル $\sigma_{t}$ を処理する。線形射影層 $\phi(\sigma_{t})$ がノイズレベルをトークンにエンコードし、これを $\bf{X}$ に組み込む。完全なMoDEアーキテクチャ $\text{MoDE}(\bf{X},\phi(\sigma_{t}))$ は、 $L$ 個のトランスフォーマーブロックで構成され、各ブロックは異なるノイズ除去フェーズに特化している。

ここで、各ブロック $f^{i}$ を自己注意（SA）層とMoE層の組み合わせとして定義する。

f^{i}(\bf{X},\phi(\sigma_{t}))=\text{MoE}(\text{SA}(\hat{\bf{X}})+\bf{X},\phi(% \sigma_{t}))+\bf{X}.

(4)

我々のアプローチにおける重要な変更点は、ノイズを考慮した位置埋め込みの統合である。自己注意の前に $\phi(\sigma_{t})$ を $\bf{X}$ のすべてのトークンに加えることで：

\hat{\bf{X}}=\phi(\sigma_{t})+\bf{X},

(5)

各トークンが現在のノイズ除去フェーズに基づいて注意パターンを適応させることが可能になる。この設計により、追加のパラメータやアーキテクチャの複雑さを導入することなく、ノイズ除去性能が向上する。

自己注意メカニズムは標準的な定式化に従う(Vaswani et al., 2017)：

\text{SA}(\hat{\bf{X}})=\text{softmax}(\frac{1}{\sqrt{D}}[\hat{\bf{X}}W_{Q}][% \hat{\bf{X}}W_{K}]^{T})[\hat{\bf{X}}W_{V}].

(6)

我々のMoE層は、新しいノイズ条件付きルーティング戦略を導入する。 $N$ 個の専門家 ${\bf{E_{i}}}{i=1}^{N}$ が与えられた場合、層の出力は：

\text{MoE}(\bf{X},\phi(\sigma_{t}))=\sum{i=1}^{N}\bf{R}(\phi(\sigma_{t}))\bf{E% }_{i}(\bf{X}),

(7)

ここで、ルーティング関数 $\bf{R}(\cdot):\mathbb{R}^{\text{tokens}\times D}\rightarrow\mathbb{R}^{\text{% tokens}\times N}$ が専門家の選択を決定する：

\bf{R}(\phi(\sigma_{t}))=\text{topk}(\text{softmax}(\phi(\sigma_{t})W_{R}),% \textit{k})

(8)

入力内容に基づいてルーティングを行う従来のMoEアプローチとは異なり、MoDEのルーティングメカニズムは特にノイズレベルを考慮する。これにより、異なるノイズ除去フェーズに特化した専門家が可能となり、性能の向上と専門家のキャッシングを通じた計算効率の向上が実現する（詳細は3.3.1項参照）。我々は(Muennighoff et al., 2024)と同じ方法を用いてルーターを初期化する。これは $\text{std}=0.02$ の切断正規分布からのものである。 topkの実装には、置換なしの多項分布サンプリングを使用し、 $\text{softmax}(\phi(\sigma_{t})W_{R})$ の確率に従ってk個の要素を選択する。微分不可能なサンプリングプロセスを通じて勾配の流れを維持するために、専門家の出力をルーティング確率でスケーリングし、選択された確率を再正規化する。専門家の崩壊を防ぐために、負荷分散損失（ $LB$ ）を組み込む(Fedus et al., 2022)：

LB(\sigma_{t})=N\sum_{n=1}^{N}\frac{1}{|\mathcal{B}|}(\sum_{i=1}^{|\mathcal{B}% |}\mathbbm{1}{\bf{R}(\phi(\sigma_{t_{i}}))n>0})\frac{1}{|\mathcal{B}|}(\sum{i=% 1}^{|\mathcal{B}|}\text{softmax}(\phi(\sigma_{t_{i}})W_{R})_{n})

(9)

ここで、 $\gamma=0.01$ はバランシング係数である。

3.3.1 Router and Expert Caching

我々の手法をより効率的にするために、MoEがノイズ条件付きであるという事実を利用する。つまり、各ノイズレベルにおいて、ルーティングパスは決定論的であり、事前に計算できる。これにより、各ノイズレベルに対して選択される専門家を事前に決定することができる。図2はこのプロセスを示している。これにより、選択された専門家MLPを単一の複合MLPに融合し、計算コストを効果的に削減することができる。各専門家を個別にループする代わりに、この融合された専門家MLPは並列で実行でき、ネットワークの全体的な遅延を大幅に減少させる。各専門家を動的に呼び出す必要性を排除することで、時間を節約するだけでなく、メモリアクセスパターンを効率化し、従来のMoEセットアップに通常伴うオーバーヘッドを削減する。我々のキャッシング戦略は、標準的なMoE展開と比較してFLOPsのオーバーヘッドを80 $\%$ 以上削減し、推論時に2倍高速化する。

3.4 Generalist Policy Pre-Training

我々は、OXEデータセットCollaboration et al. (2023)から収集した多様なマルチロボットデータセットの混合物でMoDEを事前学習する。我々の学習データは、様々なロボットプラットフォームと操作タスクを特徴とする6つの多様なデータセットから選択された196,000の軌跡を含んでいる。 MoDEの事前学習プロセスは、6台のGPUを搭載した分散クラスターで3日間にわたり300,000ステップ実行される。ファインチューニングでは、各層の事前学習されたルーターを凍結し、他のモデルコンポーネントのみをファインチューニングする。我々の事前学習データセットの構成と方法論の包括的な概要は付録（サブセクションA.1）に記載されている。

real2simベンチマークSIMPLER(Li et al., 2024b)を用いた詳細な評価において、MoDEは最先端の汎用的方策と比較して優れた性能を示している。多様な操作タスクにおいて平均成功率26.30%を達成し、OpenVLA（23.70%）とOcto（17.75%）の両方を上回っている。完全な評価の詳細はサブサブセクションA.2.1に記載されている。

4 Evaluation

我々の実験は、以下の4つの主要な問いに答えることを目的としている： (I) MoDEは、性能の面で他の方策や先行するディフュージョントランスフォーマーアーキテクチャとどのように比較されるか？ (II) 多様なロボティクスデータの大規模な事前学習は、MoDEの性能を向上させるか？ (III) MoDEの効率性と速度は、密なトランスフォーマーのベースラインと比較してどうか？ (IV) 状態ベースおよび画像ベースの環境において、ディフュージョン方策に最も効果的なトークンルーティング戦略はどれか？ (V) モデルは、行動のデノイジングプロセス中にどのように異なる専門家を活用するか？

我々は、MoDEを先行するディフュージョントランスフォーマーアーキテクチャ(Chi et al., 2023)と比較し、アクティブなパラメータ数を同程度にすることで公平な比較を確保している。 MoDEは、すべての実験において8層、4専門家、潜在次元1024を使用している。我々の事前学習済みバリアントは、12層、4専門家、同じ潜在次元1024でわずかに大きくなっている。

我々は、すべての実験において行動チャンキング長を10、履歴長を1としている。 MoDEは、早期の再計画や時間的集約なしに、生成された10個の行動をすべて実行する。詳細なハイパーパラメータは付録（表3）に記載されている。

4.1 Long-Horizon Multi-Task Experiments

我々はまず、LIBEROベンチマーク(Liu et al., 2023)のLONGチャレンジとLIBERO-90チャレンジでMoDEを評価する。LONGチャレンジは、異なる設定で10のタスクを学習することをモデルに要求する。これは、完了までに数百ステップを要する長期的な行動生成を必要とする。90バリアントは、90の多様な短期的タスクを異なる環境でポリシーをテストする。図3(a)はこれらのタスクの例をいくつか可視化している。すべての環境には2つのカメラがある：静的なものと手首に取り付けられたカメラで、これらを使用してFiLM-ResNets-18で現在の観測をエンコードする。我々は各ポリシーを各タスクで20回テストし、3シードの平均結果を報告する。 MoDEおよび他のすべての拡散アーキテクチャは、目標と画像をエンコードするためにCLIP文埋め込みを持つFiLM条件付きResNets-18を使用する。

ベースライン。我々はMoDEを3つの最先端のベースラインと比較する： 1) Diffusion Transformer (DP-T)アーキテクチャ(Chi et al., 2023)。これはクロスアテンションモジュールを使用してノイズと観測に条件付けを行う。 2) 標準的なDiffusion Policy CNN-ベースのアーキテクチャ(DP-CNN)。 3) QueST (Mete et al., 2024)。これは、アクションシーケンスのベクトル量子化埋め込みを使用して離散的なアクション表現を学習するトランスフォーマーベースのポリシーである。我々はQueST以外のすべてのベースラインを自身でテストし、QueSTの結果は論文から直接引用した。

結果。ベンチマークにおけるすべてのモデルの性能は3(b)にまとめられている。全体として、MoDEは両方のベンチマークで最高の平均性能を達成し、一方でQueSTベースラインはLIBERO-90設定で2番目に良く、CNNアーキテクチャは長期的設定で2番目に良い結果となった。これらの結果は、MoDEが高い精度で長期的タスクを学習する能力を示している。性能差は、より困難なLIBERO-10実験でより顕著であり、MoDEは $90\%$ を超える成功率を達成した最初のポリシーである。さらに、MoDEは両設定において、計算上の利点を維持しながら、以前の最良の拡散ベースラインを平均 $16\%$ 上回っている。事前学習されたMoDEバリアントは両設定でさらに高い性能を達成し、多様な事前学習の可能性を示している。これは、MoDEが計算リソースをより効率的に使用しながら最先端の性能を達成する能力を示している。

4.2 Scaling Multi-Task Experiments

Train $\rightarrow$ Test	Method	Active Params	PrT	No. Instructions in a Row (1000 chains)
Train $\rightarrow$ Test	Method	in Million	PrT	1	2	3	4	5	Avg. Len.
ABCD $\rightarrow$ D	Diff-P-CNN	321	$\times$	86.3%	72.7%	60.1%	51.2%	41.7%	3.16 $\pm$ 0.06
	Diff-P-T	194	$\times$	78.3%	53.9%	33.8%	20.4%	11.3%	1.98 $\pm$ 0.09
	RoboFlamingo	1000	$\checkmark$	96.4%	89.6%	82.4%	74.0%	66.0%	4.09 $\pm$ 0.00
	GR-1	130	$\checkmark$	94.9%	89.6%	84.4%	78.9%	73.1%	4.21 $\pm$ 0.00
	MoDE	277	$\times$	96.6%	90.6%	86.6%	80.9%	75.5%	4.30 $\pm$ 0.02
	MoDE	436	$\checkmark$	97.1%	92.5%	87.9%	83.5%	77.9%	4.39 $\pm$ 0.04
ABC $\rightarrow$ D	Diff-P-CNN	321	$\times$	63.5%	35.3%	19.4%	10.7%	6.4%	1.35 $\pm$ 0.05
	Diff-P-T	194	$\times$	62.2%	30.9%	13.2%	5.0%	1.6%	1.13 $\pm$ 0.02
	RoboFlamingo	1000	$\checkmark$	82.4%	61.9%	46.6%	33.1%	23.5%	2.47 $\pm$ 0.00
	SuSIE	860+	$\checkmark$	87.0%	69.0%	49.0%	38.0%	26.0%	2.69 $\pm$ 0.00
	GR-1	130	$\checkmark$	85.4%	71.2%	59.6%	49.7%	40.1%	3.06 $\pm$ 0.00
	MoDE	307	$\times$	91.5%	79.2%	67.3%	55.8%	45.3%	3.39 $\pm$ 0.03
	MoDE	436	$\checkmark$	96.2%	88.9%	81.1%	71.8%	63.5%	4.01 $\pm$ 0.04

表1: 2つのCALVINチャレンジにおけるパフォーマンス比較。表は、指示チェーン内の個々のタスクの平均成功率と、5つの連続した指示を完了するための平均ロールアウト長（Avg. Len.）を1000チェーンに基づいて報告している。標準偏差がゼロの場合は、平均パフォーマンスが報告されていない手法を示す。 "Prt"はポリシーの事前学習を必要とするモデルを示す。パラメータ数には言語エンコーダーは含まれない。

次に、我々は要求の厳しいCALVIN言語スキルベンチマーク (Mees et al., 2022b)においてMoDEの有効性を評価する。これは、ILのための確立された画像ベースのベンチマークである。このベンチマークには、人間が記録したデモンストレーションの大規模なデータセットが含まれている。まず、MoDEはABCD $\rightarrow$ Dチャレンジでテストされる。これは4つの環境（A, B, C, D）にわたる $22,966$ の相互作用シーケンスを含み、各シーケンスは $64$ タイムステップと $34$ の多様なタスクで構成される。これらのタスクは、複雑な連続的行動の獲得と、異なるスキルを連鎖させる能力を必要とする。 4(a)は、これらの環境内のインタラクティブな要素の多様な構成を示している。この特定のチャレンジは、複数の設定にわたる豊富な種類のデータとスキルで訓練されたポリシーのスケーリング能力を検証する。すべてのポリシーは、CALVINの公式プロトコル(Mees et al., 2022b)に従って、環境Dにおいて連続した $5$ タスクからなる $1000$ の指示チェーンでテストされる。 $5$ の異なるタスクを含む1つの例示的ロールアウトが4(b)に視覚化されている。スコアリングに関しては、モデルはタスクを完了するごとに $1$ ポイントを獲得し、前のタスクを完了した場合にのみ次のタスクに進む。我々は、 $1000$ の指示チェーンを各々含む $3$ シードにわたる平均シーケンス長を報告する。

ベースライン。我々はMODEを、言語条件付き行動の学習に特化した複数の手法と、他のベースラインの拡散ポリシーアーキテクチャと比較テストする。また、MoDEをRoboFlamingoとGR-1と比較する。 RoboFlamingoは、約 $3$ 億のパラメータを含み、多様なインターネットデータで事前学習された、微調整されたVision-Language-Actionモデルである。 GR-1は、大規模なビデオ生成で事前学習され、その後CALVINで共同微調整された生成的デコーダーのみのTransformerである(Wu et al., 2024)。可能な場合、我々はCALVINの標準評価プロトコル(Mees et al., 2022a)に基づいて、すべての先行研究の平均パフォーマンスを直接その論文から報告する。

結果。表1に概説された我々の発見は、MoDEが平均成功率の点ですべての他のポリシーを上回ることを示している。さらに、事前学習なしのMoDEは、広範なインターネットスケールの事前学習に依存するRoboFlamingoやGR-1のような確立されたベースラインを上回るパフォーマンスを示す。我々のより大規模な事前学習版はさらに高いパフォーマンスを達成する。特筆すべきは、GR-1がMoDEの277Mと比較して少ない活性パラメータ（130M）を使用しているが、各タイムステップに10および15トークンの履歴長を使用し、画像埋め込みに事前学習されたViTを使用している点である。 MoDEはより計算効率が高く、推論時に少ないFLOP（GR-1の27.5 GFLOPSに対して1.53 GFLOPS）を必要とし、6倍以上大きいにもかかわらず同等に高速である（12.6 msに対して12.2 ms）。最先端のパフォーマンスと低い計算要求の組み合わせにより、MoDEはマルチタスク設定において非常に実用的なソリューションとして位置付けられる。

4.3 Zero-shot Generalization Experiments

最後に、我々は第二段階でABC $\rightarrow$ Dチャレンジへと調査を拡張し、MoDEのゼロショット汎化能力をテストする。この実験では、モデルは最初の3つのCALVIN環境A,B,Cのデータでのみ訓練され、未見の環境Dでテストされる。環境Dでは、関連オブジェクトの位置やテーブルのテクスチャが異なる。これには、学習した行動を新しい環境構成や異なるテクスチャに汎化できる方策が必要であり、特に困難な課題である。我々は、ゼロから訓練したMoDEと、Open-X-Embodimentデータのサブセットで事前訓練したMoDEを評価する。これにより、MoDEのゼロショット性能と事前訓練の有効性を研究することが可能となる。

ベースライン。この実験では、我々はMODEを以前のCALVINベースラインと比較し、さらにSuSIE (Black et al., 2023)を追加する。 SuSIEは、ゴール画像を生成するために微調整された画像生成モデルであるInstruct2Pix (Brooks et al., 2023)を利用する階層的方策であり、生成されたゴール画像が低レベルの拡散方策を導く。高レベルのゴール生成モデルは大規模な事前訓練を必要とする。 SuSIEの結果は、サブゴール画像生成の計算コストのため、標準偏差なしで $100$ 回のロールアウトのみに基づいている。

結果。この実験の結果は表1の下部にまとめられている。 MoDEはテストされたすべてのベースラインを上回り、他のすべての拡散方策アーキテクチャを大きく凌駕している。さらに、MoDEを多様なロボティクスデータで事前訓練することで、4.01という新しい最高性能を達成した。したがって、質問(I)に対して、我々は混合エキスパートモデルがスケーリング性能を向上させるだけでなく、強力なゼロショット汎化能力も示すと肯定的に結論付ける。加えて、質問(II)に対しては、事前訓練が困難なゼロショット設定での性能を向上させると結論付けることができる。

4.4 Computational Efficiency of MoDE

我々は、MoDEを同様のパラメータを持つ密なトランスフォーマーのベースラインと比較し、バッチサイズ全体での平均推論時間とFLOPを測定した。図5に示すように、キャッシングを用いたMoDEは計算効率を大幅に向上させている - バッチサイズ1では推論が20%高速化し（12ms対15ms）、バッチサイズ512では、MoDEは16倍少ないFLOP（361対5,772 GFLOPS）を必要とし、約40%高速な推論（64ms対104ms）を達成している。これらの結果は、MoDEがその構造とキャッシングメカニズムを通じて、優れたタスクパフォーマンスと大幅な計算効率の両方を提供していることを示している。 CALVINにおける他のすべてのベースラインとの推論速度とFLOPSの詳細な比較は、付録A.4にまとめられている。

4.5 Ablation Studies

MoDEの設計選択を徹底的に評価するため、我々は一連のアブレーション実験を実施した。これらの実験は、我々の研究課題に取り組むものである：MoDEの計算効率（課題III）、ルーティング戦略の影響（課題IV）、およびトークン分布（課題V）。

4.5.1 What design decisions affect MoDE’s performance?

まず、様々な設計決定がMoDEの性能に与える影響を評価する。我々は、LIBERO-10ベンチマークにおいて、ノイズ条件付けの選択と様々なMoE戦略についてアブレーションを行った。結果は表2にまとめられている。

ノイズ注入のアブレーション。我々の実験は、MoDEにおける適切なノイズ条件付けの重要性を明らかにした。入力ノイズトークンとノイズ条件付きセルフアテンションの両方を使用する完全なMoDEモデルは、平均成功率 $0.92$ で最高の性能を達成した。入力ノイズトークンを除去すると、性能が $0.90$ にわずかに低下し、両方の条件付け方法の相補的な性質が強調された。ノイズ条件付きセルフアテンションなしで、ノイズトークンのみを条件付けに使用すると、性能がさらに $0.85$ に低下した。興味深いことに、画像拡散で一般的なアプローチであるFiLM条件付け(Perez et al., 2018) (Peebles & Xie, 2023)を使用すると、このグループで最低の性能 $0.81$ となった。これらの結果は、MoDEで提案されたノイズ条件付け戦略の有効性を裏付けており、FiLMアプローチに対して $0.08$ の明確な性能優位性を示している。

	Avg. Success.
MoDE	0.92
- Input Noise Token	0.90
- Noise-cond Attention	0.85
FiLM Noise Conditioning	0.81
TopK=1	0.91
Shared Expert	0.90
$\gamma=0.1$	0.90
$\gamma=0.001$	0.86
256 Embed Dim	0.86
512 Embed Dim	0.87

表2: LIBERO-10におけるMoDEのアブレーション実験。すべての結果は

3

シードで平均化され、各シードで

20

回のロールアウトを行っている。

MoEのアブレーション。次に、Mixture-of-Expertsに関するいくつかの設計決定についてアブレーションを行う。まず、使用する専門家のtopk数をテストする。 topkを1に設定すると、平均性能が $0.92$ から $0.91$ にわずかに低下するだけである。MoDEは単一の専門家でも堅牢な性能を維持する。我々はまた、共有専門家の使用の効果を検討する。ここでは、モデルがすべてのケースで一貫して同じ専門家を使用する。このアプローチは、 $0.90$ という同等の平均成功率を達成する。トークン分布損失の異なる選択についても検討する。 MoDEはデフォルト値として $\gamma=0.01$ を使用するが、我々は $\gamma$ の値として $0.1$ と $0.001$ を実験し、それぞれ $0.90$ と $0.86$ の平均成功率を得た。これらの結果は、 $\gamma$ の値 $0.01$ が最良の性能を示すことを示している。

潜在次元。我々はMoDEの潜在次元を変化させる影響を調査し、 $256$ 、 $512$ 、および $1024$ （デフォルト）の次元をテストした。結果は、潜在次元を $256$ から $512$ に増やすと、性能が $0.86$ から $0.87$ にわずかに向上し、さらに $1024$ に増やすと $0.92$ へとより大幅な向上が見られることを示している。これは、より大きな潜在次元によりMoDEがより複雑な表現を捉えることができ、性能向上につながることを示唆している。

4.5.2 Optimal Routing Strategy for Diffusion Transformers

次に、我々は課題（III）に答えるため、複数の環境で拡散トランスフォーマーポリシーの異なるルーティング戦略をテストする。我々は2つの異なるトークンルーティング戦略をテストする： 1）トークンのみに条件付けされたルーティング 2）ノイズのみのトークンルーティング（1）はLLMで一般的に使用され、ルーティングはトークンのみに基づいて決定される。我々は5つの実験でこれらの戦略をテストし、 $3$ シードにわたる平均性能を報告する：ノイズのみのルーティングは0.851の平均正規化性能を達成し、0.845を達成するトークンのみのルーティングをわずかに上回った。詳細な結果は付録の表7にまとめられている。結果は我々が提案するルーティング戦略の有効性を示している。性能差は小さいものの、ノイズのみのルーティングには追加の利点がある：3.3.1項で説明したように、ロールアウト前に一度、ノイズレベルに基づいて使用されるすべての専門家を予測できるため、より高速な推論が可能になる。これはロボティクスアプリケーションにとって特に有益である。

4.5.3 How does the model distribute the tokens to different experts?

課題IVに取り組むため、我々は事前学習済みモデルを使用してMoDEがトークンを異なる専門家にどのように分配するかを分析した。図6は、明確化のために10段階のデノイジングステップを使用して、様々なノイズレベルにわたる推論中の各モデル層における各専門家の平均使用率を可視化している。我々の分析は、MoDEが様々なノイズレベルに対して異なる専門家を利用することを学習していることを明らかにし、ルーターが異なるノイズ領域に特化していることを示唆している。専門家の利用における遷移が $\sigma_{8}$ 付近で発生する。最初の層では、モデルは低ノイズレベルに特化した専門家を学習し、主に $\sigma_{\text{min}}$ の最後のデノイジングステップで使用される。我々は、事前学習済みモデルとMoDEの様々な他のバージョンを用いて、付録のA.6.1項でさらなるアブレーション実験を行っている。これらの発見は課題IVに肯定的に答えるものであり、MoDEがノイズレベルに基づいて効果的に専門家間でトークンを分配していることを示している。

4.5.4 How does the model scale with more experts?

最後に、MoDEの専門家の数を増やす効果を分析する。結果は図8に示されており、2、4、6、8人の専門家を使用してCALVIN ABCDおよびCALVIN ABCベンチマークでMoDEを評価している。比較のため、2つの密なMoDEベースライン：Dense-smallとDense-largeを含めている。Dense-smallはMoDEと同じ潜在次元を共有し、Dense-largeは2024次元にスケールアップされ、MoDEの全体的なパラメータ数に一致する。我々の分析は、スケーリングが一般的な性能（C-ABCD）とゼロショット汎化（C-ABC）の両方にどのように影響するかに焦点を当てている。 ABCD環境では、4人の専門家を持つMoDEが最高の性能を達成する。興味深いことに、4人以上の専門家に増やすと性能が低下し、過学習やルーティングの複雑性の増加が原因である可能性がある。ゼロショット汎化（ABC）では、4人の専門家を持つMoDEが依然として最高の性能を示す。特筆すべきは、Dense-small変種が両方のタスクで一貫して性能が低いことであり、パラメータをより効果的に利用するMoEアーキテクチャの効率性を強調している。我々は、4人の専門家が拡散ポリシーのノイズのみのルーティングのコンテキストにおいて理想的なトレードオフを持つと仮説を立てている。図6と付録の図12で観察された異なる専門家の特化パターンは、専門家の特化がノイズ領域に基づいていることを示している。 4人以上の専門家を持つMoDEには性能上の利点がない。全体として、MoDEは密なトランスフォーマーモデルと同等またはそれ以上の性能を達成しつつ、より少ない計算リソースで実現できることを示している。

4.6 Limitations

MoDEにはまだいくつかの限界がある。我々の実験では、MoDEはベースラインと比較してわずかに高い標準偏差を示すことが分かった。我々は、ルーターの初期化が全体的な最適化に大きな影響を与えていると仮説を立てており、ルーティングモデルの安定化に関する今後の研究が必要である。さらに、エキスパートの利用状況を可視化した際、我々の実験の一部では、全エキスパートのうち一部のみが利用されているという現象が観察された。これはエキスパート崩壊として知られている現象である (Chi et al., 2022)。

5 Conclusion

本稿では、拡散政策の性能と効率を向上させるために専門家混合Transformerを活用する新しい拡散政策である、Mixture-of-Denoising Experts (MoDE)を紹介した。また、我々のモデル内で特化した専門家を学習するためのノイズ条件付きルーティング戦略を提案した。多様なベンチマークにわたる広範な実験と除去実験において、我々はMoDEが少ないパラメータ数と推論時に $90\%$ 少ないFLOPSで従来の拡散政策を上回る利点を実証した。今後の研究では、専門家選択ルーティング(Zhou et al., 2022)などのより多くのルーティング戦略を実験したいと考えている。

6 Acknowledgments

我々は、Adam Wei、Anurag Ajay、Hao-Shu Fang、Anthony Simeonov、Yilun Duの各氏による洞察に富んだ議論とフィードバックに感謝する。本研究はドイツ研究振興協会（DFG）の助成金（448648559）により支援された。著者らはまた、バーデン＝ヴュルテンベルク州科学・研究・芸術省が資金提供したHoreKaスーパーコンピュータを通じたバーデン＝ヴュルテンベルク州からの支援、および連邦教育研究省からの支援を認める。本研究はまた、陸軍研究局の後援を受け、ARO MURIグラント番号W911NF-23-1-0277のもとで遂行された。

References

Ajay et al. (2023a) Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In International Conference on Learning Representations, 2023a. URL https://openreview.net/forum?id=sP1fo2K9DFG.
Ajay et al. (2023b) Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309.08587, 2023b.
Alayrac et al. (2022) Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716–23736, 2022.
Bharadhwaj et al. (2023) Homanga Bharadhwaj, Jay Vakil, Mohit Sharma, Abhinav Gupta, Shubham Tulsiani, and Vikash Kumar. Roboagent: Generalization and efficiency in robot manipulation via semantic augmentations and action chunking, 2023.
Black et al. (2023) Kevin Black, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, and Sergey Levine. Zero-shot robotic manipulation with pretrained image-editing diffusion models. arXiv preprint arXiv:2310.10639, 2023.
Blessing et al. (2023) Denis Blessing, Onur Celik, Xiaogang Jia, Moritz Reuss, Maximilian Xiling Li, Rudolf Lioutikov, and Gerhard Neumann. Information maximizing curriculum: A curriculum-based approach for learning versatile skills. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=7eW6NzSE4g.
Brohan et al. (2022) Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong Tran, Vincent Vanhoucke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich. Rt-1: Robotics transformer for real-world control at scale. In arXiv preprint arXiv:2212.06817, 2022.
Brohan et al. (2023) Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, et al. Rt-2: Vision-language-action models transfer web knowledge to robotic control. arXiv preprint arXiv:2307.15818, 2023.
Brooks et al. (2023) Tim Brooks, Aleksander Holynski, and Alexei A. Efros. Instructpix2pix: Learning to follow image editing instructions. In CVPR, 2023.
Carvalho et al. (2023) Joao Carvalho, An T Le, Mark Baierl, Dorothea Koert, and Jan Peters. Motion planning diffusion: Learning and planning of robot motions with diffusion models. In 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 1916–1923. IEEE, 2023.
Celik et al. (2022) Onur Celik, Dongzhuoran Zhou, Ge Li, Philipp Becker, and Gerhard Neumann. Specializing versatile skill libraries using local mixture of experts. In Aleksandra Faust, David Hsu, and Gerhard Neumann (eds.), Proceedings of the 5th Conference on Robot Learning, volume 164 of Proceedings of Machine Learning Research, pp. 1423–1433. PMLR, 08–11 Nov 2022. URL https://proceedings.mlr.press/v164/celik22a.html.
Celik et al. (2024) Onur Celik, Aleksandar Taranovic, and Gerhard Neumann. Acquiring diverse skills using curriculum reinforcement learning with mixture of experts. arXiv preprint arXiv:2403.06966, 2024.
Chi et al. (2023) Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Chi et al. (2022) Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=mWaYC6CZf5.
Collaboration et al. (2023) Open X-Embodiment Collaboration, Abhishek Padalkar, Acorn Pooley, Ajinkya Jain, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Anthony Brohan, Antonin Raffin, Ayzaan Wahid, Ben Burgess-Limerick, Beomjoon Kim, Bernhard Schölkopf, Brian Ichter, Cewu Lu, Charles Xu, Chelsea Finn, Chenfeng Xu, Cheng Chi, Chenguang Huang, Christine Chan, Chuer Pan, Chuyuan Fu, Coline Devin, Danny Driess, Deepak Pathak, Dhruv Shah, Dieter Büchler, Dmitry Kalashnikov, Dorsa Sadigh, Edward Johns, Federico Ceola, Fei Xia, Freek Stulp, Gaoyue Zhou, Gaurav S. Sukhatme, Gautam Salhotra, Ge Yan, Giulio Schiavi, Hao Su, Hao-Shu Fang, Haochen Shi, Heni Ben Amor, Henrik I Christensen, Hiroki Furuta, Homer Walke, Hongjie Fang, Igor Mordatch, Ilija Radosavovic, Isabel Leal, Jacky Liang, Jaehyung Kim, Jan Schneider, Jasmine Hsu, Jeannette Bohg, Jeffrey Bingham, Jiajun Wu, Jialin Wu, Jianlan Luo, Jiayuan Gu, Jie Tan, Jihoon Oh, Jitendra Malik, Jonathan Tompson, Jonathan Yang, Joseph J. Lim, João Silvério, Junhyek Han, Kanishka Rao, Karl Pertsch, Karol Hausman, Keegan Go, Keerthana Gopalakrishnan, Ken Goldberg, Kendra Byrne, Kenneth Oslund, Kento Kawaharazuka, Kevin Zhang, Keyvan Majd, Krishan Rana, Krishnan Srinivasan, Lawrence Yunliang Chen, Lerrel Pinto, Liam Tan, Lionel Ott, Lisa Lee, Masayoshi Tomizuka, Maximilian Du, Michael Ahn, Mingtong Zhang, Mingyu Ding, Mohan Kumar Srirama, Mohit Sharma, Moo Jin Kim, Naoaki Kanazawa, Nicklas Hansen, Nicolas Heess, Nikhil J Joshi, Niko Suenderhauf, Norman Di Palo, Nur Muhammad Mahi Shafiullah, Oier Mees, Oliver Kroemer, Pannag R Sanketi, Paul Wohlhart, Peng Xu, Pierre Sermanet, Priya Sundaresan, Quan Vuong, Rafael Rafailov, Ran Tian, Ria Doshi, Roberto Martín-Martín, Russell Mendonca, Rutav Shah, Ryan Hoque, Ryan Julian, Samuel Bustamante, Sean Kirmani, Sergey Levine, Sherry Moore, Shikhar Bahl, Shivin Dass, Shuran Song, Sichun Xu, Siddhant Haldar, Simeon Adebola, Simon Guist, Soroush Nasiriany, Stefan Schaal, Stefan Welker, Stephen Tian, Sudeep Dasari, Suneel Belkhale, Takayuki Osa, Tatsuya Harada, Tatsuya Matsushima, Ted Xiao, Tianhe Yu, Tianli Ding, Todor Davchev, Tony Z. Zhao, Travis Armstrong, Trevor Darrell, Vidhi Jain, Vincent Vanhoucke, Wei Zhan, Wenxuan Zhou, Wolfram Burgard, Xi Chen, Xiaolong Wang, Xinghao Zhu, Xuanlin Li, Yao Lu, Yevgen Chebotar, Yifan Zhou, Yifeng Zhu, Ying Xu, Yixuan Wang, Yonatan Bisk, Yoonyoung Cho, Youngwoon Lee, Yuchen Cui, Yueh hua Wu, Yujin Tang, Yuke Zhu, Yunzhu Li, Yusuke Iwasawa, Yutaka Matsuo, Zhuo Xu, and Zichen Jeff Cui. Open X-Embodiment: Robotic learning datasets and RT-X models. https://arxiv.org/abs/2310.08864, 2023.
Cui et al. (2023) Zichen Jeff Cui, Yibin Wang, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. From play to policy: Conditional behavior generation from uncurated robot data. In International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=c7rM7F7jQjN.
Du et al. (2022) Nan Du, Yanping Huang, Andrew M. Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, Barret Zoph, Liam Fedus, Maarten Bosma, Zongwei Zhou, Tao Wang, Yu Emma Wang, Kellie Webster, Marie Pellat, Kevin Robinson, Kathleen Meier-Hellstern, Toju Duke, Lucas Dixon, Kun Zhang, Quoc V Le, Yonghui Wu, Zhifeng Chen, and Claire Cui. Glam: Efficient scaling of language models with mixture-of-experts, 2022.
Du et al. (2023) Yilun Du, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B Tenenbaum, Dale Schuurmans, and Pieter Abbeel. Learning universal policies via text-guided video generation. arXiv preprint arXiv:2302.00111, 2023.
Fedus et al. (2022) William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022.
Fu et al. (2024) Zipeng Fu, Tony Z Zhao, and Chelsea Finn. Mobile aloha: Learning bimanual mobile manipulation with low-cost whole-body teleoperation. arXiv preprint arXiv:2401.02117, 2024.
Go et al. (2023) Hyojun Go, Yunsung Lee, Jin-Young Kim, Seunghyun Lee, Myeongho Jeong, Hyun Seung Lee, and Seungtaek Choi. Towards practical plug-and-play diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 1962–1971, 2023.
Gu et al. (2024) Jiayuan Gu, Sean Kirmani, Paul Wohlhart, Yao Lu, Montserrat Gonzalez Arenas, Kanishka Rao, Wenhao Yu, Chuyuan Fu, Keerthana Gopalakrishnan, Zhuo Xu, Priya Sundaresan, Peng Xu, Hao Su, Karol Hausman, Chelsea Finn, Quan Vuong, and Ted Xiao. Rt-trajectory: Robotic task generalization via hindsight trajectory sketches. In International Conference on Learning Representations, 2024.
Hang et al. (2024) Tiankai Hang, Shuyang Gu, Chen Li, Jianmin Bao, Dong Chen, Han Hu, Xin Geng, and Baining Guo. Efficient diffusion training via min-snr weighting strategy, 2024.
Hansel et al. (2023) Kay Hansel, Julen Urain, Jan Peters, and Georgia Chalvatzaki. Hierarchical policy blending as inference for reactive robot control. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 10181–10188. IEEE, 2023.
Hazimeh et al. (2021) Hussein Hazimeh, Zhe Zhao, Aakanksha Chowdhery, Maheswaran Sathiamoorthy, Yihua Chen, Rahul Mazumder, Lichan Hong, and Ed H. Chi. Dselect-k: Differentiable selection in the mixture of experts with applications to multi-task learning. CoRR, abs/2106.03760, 2021. URL https://arxiv.org/abs/2106.03760.
Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840–6851, 2020.
Janner et al. (2022) Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, pp. 9902–9915. PMLR, 2022.
Jia et al. (2024) Xiaogang Jia, Denis Blessing, Xinkai Jiang, Moritz Reuss, Atalay Donat, Rudolf Lioutikov, and Gerhard Neumann. Towards diverse behaviors: A benchmark for imitation learning with human demonstrations. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=6pPYRXKPpw.
Jiang et al. (2024) Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint arXiv:2401.04088, 2024.
Karras et al. (2022) Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Ke et al. (2024) Tsung-Wei Ke, Nikolaos Gkanatsios, and Katerina Fragkiadaki. 3d diffuser actor: Policy diffusion with 3d scene representations. arXiv preprint arXiv:2402.10885, 2024.
Kim et al. (2024) Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti, Ted Xiao, Ashwin Balakrishna, Suraj Nair, Rafael Rafailov, Ethan Foster, Grace Lam, Pannag Sanketi, et al. Openvla: An open-source vision-language-action model. arXiv preprint arXiv:2406.09246, 2024.
Ko et al. (2023) Po-Chen Ko, Jiayuan Mao, Yilun Du, Shao-Hua Sun, and Joshua B Tenenbaum. Learning to Act from Actionless Video through Dense Correspondences. arXiv:2310.08576, 2023.
Le et al. (2023) An Thai Le, Kay Hansel, Jan Peters, and Georgia Chalvatzaki. Hierarchical policy blending as optimal transport. In Learning for Dynamics and Control Conference, pp. 797–812. PMLR, 2023.
Lee et al. (2024a) Seungjae Lee, Yibin Wang, Haritheja Etukuru, H. Jin Kim, Nur Muhammad Mahi Shafiullah, and Lerrel Pinto. Behavior generation with latent actions. arXiv preprint arXiv:2403.03181, 2024a.
Lee et al. (2024b) Yunsung Lee, JinYoung Kim, Hyojun Go, Myeongho Jeong, Shinhyeok Oh, and Seungtaek Choi. Multi-architecture multi-expert diffusion models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 13427–13436, 2024b.
Lewis et al. (2021) Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models, 2021.
Li et al. (2023a) Maximilian Xiling Li, Onur Celik, Philipp Becker, Denis Blessing, Rudolf Lioutikov, and Gerhard Neumann. Curriculum-based imitation of versatile skills. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 2951–2957, 2023a. doi: 10.1109/ICRA48891.2023.10160543.
Li et al. (2023b) Xiang Li, Varun Belagali, Jinghuan Shang, and Michael S Ryoo. Crossway diffusion: Improving diffusion-based visuomotor policy via self-supervised learning. arXiv preprint arXiv:2307.01849, 2023b.
Li et al. (2024a) Xinghang Li, Minghuan Liu, Hanbo Zhang, Cunjun Yu, Jie Xu, Hongtao Wu, Chilam Cheang, Ya Jing, Weinan Zhang, Huaping Liu, et al. Vision-language foundation models as effective robot imitators. In International Conference on Learning Representations, 2024a.
Li et al. (2024b) Xuanlin Li, Kyle Hsu, Jiayuan Gu, Karl Pertsch, Oier Mees, Homer Rich Walke, Chuyuan Fu, Ishikaa Lunawat, Isabel Sieh, Sean Kirmani, Sergey Levine, Jiajun Wu, Chelsea Finn, Hao Su, Quan Vuong, and Ted Xiao. Evaluating real-world robot manipulation policies in simulation. arXiv preprint arXiv:2405.05941, 2024b.
Liu et al. (2023) Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, Qiang Liu, Yuke Zhu, and Peter Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. arXiv preprint arXiv:2306.03310, 2023.
Lynch et al. (2019) Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play, 2019.
Mees et al. (2016) Oier Mees, Andreas Eitel, and Wolfram Burgard. Choosing smartly: Adaptive multimodal fusion for object detection in changing environments. In 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 151–156. IEEE, 2016.
Mees et al. (2022a) Oier Mees, Lukas Hermann, and Wolfram Burgard. What matters in language conditioned robotic imitation learning over unstructured data. IEEE Robotics and Automation Letters (RA-L), 7(4):11205–11212, 2022a.
Mees et al. (2022b) Oier Mees, Lukas Hermann, Erick Rosete-Beas, and Wolfram Burgard. Calvin: A benchmark for language-conditioned policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters, 2022b.
Mete et al. (2024) Atharva Mete, Haotian Xue, Albert Wilcox, Yongxin Chen, and Animesh Garg. Quest: Self-supervised skill abstractions for learning continuous control. arXiv preprint arXiv:2407.15840, 2024.
Muennighoff et al. (2024) Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02060, 2024.
Nair et al. (2017) Ashvin Nair, Dian Chen, Pulkit Agrawal, Phillip Isola, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Combining self-supervised learning and imitation for vision-based rope manipulation. In 2017 IEEE international conference on robotics and automation (ICRA), pp. 2146–2153. IEEE, 2017.
Obando-Ceron et al. (2024) Johan Obando-Ceron, Ghada Sokar, Timon Willi, Clare Lyle, Jesse Farebrother, Jakob Foerster, Gintare Karolina Dziugaite, Doina Precup, and Pablo Samuel Castro. Mixtures of experts unlock parameter scaling for deep rl. arXiv preprint arXiv:2402.08609, 2024.
Octo Model Team et al. (2023) Octo Model Team, Dibya Ghosh, Homer Walke, Karl Pertsch, Kevin Black, Oier Mees, Sudeep Dasari, Joey Hejna, Charles Xu, Jianlan Luo, Tobias Kreiman, You Liang Tan, Dorsa Sadigh, Chelsea Finn, and Sergey Levine. Octo: An open-source generalist robot policy. https://octo-models.github.io, 2023.
Pari et al. (2021) Jyothish Pari, Nur Muhammad Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The surprising effectiveness of representation learning for visual imitation, 2021.
Pari et al. (2022) Jyothish Pari, Nur Muhammad (Mahi) Shafiullah, Sridhar Pandian Arunachalam, and Lerrel Pinto. The Surprising Effectiveness of Representation Learning for Visual Imitation. In Proceedings of Robotics: Science and Systems, New York City, NY, USA, June 2022. doi: 10.15607/RSS.2022.XVIII.010.
Park et al. (2023) Byeongjun Park, Sangmin Woo, Hyojun Go, Jin-Young Kim, and Changick Kim. Denoising task routing for diffusion models. arXiv preprint arXiv:2310.07138, 2023.
Park et al. (2024) Byeongjun Park, Hyojun Go, Jin-Young Kim, Sangmin Woo, Seokil Ham, and Changick Kim. Switch diffusion transformer: Synergizing denoising tasks with sparse mixture-of-experts. arXiv preprint arXiv:2403.09176, 2024.
Peebles & Xie (2023) William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195–4205, 2023.
Perez et al. (2018) Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Pomerleau (1989) Dean Pomerleau. Alvinn: An autonomous land vehicle in a neural network. In D.S. Touretzky (ed.), Proceedings of (NeurIPS) Neural Information Processing Systems, pp. 305 – 313. Morgan Kaufmann, December 1989.
Reuss et al. (2023) Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal conditioned imitation learning using score-based diffusion policies. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Reuss et al. (2024) Moritz Reuss, Ömer Erdinç Yağmurlu, Fabian Wenzel, and Rudolf Lioutikov. Multimodal diffusion transformer: Learning versatile behavior from multimodal goals. In Robotics: Science and Systems, 2024.
Riquelme et al. (2021) Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts, 2021.
Roller et al. (2021) Stephen Roller, Sainbayar Sukhbaatar, Arthur Szlam, and Jason Weston. Hash layers for large sparse models. CoRR, abs/2106.04426, 2021. URL https://arxiv.org/abs/2106.04426.
Scheikl et al. (2023) Paul Maria Scheikl, Nicolas Schreiber, Christoph Haas, Niklas Freymuth, Gerhard Neumann, Rudolf Lioutikov, and Franziska Mathis-Ullrich. Movement primitive diffusion: Learning gentle robotic manipulation of deformable objects. arXiv preprint arXiv:2312.10008, 2023.
Shafiullah et al. (2022) Nur Muhammad Mahi Shafiullah, Zichen Jeff Cui, Ariuntuya Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning $k$ modes with one stone. In Thirty-Sixth Conference on Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id=agTr-vRQsa.
Shazeer et al. (2017) Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017.
Song et al. (2021) Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021.
Song & Ermon (2019) Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
Song et al. (2020) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
Urain et al. (2023) Julen Urain, Niklas Funk, Jan Peters, and Georgia Chalvatzaki. Se (3)-diffusionfields: Learning smooth cost functions for joint grasp and motion optimization through diffusion. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 5923–5930. IEEE, 2023.
Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Vincent (2011) Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23(7):1661–1674, 2011. doi: 10.1162/NECO_a_00142.
Walke et al. (2023) Homer Rich Walke, Kevin Black, Tony Z Zhao, Quan Vuong, Chongyi Zheng, Philippe Hansen-Estruch, Andre Wang He, Vivek Myers, Moo Jin Kim, Max Du, et al. Bridgedata v2: A dataset for robot learning at scale. In Conference on Robot Learning, pp. 1723–1736. PMLR, 2023.
Wang et al. (2023) Tsun-Hsuan Wang, Juntian Zheng, Pingchuan Ma, Yilun Du, Byungchul Kim, Andrew Everett Spielberg, Joshua B. Tenenbaum, Chuang Gan, and Daniela Rus. Diffusebot: Breeding soft robots with physics-augmented generative diffusion models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview.net/forum?id=1zo4iioUEs.
Wang et al. (2024) Yixiao Wang, Yifei Zhang, Mingxiao Huo, Ran Tian, Xiang Zhang, Yichen Xie, Chenfeng Xu, Pengliang Ji, Wei Zhan, Mingyu Ding, et al. Sparse diffusion policy: A sparse, reusable, and flexible policy for robot learning. arXiv preprint arXiv:2407.01531, 2024.
Wu et al. (2024) Hongtao Wu, Ya Jing, Chilam Cheang, Guangzeng Chen, Jiafeng Xu, Xinghang Li, Minghuan Liu, Hang Li, and Tao Kong. Unleashing large-scale video generative pre-training for visual robot manipulation. In International Conference on Learning Representations, 2024.
Xian et al. (2023) Zhou Xian, Nikolaos Gkanatsios, Theophile Gervet, Tsung-Wei Ke, and Katerina Fragkiadaki. Chaineddiffuser: Unifying trajectory diffusion and keypose prediction for robotic manipulation. In 7th Annual Conference on Robot Learning, 2023. URL https://openreview.net/forum?id=W0zgY2mBTA8.
Zhao et al. (2023) Tony Z Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304.13705, 2023.
Zhou et al. (2022) Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew Dai, Zhifeng Chen, Quoc Le, and James Laudon. Mixture-of-experts with expert choice routing, 2022.
Zitkovich et al. (2023) Brianna Zitkovich, Tianhe Yu, Sichun Xu, Peng Xu, Ted Xiao, Fei Xia, Jialin Wu, Paul Wohlhart, Stefan Welker, Ayzaan Wahid, Quan Vuong, Vincent Vanhoucke, Huong Tran, Radu Soricut, Anikait Singh, Jaspiar Singh, Pierre Sermanet, Pannag R Sanketi, Grecia Salazar, Michael S Ryoo, Krista Reymann, Kanishka Rao, Karl Pertsch, Igor Mordatch, Henryk Michalewski, Yao Lu, Sergey Levine, Lisa Lee, Tsang-Wei Edward Lee, Isabel Leal, Yuheng Kuang, Dmitry Kalashnikov, Ryan Julian, Nikhil J Joshi, Alex Irpan, brian ichter, Jasmine Hsu, Alexander Herzog, Karol Hausman, Keerthana Gopalakrishnan, Chuyuan Fu, Pete Florence, Chelsea Finn, Kumar Avinava Dubey, Danny Driess, Tianli Ding, Krzysztof Marcin Choromanski, Xi Chen, Yevgen Chebotar, Justice Carbajal, Noah Brown, Anthony Brohan, Montserrat Gonzalez Arenas, and Kehang Han. RT-2: Vision-language-action models transfer web knowledge to robotic control. In 7th Annual Conference on Robot Learning, 2023. URL https://openreview.net/forum?id=XMQgwiJ7KSX.

Appendix A Appendix / supplemental material

Hyperparameter	CALVIN ABCD	CALVIN ABC	LIBERO-10	LIBERO-90	Pret-MoDE
Number of Transformer Layers	8	8	8	8	12
Number Experts	4	4	4	4	4
Attention Heads	8	8	8	8	8
Action Chunk Size	10	10	10	10	10
History Length	1	1	1	1	1
Embedding Dimension	1024	1024	1024	1024	1024
Image Encoder	FiLM-ResNet18	FiLM-ResNet50	FiLM-ResNet18	FiLM-ResNet18	FiLM-ResNet50
Goal Lang Encoder	CLIP ViT-B/32	CLIP ViT-B/32	CLIP ViT-B/32	CLIP ViT-B/32	CLIP ViT-B/32
Attention Dropout	0.3	0.3	0.3	0.3	0.3
Residual Dropout	0.1	0.1	0.1	0.1	0.1
MLP Dropout	0.1	0.1	0.1	0.1	0.1
Optimizer	AdamW	AdamW	AdamW	AdamW	AdamW
Betas	[0.9, 0.95]	[0.9, 0.95]	[0.9, 0.95]	[0.9, 0.95]	[0.9, 0.95]
Learning Rate	1e-4	1e-4	1e-4	1e-4	1e-4
Transformer Weight Decay	0.05	0.05	0.05	0.05	0.1
Other weight decay	0.05	0.05	0.05	0.05	0.1
Batch Size	512	512	512	512	512
Train Steps in Thousands	30	25	15	30	300
$\sigma_{\text{max}}$	80	80	80	80	80
$\sigma_{\text{min}}$	0.001	0.001	0.001	0.001	0.001
$\sigma_{t}$	0.5	0.5	0.5	0.5	0.5
EMA	True	True	True	True	True
Time steps	Exponential	Exponential	Exponential	Exponential	Exponential
Sampler	DDIM	DDIM	DDIM	DDIM	DDIM
Parameter Count (Millions)	460	460	460	460	685

表3: 我々の実験で使用したMoDEポリシーのすべてのハイパーパラメータの概要。

A.1 Pretraining Details

Dataset	Weight
BC-Z	0.258768
LIBERO-10	0.043649
BRIDGE	0.188043
CMU Play-Fusion	0.101486
Google Fractal	0.162878
DOBB-E	0.245176
Total	1.000000

表4: 軌道の小さなサブセットでMoDEを訓練するために使用されたデータセットのサンプリング重み。全データセットは196,000の軌道で構成されている。

我々は、MoDEの大規模バリアントをOpen-X-Embodiment Collaboration et al. (2023)から入手可能なデータセットのサブセットで事前学習し、MoDEの汎化能力を研究する。使用したデータセットの概要は表4にまとめられている。我々の事前学習データセットは、6つの異なるソースからの196,000の軌道で構成されており、BC-Z（0.259）、LIBERO-10（0.044）、BRIDGE（0.188）、CMU Play-Fusion（0.101）、Google Fractal（0.163）、DOBB-E（0.245）にわたって重み付けサンプリングを行っている。このデータセットには、Googleロボット、Franka Panda、Hello-Stretchロボットを含む多様なロボットプラットフォームからのデモンストレーションが含まれており、幅広い操作タスクをカバーしている。事前学習は、40GB VRAMを搭載した6台のNVIDIA A6000 GPUで3日間にわたって実施され、300,000の訓練ステップを完了した。我々は、バッチサイズ1024、学習率1e-4、重み減衰0.1を使用した。訓練中のバランスの取れたデータセット混合を確保するために、400,000サンプルの大規模なシャッフルバッファを実装した。各データセットは、様々なロボットプラットフォーム間の異なるスケールと範囲を考慮して個別に正規化された。この多様な事前学習により、MoDEのゼロショット汎化が大幅に改善され、特にCALVIN ABC→Dのような困難なベンチマークでは、平均ロールアウト長4.01という新しい最先端の性能を達成した。再現性を確保するために、我々は事前学習済みモデルの重みと前処理コードをコミュニティにリリースする予定である。

ファインチューニングでは、モデルのルーターを凍結し、負荷分散損失を除去して、LIBEROでは10,000ステップ、CALVINでは15,000ステップ、GPU当たり64のバッチサイズで4台のGPUを使用してローカルドメインで訓練を行う。

A.2 Experiments Details

Benchmark	MoDE	DP-T	DP-CNN	Avg. Baseline	Improvement
CALVIN ABC→D (norm.)	0.678	0.226	0.270	0.248	+151.1%
CALVIN ABCD→D (norm.)	0.860	0.396	0.632	0.514	+36.1%
LIBERO-90	0.910	0.690	0.780	0.735	+16.7%
LIBERO-10	0.920	0.510	0.730	0.620	+26.0%
Average Improvement Over Second-Best:					57.5%

表5: 詳細な性能改善分析。CALVINのスコアはLIBEROのスケールに合わせるために5で割って正規化されている。改善率は以下のように計算: (MoDE - 平均ベースライン) / 平均ベースライン × 100。最終平均は、4つのベンチマークすべてにおいて、各ベンチマークで2番目に優れたDiffusion Policy変種と比較した改善率の平均である。

平均性能向上。 MoDEの既存のDiffusion Policiesに対する優位性を定量化するために、我々はすべてのベンチマークにおいて2番目に優れた手法と比較した。 MoDEは大幅な改善を示し、特にCALVIN ABC→Dのような困難な転移シナリオでは、次に優れた手法を151.1%上回った。より標準化されたLIBEROベンチマークにおいても、MoDEは16.7%から26.0%の大きな優位性を維持した。すべてのタスクの平均では、MoDEは2番目に優れた手法と比較して57.5%の改善を達成し、同時に同程度のパラメータ数の密なトランスフォーマーと比較して90%少ないFLOPsで計算効率を維持した。

A.2.1 MoDE Evaluation on SIMPLER

我々は、MoDEの汎用的なポリシーとしての能力を、Open-X-Embodimentからの大幅に大きなデータセットで訓練された2つの最先端モデルと比較して評価する：Octo（80万のトラジェクトリ）(Octo Model Team et al., 2023)とOpenVLA (Kim et al., 2024)（100万のトラジェクトリ）。この比較は、RT-1 (Brohan et al., 2023)の訓練に使用されたBridgeV2 (Walke et al., 2023)およびGoogle Fractalデータセットのreal2sim変種を提供するSIMPLERベンチマークを使用して行う。このベンチマークは、図7に示されているように、複数の環境にわたる多様な操作タスクを包含している。

	OpenVLA		Octo Base		MoDe (ours)
Metric	Score	Rank	Score	Rank	Score	Rank
Drawer Open	16%	1	0%	3	4.23%	2
Drawer Close	20%	2	2%	3	34.92%	1
Pick Can Horizontal	71%	1	0%	3	33.78%	2
Pick Can Vertical	27%	2	0%	3	29.78%	1
Pick Can Standing	65%	1	0%	3	36.44%	2
Move Near	48%	1	3%	3	30%	2
Drawer Open	19%	2	1%	3	21.30%	1
Drawer Close	52%	2	44%	3	76.85%	1
Pick Can Horizontal	27%	1	21%	3	22%	2
Pick Can Vertical	3%	3	21%	2	40%	1
Pick Can Standing	19%	2	9%	3	35%	1
Move Near	46%	1	4%	3	45.42%	2
Partial Put Spoon on Tablecloth	4%	3	35%	1	29.17%	2
Put Spoon on Tablecloth	0%	3	12%	1	12.5%	1
Partial Put Carrot on Plate	33%	2	53%	1	29.17%	3
Put Carrot on Plate	0%	3	8%	1	8.33%	1
Partial Stack Green Block on Yellow Block	12%	2	32%	1	8.33%	3
Stack Green Block on Yellow Block	0%	2	0%	2	0%	2
Partial Put Eggplant in Basket	8%	3	67%	1	37.5%	2
Put Eggplant in Basket	4%	3	43%	1	8.33%	2
Average	23.70%	1.95	17.75%	2.1	26.30%	1.65

表6: 2952回の評価で全SIMPLERタスクでテストされた2つの最先端汎用ポリシーOpenVLA Kim et al. (2024)とOcto Octo Model Team et al. (2023)に対するMoDEの詳細な比較。

評価の結果は表6にまとめられている。平均してMoDEは26.30%の最高平均成功率と1.65の最良平均ランキングをすべてのタスクにわたって達成し、Octo（17.75%の成功率、2.1のランク）と7.7Bパラメータを持つOpenVLAモデル（23.70%の成功率、1.95のランク）の両方を上回っている。 MoDEは、引き出し操作（引き出しを閉じるタスクで34.92%）や精密な物体操作（垂直な缶のピッキングで40%）などの困難な操作タスクで特に強力な性能を示している。ブロックの積み重ねのような特殊なタスクはすべてのモデルにとって依然として困難であるが、MoDEの多様なタスクにわたる一貫した性能は、汎用ポリシーのためのスケーラブルなアーキテクチャとしての有効性を示している。

A.2.2 CALVIN Benchmark

CALVINベンチマーク (Mees et al., 2022b) は、人間のプレイデータから言語条件付き行動を学習するための確立されたILベンチマークである。他のベンチマークとは対照的に、このデータにはロボットが1つのタスクを完了する構造化されたデモンストレーションは含まれておらず、代わりに、ランダムに環境と相互作用する人間によってデータセットが収集された。 $4$ の異なる設定にわたるこれらの長期的なトラジェクトリから、著者らはランダムに $64$ フレームの短いシーケンスを切り出し、タスクラベルを付けた。データセットはラベルのない部分でもモデルを訓練するオプションを提供しているが、我々はMoDEをラベル付けされた部分のみで訓練することに制限した。 Franke Emika Pandaロボットは、離散的なグリッパーを持つデルタエンドエフェクタ空間を使用して制御される。我々は現在のシーンをエンコードするために2つのカメラを使用する：静的カメラと手首カメラである。そして次の $10$ アクションを予測し、次の観察を受け取った後、別の $10$ アクションのセットを生成する。

CALVIN ABC。我々はMoDEと密なトランスフォーマーのベースラインを、 $25$ kの訓練ステップ、バッチサイズ $512$ で、 $4$ GPUクラスタノード上の $4$ 台のA $6000$ NVIDIA GPUを使用して $6.5$ 時間訓練し、訓練終了時に $1000$ 回のロールアウトを行った。我々は、関連する先行研究すべてで行われているように、 $3$ シードにわたる平均結果を報告する。すべてのベースラインは、CALVINの標準化された評価プロトコル (Mees et al., 2022b) に基づいて元の論文から報告されている。

CALVIN ABCD。我々はMoDEと密なトランスフォーマーのベースラインを、 $30$ kの訓練ステップ、バッチサイズ $512$ で、 $4$ GPUクラスタノード上の $4$ 台のA $6000$ NVIDIA GPUを使用して $7.5$ 時間訓練し、訓練終了時に $1000$ 回のロールアウトを行った。我々は、関連する先行研究すべてで行われているように、 $3$ シードにわたる平均結果を報告する。

A.2.3 LIBERO Benchmark

LIBERO-10。 LIBERO-10ベンチマークは、 $10$ の異なるタスクに対して $50$ のデモンストレーションで構成され、すべてテキスト指示でラベル付けされている。 Franka Emika Pandaロボットはエンドエフェクタコントローラーを使用して制御される。 CALVINと同様に、すべてのモデルは2つのカメラ入力にアクセスできる：静的なものと手首カメラである。我々はMoDEと密なトランスフォーマーのベースラインを $50$ エポック、バッチサイズ $512$ で、 $4$ GPUクラスタノード上の $4$ 台のA $6000$ NVIDIA GPUを使用して $2$ 時間訓練し、訓練終了時に $200$ 回のロールアウトを行った。このベンチマークでは、 $10$ の異なる長期的タスクでモデルをテストする必要がある。我々は各タスクを各モデルに対して $20$ 回テストし、 $10$ タスク全体の最終的な平均性能を報告する。

LIBERO-90。 LIBERO-10ベンチマークは、 $90$ の異なるタスクに対して $50$ のデモンストレーションで構成され、すべてテキスト指示でラベル付けされている。 Franka Emika Pandaロボットはエンドエフェクタコントローラーを使用して制御される。我々はMoDEと密なトランスフォーマーのベースラインを $50k$ ステップ、バッチサイズ $512$ で、 $4$ GPUクラスタノード上の $4$ 台のA $6000$ NVIDIA GPUを使用して $12$ 時間訓練し、訓練終了時に $1800$ 回のロールアウトを行った。このベンチマークでは、多くの異なる環境で $90$ の異なるタスクでモデルをテストする必要がある。我々は各タスクを各モデルに対して $20$ 回テストし、 $90$ タスク全体の最終的な平均性能を報告する。

Model	Block Push	Relay Kitchen	CAL ABC	CAL ABCD	L-10	Average
Dense T	0.96 $\pm$ 0.02	3.73 $\pm$ 0.12	2.83 $\pm$ 0.19	4.13 $\pm$ 0.11	0.91 $\pm$ 0.02	0.839 $\pm$ 0.144
Token-Router	0.97 $\pm$ 0.01	3.85 $\pm$ 0.03	2.67 $\pm$ 0.04	4.29 $\pm$ 0.08	0.90 $\pm$ 0.01	0.845 $\pm$ 0.161
$\sigma_{t}$ -Router	0.97 $\pm$ 0.01	3.79 $\pm$ 0.04	2.79 $\pm$ 0.16	4.30 $\pm$ 0.02	0.92 $\pm$ 0.02	0.851 $\pm$ 0.151

表7: MoDEに使用されたすべての異なるトークンルーティング戦略の

5

ベンチマークにわたる性能の概要。各環境で最良の結果を太字で、2番目に良い結果を斜体で示す。CALはCALVINを表す。結果を平均化するために、我々はすべてのスコアを正規化し、すべての環境にわたる平均を計算する。

A.3 Baselines

以下に、実験で使用した複数のベースラインについて詳細に説明する：

Diffusion Policy-CNN/T (Chi et al., 2023)に触発され、我々はDDPMベースのDiffusion Policyフレームワークを目標条件付きマルチタスク学習に拡張したものを評価する。 CNNベースの変種とDiffusion-Transformer変種の2つのバージョンを評価する。後者はコンテキストとノイズをクロスアテンションを用いて条件付けている。我々の実験では、MoDEとの公平な比較のために、EDMベースのDiffusionフレームワークも使用する。 Transformerベースラインの理想的な層数と潜在次元を最適化し、最終バージョンでは $8$ 層、潜在次元 $1024$ を使用する。より大きいまたは小さい変種では、平均性能が低下した。

RoboFlamingo RoboFlamingo (Li et al., 2024a)は、行動生成のためにファインチューニングされたVision-Language-Models (VLM)である。著者らは $3$ 億パラメータのFlamingoモデル(Alayrac et al., 2022)を使用し、CALVINでファインチューニングを行っている。その際、フォワードブロックを凍結し、新しいPerceiver Resamplerモジュールのみをファインチューニングして、凍結されたvision-transformer画像エンコーダーから特徴を抽出し、クロスアテンション層で画像特徴を処理している。最後に、行動を生成するための新しい行動ヘッドが学習される。全体として、ファインチューニングには約 $1$ 億のパラメータの訓練が必要である。我々は、標準のCALVIN評価スイートを使用しているため、論文から報告された結果を記載する。

SuSIE このモデルは、まず別の画像とテキスト記述に条件付けられた画像を生成する画像生成拡散モデルであるInstruct2Pix (Brooks et al., 2023)をローカルのCALVINロボティクスドメインでファインチューニングし、高レベルの目標生成器として使用する。低レベルのポリシーはCNNベースのDiffusion Policyであり、現在の状態埋め込みと高レベルポリシーからの望ましいサブゴールが与えられた次の $4$ 個の行動を予測する(Black et al., 2023)。

GR-1 因果的なGPT-Transformerモデル(Wu et al., 2024)であり、人間の動画の大規模な生成的ビデオ予測で事前訓練されている。その後、CALVINで行動予測とビデオ予測の共同訓練を用いてモデルがファインチューニングされる。我々は、CALVINベンチマークについて、彼らの論文から直接結果を報告する。

A.4 Average FLOPs computation and Inference Speed

Method	Active Params (M)	Total Params (M)	GFLOPS	PrT	Avg. Length	SF-Ratio	Inf. Time [ms]
Diff-P-CNN	321	321	1.28	$\times$	1.35	1.05	11.7
Diff-P-T	194	194	2.16	$\times$	1.13	0.53	16.2
RoboFlamingo	1000	1000	690	$\checkmark$	2.47	0.004	65
SuSIE	860+	860+	60	$\checkmark$	2.69	0.045	199
GR-1	130	130	27.5	$\checkmark$	3.06	0.11	12.6
MoDE (ours)	436	740	1.53	$\checkmark$	4.01	2.6	12.2

表8: CALVINベンチマークで使用された手法のパラメータ総数とアクティブパラメータ数の比較。異なる手法に必要な平均FLOPSの概要と、ABCベンチマークにおける平均性能も併せて示す。SF比は平均ロールアウト長とGFLOPSを比較している。

我々は表8において、各手法で使用される総パラメータ数とFLOPsの詳細な比較を提供する。さらに、様々な手法を比較し、単一のアクションの平均予測時間を測定するために、GFLOPS（ $10^{9}$ FLOPS）あたりの計算効率指標を提供する。以下では、CALVIN ABCベンチマークにおける全ての関連するベースラインの平均GFLOPS計算について詳述する。具体的には、単一のアクションを予測するために必要な平均GFLOPSを比較する。

公平な比較を保証するために、我々は全ての手法を同じNVIDIA A6000 GPU（40 GB VRAM）で評価した。平均推論速度を計算するために、各手法を100回テストし、大きな外れ値を除去して平均時間を計算した。

MoDE。我々は12層、4エキスパート、隠れ次元1024の大規模な事前学習済みバリアントをベンチマークした。順伝播の平均GFLOPSは0.7 GFLOPSである。ルーターキャッシングがなければ、モデルは5 GFLOPSを必要とするが、ルーターキャッシングにより全体の計算コストが90 $\%$ 以上削減されることを示している。このアーキテクチャは合計14トークン（1ノイズ + 1ゴール + 2画像 + 10ノイズ付きアクション）を処理する。 MoDEは10回のデノイジングパスで10アクションのシーケンスを予測する。 ResNet-50を使用するバリアントでは、画像エンコーダーに8.27 GFLOPSが必要である。単一のアクションに対して平均的に、MoDEはトランスフォーマーで10回の順伝播とResNet-50で1回のパスを必要とする。結果として、MoDEの事前学習済みバリアントは、単一のアクションを予測するために平均1.53 GFLOPSを必要とする。このモデルの推論時間はハードウェアに依存する。我々はアクションあたりの平均推論時間を12.2ミリ秒と測定した。

DP-CNN/T。 DP-CNNは平均順伝播に0.8 GFLOPSを使用する。 ResNet-18は3.62 GFLOPSを必要とする。このモデルは10回のデノイジングステップで10アクションを予測し、再計画なしで10アクションを実行する。これにより、CNN版は単一のアクションを予測するために1.28 GFLOPSを必要とする。トランスフォーマー版では、アーキテクチャは10回のデノイジングステップを使用して10アクションを予測し、MoDEと同様に合計14トークン（1ノイズ + 1ゴール + 2画像 + 10ノイズ付きアクション）を処理する。トランスフォーマーを通る順伝播で平均1.8 GFLOPSの使用を達成する。 DP-Tベースラインは単一のアクションを予測するために平均2.16 GFLOPSを必要とする。 CNN版は単一のアクションを予測するのに平均11.7ミリ秒を要し、クロスアテンション条件付けを持つトランスフォーマー版は16.2ミリ秒を要する。

RoboFlamingo。計算分析において、このモデルはViTで単一の画像をエンコードするために34 GFLOPSを必要とする。ポリシーバックボーンについては、論文で使用されている"mpt-1b-redpajama-200b-dolly"バリアントを評価した。このアーキテクチャは、32トークンの平均シーケンスに対して順伝播あたり656 GFLOPSを必要とする。 RoboFlamingoには複数のバリアントが存在するが、これは必要なGFLOPSの大まかな見積もりを提供する。合計で、CALVINでアクションを予測するために平均690 GFLOPSが必要と推定される。単一のアクションを予測するために、このモデルは平均65ミリ秒を要する。

SuSIE。我々の計算分析では、SuSIEによって実装された50回のデノイジングステップを持つInstruct2Pixをテストした。結果として得られた1026 GFLOPSは20で割られる。これは、モデルが20タイムステップごとに新しいサブゴールを生成するためである。低レベルポリシーは8.27 GFLOPSのResNet-50画像エンコーダーを使用する。他のポリシーとは対照的に、SuSIEは状態ごとに1つの画像のみを計算し、毎タイムステップでアクションを予測する。これらは指数平均を用いて平均化される。したがって、我々は小さな拡散ヘッドを省略し、アクションあたり60 GFLOPSの見積もりを得る。平均推論速度については、単一のゴール画像を生成する時間を測定し、それを20で割り、次にResNet-50で2つの画像をエンコードする平均時間と小さなMLPを通る10回の順伝播を加える。 20タイムステップごとに、SuSIEが新しい画像を生成する際、単一のアクション生成に3777.62ミリ秒を要する。それ以外の場合は10.7ミリ秒とはるかに高速である。平均して、SuSIEは単一のアクションを生成するのに199ミリ秒を要し、これは全体で最も遅いポリシーとなる。

GR-1。事前学習済みMAE Vision Transformerは、単一の画像をエンコードするために約17.5 GFLOPSを必要とする。トランスフォーマーバックボーンは、履歴長10と1タイムステップあたり15トークン（10画像トークン + 1ゴールトークン + 1自己受容性トークン + 2ビデオ読み出しトークン + 1アクショントークン）で150トークンを処理する。結果として、潜在次元384と12層を持つデコーダーを使用した単一のアクション予測の平均GFLOPSは10 GFLOPSである。合計すると、これはCALVINで単一のアクションを予測するための平均計算コストが27.5 GFLOPSになる。平均的な単一アクション予測に、このモデルは12.6ミリ秒を要する。

分析。全体として、MoDEはテストされた全てのベースラインの中で最高のパフォーマンス対GFLOPS比を提供する。 MoDEは他の拡散ポリシーアーキテクチャと比較して総サイズが大幅に大きいにもかかわらず、同様の推論速度と低いFLOPカウントを必要とする。さらに、計算リソースの観点から優れた効率性を示しつつ、CALVINベンチマークタスクで高いパフォーマンスを維持している。推論速度に関しては、MoDEは総パラメータ数が多いにもかかわらず2番目に高速である。

A.5 Detailed Experimental Results

我々はルーティングの選択に関する比較実験の結果を表7にまとめている。したがって、 $5$ のベンチマークにわたって $2$ の異なるルーティング戦略を検証した。

A.6 State-based Experiments

我々は、2つの確立されたマルチタスク状態ベース環境でMoDEを用いて追加実験を行った：

リレーキッチン。我々は(Lynch et al., 2019)のFrankaキッチン環境をモデル評価に使用した。この仮想キッチン環境では、人間の参加者がVRインターフェースを使用して7つのオブジェクト（ケトル、電子レンジ、スライドドア、蝶番ドア、ライトスイッチ、2つのバーナー）を操作できる。結果として得られたデータセットは、元の研究者によって収集された $566$ のデモンストレーションで構成されており、各参加者はエピソードごとに4つの所定の操作タスクを実行した。 Franka Emika Pandaロボットは、ロボットの関節とエンドエフェクタの位置を表す $9$ 次元のアクション空間を介して制御される。 $30$ 次元の観測空間には、環境内の関連オブジェクトの現在の状態に関する情報が含まれる。目標状態として、我々は到達すべき目標として将来の状態をランダムにサンプリングする。

この実験では、我々はすべてのモデルを $40$ kのトレーニングステップで訓練し、バッチサイズは $1024$ とした。そして、公平な評価を保証するために、先行研究(Shafiullah et al., 2022; Cui et al., 2023; Reuss et al., 2023)と同様に $100$ 回評価を行った。報告されたすべての結果は $4$ シードの平均である。我々のモデルは、ローカルPCのRTX $3070$ GPUで、追加の実験ロールアウトを含めて各実行につき約 $2$ 時間訓練した。

ブロックプッシュ。 PyBullet環境では、XArmロボットが平面内の2つの正方形のターゲットに2つのブロックを押し込むタスクを行う。ブロックを押す順序と特定のブロック-ターゲットの組み合わせは、 $1000$ のデモンストレーションセットから目標状態としてサンプリングされる。我々のモデルのトレーニングに使用されたデモンストレーションは、最初に押すブロックを選択し、そのブロックのターゲットを独立して選択するハードコードされたコントローラーを使用して収集された。最初のブロックをターゲットに押した後、コントローラーは2番目のブロックを残りのターゲットに押す。このアプローチにより、4つの可能な行動モードが生まれ、ブロックをターゲットに押し込む様々な方法から追加の確率性が生じる。モデルは、ブロックが正しいターゲット位置と順序で押された場合にのみクレジットを得る。我々は、ブロックの中心がターゲット正方形の0.05単位以内にある場合、そのブロックが正常に押されたと見なす。

すべてのモデルは、これらのランダム化された条件下で $1000$ のコントローラー生成デモンストレーションのデータセットで訓練された。すべてのモデルは $60$ kステップ、バッチサイズ $1024$ で訓練された。評価のために、我々は先行研究(Shafiullah et al., 2022; Cui et al., 2023; Reuss et al., 2023)に従い、 $100$ の異なる指示でテストし、 $4$ シードの平均結果を報告する。我々のモデルは、ローカルPCのRTX $3070$ GPUで、最終評価を含めて各実行につき約 $3$ 時間訓練した。デモンストレーションは、最初にランダムに選択されたブロックを選択された正方形に押し、その後、別のブロックを別の正方形に押すスクリプト化されたオラクルから取得される。ポリシーは、目標状態ベクトルを使用して、ブロックを望ましい構成に押すように条件付けられる。我々は、これらの実験のために、密なディフュージョントランスフォーマーのベースラインであるBESO (Reuss et al., 2023)にインスパイアされた、履歴長 $4$ に対してアクションシーケンス長 $1$ を選択した。

	Block Push	Relay Kitchen
C-BeT	0.87 $\pm$ (0.07)	3.09 $\pm$ (0.12)
VQ-BeT	0.87 $\pm$ (0.02)	3.78 $\pm$ (0.04)
BESO	0.96 $\pm$ (0.02)	3.73 $\pm$ (0.05)
MoDE	0.97 $\pm$ (0.01)	3.79 $\pm$ (0.02)

表9: 状態ベースの目標条件付きリレーキッチンおよびブロックプッシュ環境における異なるポリシーのパフォーマンス比較（

4

シードの平均）。MoDEは、密なトランスフォーマー変種BESOおよび他のポリシー表現をすべてのベースラインで上回っている。

ベースライン。この設定では、我々はMoDEをいくつかの最先端の目標条件付きポリシーと比較する。我々は、オフセットを持つ離散化されたアクションを予測する2つのトランスフォーマーアーキテクチャ、C-BeT (Cui et al., 2023)とVQ-BeT (Lee et al., 2024a)をテストする。 C-BeTはk-meansクラスタリングとオフセットベクトルを使用し、VQ-BeTは残差ベクター量子化を利用してアクションを階層的な潜在空間に埋め込む。さらに、我々は密なディフュージョンポリシートランスフォーマーモデルBESO (Reuss et al., 2023)と比較する。 BESOは、同じ連続時間ディフュージョンポリシーを密なトランスフォーマーと組み合わせて、先行状態のシーケンスを与えられた単一のアクションを予測する。公平な比較を可能にするために、我々は両方の設定でBESOとMoDEに同じハイパーパラメータを選択した。我々はすべてのモデルを4シードで平均化してテストし、先行研究(Lee et al., 2024a)から直接平均値を報告する。

結果。両実験の結果は表9にまとめられている。 MoDEは両方のベンチマークで新しい最先端のパフォーマンスを達成し、両方の設定でBESOの密なトランスフォーマー変種を上回っている。さらに、MoDEはVQ-BeTやC-BeTなどの他のポリシー表現方法と比較してより高いパフォーマンスを達成している。

A.6.1 Mixture-of-Experts Ablations

Q: 負荷分散損失はエキスパートの分布にどのように影響するか？

我々は、負荷分散損失がノイズレベル間のエキスパート分布にどのように影響するかを分析するために、LIBERO-10で様々な負荷分散重み $\gamma_{LB}\in[0.1,0.01,0.001,0.0001]$ を用いてMoDEを訓練した。図9は、結果として得られたエキスパート分布を視覚化している。

高い負荷分散損失 $0.1$ では、8層のうち2層でわずかな変動があるものの、エキスパートはほぼすべての層で均一に使用されている（9(a)）。しかし、この強制された均一性には代償がある - 平均パフォーマンスは0.9に低下する。この結果は、ノイズレベル間でエキスパートの使用を均等に強制することが、モデルの学習能力を制限する可能性があることを示唆している。

$0.01$ では、全体的に良好なエキスパート利用を維持しながら、より柔軟な分布が観察される（9(b)）。個々の層内では、エキスパートの使用率は様々な層で

A.7 Extended Related Work

ロボティクスにおけるMoE。ロボティクスの文脈では、MoEモデルはトランスフォーマーアーキテクチャと組み合わせることなく、多くの設定で使用されている。いくつかの研究では、強化学習における異なるスキルに焦点を当てた小規模なMLPポリシーの混合(Obando-Ceron et al., 2024; Celik et al., 2022; 2024)や、ロボットの動作生成(Hansel et al., 2023; Le et al., 2023)に使用されており、別の研究群ではロボットの知覚に小規模なCNNの組み合わせを利用している(Riquelme et al., 2021; Mees et al., 2016)。さらなる応用として、ガウシアンポリシーの混合を用いたマルチモーダル行動の学習がある(Blessing et al., 2023; Li et al., 2023a)。多くの領域でMoEが広く使用されているにもかかわらず、スケーラブルでより効率的な拡散ポリシーのためにMoEを拡散ポリシーと組み合わせて利用しようとした先行研究は存在しない。

ロボット学習のためのトランスフォーマー。トランスフォーマーモデルは、ここ数年、多くのエンドツーエンドのロボット学習ポリシーにおいて標準的なネットワークアーキテクチャとなっている。これらはILの文脈で異なるポリシー表現と組み合わされている。研究の一分野は、変分オートエンコーダー（VAE）モデルを用いた行動シーケンスの生成に焦点を当てている(Bharadhwaj et al., 2023; Zhao et al., 2023)。これらのアクションチャンキングトランスフォーマーモデルは、通常、エンコーダー-デコーダートランスフォーマーをポリシーアーキテクチャとして使用する。 Octo (Octo Model Team et al., 2023)、BESO (Reuss et al., 2023)、ChainedDiffuser (Xian et al., 2023)、3D-Diffusion-Actorなどのいくつかの拡散ポリシーは、ポリシーのバックボーンとしてトランスフォーマーモデルを活用している。研究のもう一つの方向性は、自動回帰的言語生成と同様に、行動生成を離散的な次トークン予測として扱っている(Touvron et al., 2023)。 C-Bet、RT-1、RT-2は、観測された行動を

k

クラスに分割する離散化された行動ビニングを使用している(Cui et al., 2023; Shafiullah et al., 2022; Brohan et al., 2022; Zitkovich et al., 2023)。一方、VQ-BeT (Lee et al., 2024a)は残差ベクトル量子化を用いて潜在行動を学習する。いくつかの研究では、事前学習されたLLMやVLMをポリシーのバックボーンとして使用し、その後行動生成のためにファインチューニングすることの利点が示されている(Brohan et al., 2023; Gu et al., 2024; Collaboration et al., 2023; Li et al., 2024a)。最近の研究では、ポリシー学習にMixture-of-Expertアーキテクチャを使用することは考慮されていない。 MoDEは、行動生成のためにMoEアーキテクチャを拡散と組み合わせた最初のアーキテクチャである。