World Action Modelsのタスク未完了を抑えるCompletion Aware Guidance
本研究は、World Action Models（WAM）が視覚的に妥当で行動とも整合するもののタスク完了に必要な遷移を省く「task-incomplete imagination」を抑える学習不要のサンプリング手法Completion Aware Guidance（CAG）を提案した。CAGはDreamZeroとFast-WAMで平均成功率をRoboTwin 2.0の一部タスクで64.4%から70.0%に、ゼロショットシミュレーションで69%から75%に向上させ、失敗ロールアウトに占めるtask-incomplete imaginationの割合を79%から40%に低減した。本論文は査読前のプレプリントである。
WAMがタスク未完了の想像をする理由は何か、またそれを再学習なしに解決するにはどうすればよいか。
WAMは固定長の映像生成を短いチャンクでの制御に適応させたため、現在のチャンク内で完了すべき状態遷移を繰り返し先送りする。その結果、物体を目標の上まで運びながら離さないといった、指示と整合し局所的にも妥当だがタスクを完了しない映像と行動の組が生成される。
従来のWAM（DreamZero、Fast-WAMなど）は映像拡散バックボーンを短い動画行動チャンクの予測にファインチューニングし、生成したチャンクを実行しては再生成する。外部からサブタスク指示やサブゴール画像を与える手法もあるが、指示を与えても完了に必要な状態遷移を実際に生成する保証はなかった。
CAGは学習不要の推論時サンプリング手法である。映像行動クエリと指示トークンのクロスアテンションを観測し、現在の生成が既に依存している指示トークンを特定する。そのトークンのキーと値をスパースな二値ゲートで選択的に増幅し、現在のチャンクをタスク完了へ誘導する。介入はクロスアテンション経路に限定し、WAMの事前分布から大きく外れないように正則化する。
評価はDreamZeroの公式DROIDシミュレーション3タスクで各50ロールアウト、Fast-WAMのRoboTwin 2.0の成功率90%未満の9タスクで各30ロールアウトを実施した。CAGはRoboTwin 2.0部分集合の平均成功率を64.4%から70.0%に改善し、DreamZeroのゼロショット3タスクでも69%から75%に改善、全タスクで上昇した。task-incomplete imaginationに起因する失敗の割合は79%から40%に低下した。介入強度αは全タスクとバックボーンで0.15に固定し、タスクごとの調整は行っていない。なおRoboTwin 2.0のタスク別ではOpen Microwaveが33%から30%、Place Can Basketが63%から53%に低下する例もあった。
著者はCAGがサンプリング時介入に限定されており、WAMの学習段階への拡張が今後の課題と述べている。評価はシミュレーションのみで、実機での検証は報告されていない。追加の推論コストやコード公開の有無も論文では報告されていない。
WAMを操作タスクに利用するロボット制御スタックに対して、再学習なしで即座に追加できるため、倉庫やサービスロボット向けの操作成功率向上に寄与する可能性がある。現時点ではシミュレーションでの有効性が示された段階であり、実機展開には追加検証が必要である。条件が整えば1〜3年程度で実ロボットへの応用が進むと考えられる。
論文全文
Completion Aware Guidance for World Action Models
CC BY 4.0 のもとで公開された論文です。出典を明記して転載しています。原文は arXiv:2610.01559（PDF）。
\workshoptitle
Robot Learning with World Models: Capabilities, Frontiers, and Challenges
Abstract
World Action Models (WAMs) predict visual futures and robot actions, yet they remain susceptible to task-incomplete imagination, where plausible, action-consistent predictions omit the transition needed for task completion. In this paper, we show that this failure is not inherent to the world model backbone, but emerges when adapted for short-chunk control, which can repeatedly favor plausible local continuations over task-completing transitions. To address this, we introduce Completion Aware Guidance (CAG), a training-free sampling method that guides generation toward task completion. Across representative WAMs, CAG improves success from 64% to 70% on a RoboTwin 2.0 subset and from 69% to 75% in zero-shot simulation, while reducing task-incomplete imagination from 79% to 40%.
1 Introduction
Robot tasks unfold through changes in the world state, and therefore require temporal understanding of what is likely to happen after an action is executed. Vision–language–action policies [1, 2, 3, 4, 5] predict actions from the current observation, yet the resulting future changes are difficult to infer from observed images alone. World Action Models (WAMs) address this limitation by using video generative models to jointly predict a visual future and the actions that realize it [6, 7, 8]. However, generating visual futures with aligned actions does not guarantee task success.
Interestingly, as shown in Fig. 1, we find that WAMs exhibit task-incomplete imagination, generating visually coherent, instruction-relevant futures while omitting the state transitions required for task completion. For example, given a pick-and-place instruction, the imagined video may move an object above its target but continue holding it instead of releasing it, and the predicted actions follow this future, resulting in task failure. This differs from a corrupted video or video–action misalignment in that both modalities agree on a plausible continuation that does not complete the task.
In this paper, we ask why WAMs exhibit task-incomplete imagination and how this failure can be addressed. We argue that it arises from an objective mismatch between the fixed-horizon video-generation objective for video-backbone pretraining and the short-horizon video-action prediction objective for WAM fine-tuning. Under the former objective, video diffusion models jointly predict an instruction-following trajectory over a fixed horizon, whereas WAMs predict video-action chunks that are executed and regenerated from the observation. Because the WAM objective does not require the current phase to be completed within each horizon, the state transition required for task completion can be repeatedly deferred. As a result, although the video backbone can generate instruction-following trajectories, this capability is not reliably elicited when needed.
To address this mismatch, we introduce Completion Aware Guidance (CAG), a training-free sampling method for WAMs. CAG strengthens the conditioning of instruction tokens most relevant to the current video–action prediction, steering the current chunk toward the required state transition without retraining the WAM. Across DreamZero and Fast-WAM [6, 7], CAG improves the average task success rate from 64.4% to 70.0% on a RoboTwin 2.0 subset and from 69% to 75% in zero-shot simulation. More importantly, in a representative manipulation setting exhibiting task-incomplete imagination, CAG further reduces the incidence of this failure from 79% to 40%.
2 Related Work
Vision-Language-Action Models. Generalist robot policies have largely been developed as vision–language–action models, which map visual observations and language instructions directly to actions or action chunks [1, 2, 3, 4, 5]. For temporally extended tasks, however, a global instruction may not specify the next semantic subtask or desired near-future state. Recent approaches [9] therefore provide language subtask instructions or visual subgoal images to make task progress explicit [10, 11]. Yet specifying an intermediate target and reaching it are distinct problems: external guidance does not ensure that the policy produces the state transition required to complete the subtask.
World Models and World Action Models. Video generative models capture temporal dynamics in observations, making them natural backbones for predictive world models in robot planning and control [12, 13, 14, 15, 16, 17, 18, 19, 20]. WAMs extend this idea by coupling the imagined visual trajectory with actions, allowing a generated future to act as an implicit temporal plan [6, 7, 21, 22, 8]. Prior works[6, 7, 23, 24] do not isolate the case in which the imagined video and predicted actions remain mutually consistent and locally plausible while omitting a phase-completing transition. We characterize this failure as task-incomplete imagination and mitigate it with training-free sampling guidance.
3 Method
Task-Incomplete Imagination. As illustrated in Fig. 1, WAMs can repeatedly produce locally reasonable video-action continuations while failing to progress to the transition required for the next task phase. We refer to this failure as task-incomplete imagination; for example, the model may reach the target region yet continue holding the object instead of releasing it. Does this failure simply reflect a limitation of the underlying video diffusion backbone?
Using the same video diffusion backbone, we compare its native generation horizon with the short chunk horizon used by the WAM. As shown in Fig. 2 (a), the native horizon produces the required completion event, whereas short-chunk autoregressive generation remains visually coherent but repeatedly omits it. This suggests that task-incomplete imagination is not an intrinsic limitation of the backbone, but arises from its adaptation to short-horizon control generation as follows:
\underbrace{(o_{0},c)\longmapsto o_{1:T}}_{\text{WM: task-level visual trajectory}}\quad\xrightarrow{\text{WAM adaptation}}\quad\underbrace{(o_{\leq t},q_{t},c)\longmapsto\bigl(o_{t+1:t+H},a_{t:t+H-1}\bigr)}_{\text{WAM: local video--action continuation}}.Here, o_{t} denotes the observation at time t, a_{t} the robot action, c the language instruction, q_{t} the proprioceptive state, T the native video-generation horizon, and H the WAM chunk horizon. Once the task reaches a completion-relevant phase, the WAM must select the required transition within a short horizon. However, a holding or hovering continuation can remain locally plausible and compatible with the instruction, providing no explicit preference for realizing the completion in the current chunk. Repeated generation can therefore defer the required transition. This raises the central question: how can we guide short-horizon generation to include the completion required by the current task phase?
Completion Aware Guidance. Our algorithm, Completion Aware Guidance (CAG), tackles task-incomplete imagination as an elicitation problem. The video diffusion backbone represents the task-completing transition specified by the instruction, but short-horizon control can favor a plausible continuation that preserves the pre-completion state. CAG biases this selection toward a continuation that completes the current manipulation phase while remaining close to the WAM prior.
Let h_{t}=(o_{\leq t},q_{t}) denote the current WAM context and z_{t}=(o_{t+1:t+H},a_{t:t+H-1}) the predicted video-action chunk. We introduce a phase-completion indicator E_{t}(z_{t})\in\{0,1\}, with E_{t}(z_{t})=1 when the chunk contains the transition that completes the active task phase. Let p_{\theta}(\cdot\mid h_{t},c) denote the chunk distribution of the WAM with parameters \theta. The ideal completion aware distribution is
p_{t}^{\star}=\arg\max_{p(\cdot\mid h_{t},c)}\;\mathbb{E}_{z_{t}\sim p(\cdot\mid h_{t},c)}[E_{t}(z_{t})]-\lambda D_{\mathrm{KL}}\!\left(p(\cdot\mid h_{t},c)\,\|\,p_{\theta}(\cdot\mid h_{t},c)\right).The first term favors video-action chunks that complete the active phase, while the KL regularizer keeps the guided distribution close to the original WAM prior. Here, p(\cdot\mid h_{t},c) denotes the candidate chunk distribution and \lambda>0 controls the strength of this regularization. The frozen model does not allow direct control over p_{t}^{\star}. Its instruction signal enters video-action generation through cross-attention. At denoising layer \ell, video-action queries Q^{va}_{t,\ell} receive the conditional context
C^{va\leftarrow c}_{t,\ell}=\operatorname{softmax}\!\left(\frac{Q^{va}_{t,\ell}(K^{c}_{\ell})^{\top}}{\sqrt{d}}\right)V^{c}_{\ell},where K_{\ell}^{c} and V_{\ell}^{c} are the keys and values of the instruction tokens, and d is the attention key dimension. This pathway identifies where the instruction components used by the current video-action generation enter the model. We restrict the intervention to this pathway. For a token-wise scaling vector \gamma_{t}, we modify the key and value associated with each instruction token j as
\widetilde{K}_{\ell,j}^{c}=\gamma_{t,j}K_{\ell,j}^{c},\qquad\widetilde{V}_{\ell,j}^{c}=\gamma_{t,j}V_{\ell,j}^{c},which induces the conditional chunk distribution p_{\theta}^{\gamma_{t}}(z_{t}\mid h_{t},c). Ideal intervention is
\gamma_{t}^{\star}=\arg\max_{\gamma_{t}}J_{t}(\gamma_{t}),\qquad J_{t}(\gamma_{t})=\mathbb{E}_{z_{t}\sim p_{\theta}^{\gamma_{t}}(\cdot\mid h_{t},c)}[E_{t}(z_{t})]-\lambda D_{\mathrm{KL}}\!\left(p_{\theta}^{\gamma_{t}}(\cdot\mid h_{t},c)\,\|\,p_{\theta}(\cdot\mid h_{t},c)\right).Around the original conditioning \gamma_{t}=\mathbf{1}, the effect of an intervention is locally characterized by
J_{t}(\gamma_{t})-J_{t}(\mathbf{1})\approx\left\langle g_{t},\,\gamma_{t}-\mathbf{1}\right\rangle,\qquad g_{t,j}=\left.\frac{\partial J_{t}}{\partial\gamma_{t,j}}\right|_{\gamma_{t}=\mathbf{1}}.The vector g_{t} specifies which instruction components should be strengthened to increase completion elicitation. However, neither the phase-completion indicator E_{t} nor the resulting sensitivity g_{t} is available during sampling. CAG instead uses the cross-attention pathway to identify the instruction components that the current video-action generation already relies on. For each instruction token w_{j}, we aggregate its attention affinity over video-action queries and a set of selected denoising layers \mathcal{L},
r_{t,j}=\frac{1}{|\mathcal{L}|}\sum_{\ell\in\mathcal{L}}\frac{1}{|Q^{va}_{t,\ell}|}\sum_{\mathbf{u}\in Q^{va}_{t,\ell}}\left[\operatorname{softmax}\left(\frac{\mathbf{u}(K_{\ell}^{c})^{\top}}{\sqrt{d}}\right)\right]_{j}.CAG converts this observable signal into a bounded conditional intervention,
\widehat{\gamma}_{t,j}=1+\alpha\,\psi(r_{t,j}),where \psi is a sparse binary gating function. Applying Eq. (8) through Eq. (4) strengthens the instruction signal already used for the current action, biasing the next video-action chunk toward the completion required by the active phase.
4 Experiments
Setting. CAG is a training-free inference-time method applicable to WAMs built on video diffusion backbones. To test its generality across different WAM designs, we evaluate it on DreamZero and Fast-WAM, which respectively use explicit future generation and implicit world modeling for action prediction. We report task success over 50 rollouts on each of the three official DreamZero DROID simulation scenes and 30 rollouts per task on nine non-saturated RoboTwin 2.0 tasks with published success rates below 90% under the Fast-WAM Random setting. We follow the corresponding evaluation protocols, changing only inference-time guidance between the baseline and CAG, and fix \alpha=0.15 across all tasks and backbones without per-task tuning.
- Open Microwave
- 37
- 14
- 33
- 30
- Hanging Mug
- 3
- 0
- 60
- 63
- Turn Switch
- 6
- 30
- 53
- 67
- Place Can Basket
- 25
- 37
- 63
- 53
- Move Stapler Pad
- 18
- 4
- 57
- 73
- Stack Bowls Three
- 35
- 7
- 67
- 87
- Pick Diverse Bottles
- 3
- 56
- 80
- 83
- Place Mouse Pad
- 26
- 10
- 80
- 83
- Place Object Basket
- 36
- 49
- 87
- 90
Results. CAG improves task success on both backbones. On the nine-task RoboTwin 2.0 subset, CAG increases average success from 64.4\% to 70.0\% (Tab. 1). On three DreamZero zero-shot tasks, CAG improves average success from 69\% to 75\%, with gains on every task (Fig. 2 (b)). The consistent gains across DreamZero and Fast-WAM suggest that CAG is not tied to explicit future-video generation, but generalizes across different ways of incorporating video-based world modeling into action prediction.
Failure Analysis. To examine whether CAG directly addresses task-incomplete imagination, we analyze failure cases in the DreamZero zero-shot evaluation. As shown in Fig. 2 (c), the fraction attributable to task-incomplete imagination decreases from 79\% to 40\% with CAG, indicating that CAG gains from mitigating task-incomplete imagination.
5 Conclusion
We identify task-incomplete imagination, in which a WAM produces a visually plausible, action-consistent continuation while omitting the state transition required to complete the current task phase. We attribute this failure to the mismatch between fixed-horizon video generation and short-horizon WAM control. CAG addresses this mismatch by strengthening instruction conditioning relevant to the current prediction, eliciting task-completing transitions during sampling. Across representative WAMs, CAG improves task success and reduces task-incomplete imagination without retraining.
Limitations. CAG is currently limited to sampling-time intervention, and extending completion aware guidance to WAM training is a promising direction for future work.