Completion Aware Guidance verbessert World Action Models ohne Nachtraining
Der Preprint stellt Completion Aware Guidance (CAG) vor, eine trainingsfreie Sampling-Methode für World Action Models. CAG verstärkt während der Inferenz die für die aktuelle Vorhersage relevanten Instruktionsanteile und reduziert so die als task-incomplete imagination bezeichnete Fehlerart. In Simulationen steigt die Erfolgsquote bei DreamZero und Fast-WAM auf einem RoboTwin-2.0-Subset von 64,4 % auf 70,0 % und in Zero-Shot-Aufgaben von 69 % auf 75 %.
Warum erzeugen World Action Models trotz visuell plausibler und aktionskonsistenter Vorhersagen die für die Aufgabenerfüllung nötige Zustandsänderung nicht, und wie lässt sich dieses Problem ohne Nachtraining beheben?
World Action Models sagen visuelle Zukunft und Roboteraktionen gemeinsam voraus, neigen aber zu task-incomplete imagination: Sie generieren lokal plausible, zur Instruktion passende Fortsetzungen, in denen der für den nächsten Aufgabenschritt erforderliche Zustandswechsel fehlt, etwa das Loslassen eines Objekts über dem Ziel. Die Ursache liegt in einem Zielkonflikt zwischen dem Fixed-Horizon-Videotraining des Backbones und dem Short-Horizon-Feintuning für Steuerung, wodurch kurze Chunks wiederholt plausible lokale Fortsetzungen bevorzugen und den Abschluss der aktuellen Phase verschieben.
Bisherige World Action Models wie DreamZero und Fast-WAM nutzen Video-Diffusions-Backbones und sagen kurze Video-Aktions-Chunks voraus, ohne die Konsistenz zwischen plausibler Fortsetzung und tatsächlicher Aufgabenphase explizit zu prüfen. Vision-Language-Action-Policies sagen Aktionen direkt aus Beobachtungen voraus, erfassen aber keine zukünftigen Zustandsänderungen. Ansätze mit externen Subtask- oder Subgoal-Vorgaben spezifizieren Zwischenziele, garantieren jedoch nicht, dass die Politik den nötigen Zustandswechsel im aktuellen Schritt realisiert.
Completion Aware Guidance ist ein reines Inferenzverfahren ohne Nachtraining. Es nutzt den Cross-Attention-Pfad zwischen Video-Aktions-Queries und Instruktionstokens, berechnet für jedes Token eine Aufmerksamkeitsaffinität und verstärkt über einen binären Gating-Mechanismus gezielt die Schlüssel- und Wertvektoren relevanter Instruktionstokens mit einem festen Faktor Alpha von 0,15. Dadurch wird die Generierung des aktuellen Video-Aktions-Chunks in Richtung des für die aktuelle Phase erforderlichen Zustandswechsels gelenkt, während die Verteilung nahe am trainierten WAM-Prior bleibt. Angaben zu spezifischer Hardware oder Rechenaufwand macht das Paper nicht.
Auf einem Subset von neun nicht gesättigten RoboTwin-2.0-Aufgaben mit je 30 Rollouts pro Aufgabe unter der Fast-WAM-Random-Einstellung steigt die durchschnittliche Erfolgsquote mit CAG von 64,4 % auf 70,0 %. Einzelne Aufgaben verbessern sich deutlich, etwa Stack Bowls Three von 67 % auf 87 % und Move Stapler Pad von 57 % auf 73 %, während Open Microwave von 33 % auf 30 % und Place Can Basket von 63 % auf 53 % zurückgehen. In drei Zero-Shot-DreamZero-Simulationsaufgaben mit je 50 Rollouts steigt der Durchschnitt von 69 % auf 75 %, mit Gewinnen in allen Aufgaben: Cube to Bowl von 90 % auf 92 %, Can to Mug von 60 % auf 72 % und Banana to Bin von 56 % auf 60 %. Der Anteil fehlgeschlagener Rollouts, der auf task-incomplete imagination zurückgeführt wird, sinkt von 79 % auf 40 %.
Die Autoren nennen als Grenze, dass CAG nur während des Samplings eingreift und eine Erweiterung auf das Training künftiger Arbeit bedarf. Die Evaluation beschränkt sich auf Simulation und zwei WAM-Varianten; ein realer Roboter wird nicht berichtet. Zudem wird Alpha mit 0,15 global fixiert, ohne aufgabenspezifische Optimierung. Angaben zu veröffentlichtem Code, Hardware oder Rechenaufwand fehlen im Paper.
Unternehmen, die World Action Models für Manipulationsaufgaben einsetzen, könnten CAG kurzfristig als Inferenz-Add-on nutzen, da es kein Nachtraining erfordert und mit bestehenden Modellen wie DreamZero oder Fast-WAM kompatibel ist. Für den produktiven Einsatz in Logistik- oder Montageanwendungen sind zunächst eine Validierung auf realen Robotern, die Veröffentlichung des Codes und eine breitere Aufgabeabdeckung nötig. Innerhalb von ein bis drei Jahren könnte CAG die Zuverlässigkeit von WAM-basierten Politiken ohne zusätzliche Trainingskosten erhöhen.
Vollständige Studie
Completion Aware Guidance for World Action Models
Veröffentlicht unter CC BY 4.0. Wiedergabe mit Namensnennung, Original unter 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.