Completion Aware Guidance, 세계 행동 모델의 과제 미완료 상상을 추론 단계에서 줄인다
서울대, KAIST 연구진이 공개한 프리프린트에서 세계 행동 모델(WAM)이 과제 완료에 필요한 상태 전이를 빠뜨리는 '과제 미완료 상상' 문제를 지적하고, 학습 없이 추론 시점에 안내하는 Completion Aware Guidance(CAG)를 제안했다. DreamZero와 Fast-WAM에 적용했을 때 RoboTwin 2.0 아홉 과제 평균 성공률을 64.4%에서 70.0%로, DreamZero 제로샷 시뮬레이션 평균 성공률을 69%에서 75%로 높였다. 과제 미완료 상상에 기인한 실패 비율은 79%에서 40%로 줄었다.
WAM이 지역적으로 그럴듯하지만 과제를 완료하지 못하는 상상을 하는 이유가 무엇이고, 이를 어떻게 완화할 수 있는가?
WAM은 시각적 미래와 행동을 함께 예측하지만, 과제 완료에 필요한 상태 전이를 빠뜨리는 과제 미완료 상상을 보인다. 예를 들어 물건을 목표 지점까지 옮기고도 놓지 않고 계속 들고 있는 영상과 행동을 함께 생성한다. 이는 고정 수평선 비디오 생성 목표와 짧은 청크 제어 목표의 불일치에서 비롯되며, 짧은 청크에서는 지역적으로 타당한 계속이 반복 선택되어 완료 전이가 계속 미뤄진다.
기존 WAM(DreamZero, Fast-WAM 등)은 사전 학습된 비디오 확산 백본을 짧은 청크의 비디오-행동 예측에 적응시켰지만, 완료 전이를 유도하는 명시적 메커니즘이 없었다. 언어 하위 과제나 시각적 서브골을 제공하는 방법도 외부 안내가 실제 상태 전이를 보장하지는 못했다.
CAG는 학습 없는 샘플링 기법으로, 확산 과정의 크로스 어텐션에서 현재 예측에 이미 사용 중인 명령어 토큰의 중요도를 계산하고, 그 토큰들의 키, 밸류에 경계 있는 스케일링(γ=1+αψ(r))을 적용해 명령어 조건부 신호를 강화한다. 이를 통해 현재 청크가 활성 단계를 완료하는 전이를 포함하도록 편향시킨다. 별도 재학습 없이 DreamZero와 Fast-WAM에 적용했으며 α=0.15를 모든 과제와 백본에 고정했다.
RoboTwin 2.0 아홉 과제(Fast-WAM Random 기준 성공률 90% 미만 비포화 과제)에서 평균 성공률이 64.4%에서 70.0%로 올랐다. 단, Open Microwave는 33%에서 30%, Place Can Basket은 63%에서 53%로 하락했다. DreamZero 제로샷 3개 시뮬레이션 과제에서는 평균 성공률이 69%에서 75%로 올랐다. 세부로 Cube to Bowl 90%→92%, Can to Mug 60%→72%, Banana to Bin 56%→60%였다. DreamZero 평가에서 과제 미완료 상상에 의한 실패 비율은 79%에서 40%로 감소했다. DreamZero는 과제당 50롤아웃, RoboTwin 2.0은 과제당 30롤아웃을 사용했다.
저자는 CAG가 샘플링 시점 개입에 국한되며, 완료 인지 안내를 WAM 학습으로 확장하는 것을 후속 연구로 남겼다. 평가는 시뮬레이션에 그쳤고 실물 로봇 검증은 보고되지 않았다. 논문에 코드 공개 여부도 보고되지 않았다. 일부 RoboTwin 과제에서는 성능이 하락해 모든 과제에서 일관된 개선을 보이지는 않았다.
WAM 기반 로봇 조작 정책에 추가 학습 없이 추론 단계에서 적용할 수 있어 기존 배포 시스템에 비교적 쉽게 통합할 수 있다. 다만 현재 평균 성능 향상이 5~6%p 수준이고 일부 과제에서는 하락하므로, 실물 로봇 검증과 더 넓은 과제군에서의 안정성 확인이 선행되어야 한다. 물류 피킹이나 조립처럼 단계 완료가 중요한 장기 조작 작업에서 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.