Learning Long-Context Diffusion Policies via Past-Token Prediction

要約

多くのロボットタスクにとって、観察とアクションの長いシーケンスにわたる推論が不可欠です。
しかし、デモンストレーションから効果的なロングコンテキストポリシーを学ぶことは依然として困難です。
コンテキストの長さが増加すると、メモリの需要が高まるため、トレーニングはますます高くなり、ポリシーのパフォーマンスは、偽の相関の結果として多くの場合劣化します。
最近の方法は、通常、コンテキストの長さを切り捨て、その後の決定に重要な履歴情報を破棄することにより、これらの問題を回避します。
この論文では、過去の情報の保持を明示的に正規化する代替アプローチを提案します。
最初に、模倣学習の模倣問題を再検討し、最近の拡散ポリシーで反対の課題を特定します。以前の行動を過剰に依存するのではなく、過去と将来の行動の間の本質的な依存関係をキャプチャできないことがよくあります。
これに対処するために、過去のトークン予測（PTP）を紹介します。これは、ポリシーが将来のアクショントークンと一緒に過去のアクショントークンを予測することを学ぶ補助タスクです。
この正則化は、視覚的表現への依存を最小限に抑えて、ポリシーヘッドの時間モデリングを大幅に改善します。
この観察結果に基づいて、マルチステージトレーニング戦略をさらに紹介します。短いコンテキストで視覚エンコーダーを事前トレーニングし、キャッシュされた長いコンテキスト埋め込みを使用してポリシーヘッドを微調整します。
この戦略は、メモリと計算オーバーヘッドを大幅に削減しながら、PTPの利点を維持します。
最後に、PTPをテスト時に自己検証メカニズムに拡張し、推論中の過去のアクションと一致する候補者を採点および選択できるようにします。
4つの現実世界と6つのシミュレートされたタスクにわたる実験は、提案された方法が長いコンテキスト拡散ポリシーのパフォーマンスを3倍に改善し、ポリシートレーニングを10倍以上加速することを示しています。

要約(オリジナル)

Reasoning over long sequences of observations and actions is essential for many robotic tasks. Yet, learning effective long-context policies from demonstrations remains challenging. As context length increases, training becomes increasingly expensive due to rising memory demands, and policy performance often degrades as a result of spurious correlations. Recent methods typically sidestep these issues by truncating context length, discarding historical information that may be critical for subsequent decisions. In this paper, we propose an alternative approach that explicitly regularizes the retention of past information. We first revisit the copycat problem in imitation learning and identify an opposite challenge in recent diffusion policies: rather than over-relying on prior actions, they often fail to capture essential dependencies between past and future actions. To address this, we introduce Past-Token Prediction (PTP), an auxiliary task in which the policy learns to predict past action tokens alongside future ones. This regularization significantly improves temporal modeling in the policy head, with minimal reliance on visual representations. Building on this observation, we further introduce a multistage training strategy: pre-train the visual encoder with short contexts, and fine-tune the policy head using cached long-context embeddings. This strategy preserves the benefits of PTP while greatly reducing memory and computational overhead. Finally, we extend PTP into a self-verification mechanism at test time, enabling the policy to score and select candidates consistent with past actions during inference. Experiments across four real-world and six simulated tasks demonstrate that our proposed method improves the performance of long-context diffusion policies by 3x and accelerates policy training by more than 10x.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google