World Action Modeling with Progressive Visual Planning

AI summary

Overview

Research area: Robot learning — specifically world action models (WAMs) that combine video-generation priors with action policies for robotic manipulation. The work sits at the intersection of diffusion/flow-matching generative modeling, vision-language-action (VLA) policies, and hierarchical long-horizon control.

Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion transformers (DiTs), mixture-of-transformers (MoT) architectures, VAE latents, and standard robot-manipulation benchmarking.

One-sentence scope: ProWAM is a world action model that replaces dense future-video rollouts with a sparse, ordered sequence of progress-indexed visual sub-goals that directly condition action generation, improving long-horizon and out-of-distribution robotic control while cutting inference cost.

What This Paper Is About

World action models predict future visual dynamics and actions together, which gives robot policies useful visual foresight — but generating a full dense video rollout at every control step is slow and inefficient. Existing shortcuts either drop visual foresight entirely (fast but unguided over long horizons) or predict only a single endpoint frame (which specifies where the task ends but not how to get there). ProWAM's goal is to give the policy explicit intermediate visual guidance at multiple points along the remaining trajectory, without paying the cost of dense video generation or adding a separate high-level planner module.

Key Contributions

1. Progressive sub-goal imagination. A new visual-planning paradigm for WAMs that represents a long-horizon future as an ordered sequence of k sparse visual sub-goals, each indexed by a normalized task-progress coordinate r ∈ [0, 1], where r = 0 is the current observation and r = 1 is the trajectory endpoint. This supplies intermediate guidance without dense rollouts and without an external planner.

2. The ProWAM architecture. A vision-action mixture-of-transformers (MoT) with a video generation expert and a lightweight diffusion-based action expert. Sub-goal slots are appended to the video latent sequence, conditioned on their progress value via adaptive LayerNorm (adaLN) modulation, and connected to actions through an asymmetric attention mask (video tokens never attend to action tokens). This design permits action-free video pretraining followed by joint vision-action fine-tuning.

3. Efficient inference via sub-goal caching. Instead of joint video-action denoising across T steps, ProWAM runs the video expert once per replanning step, caches the sub-goal key–value features, and reuses them across all action-denoising steps — reducing video-expert forward passes from T to 1 (approximately 10× fewer evaluations).

4. Broad empirical validation. New state-of-the-art results on LIBERO-Plus (85.8%) and randomized RoboTwin (75.7%), a 4th-place overall result on RoboCasa365 (48.1% overall, 18.2% on Composite-Unseen), and 70.0% zero-shot success on four real-world Franka tasks.

Main Findings

• LIBERO / LIBERO-Plus: Both ProWAM variants reach near-saturated in-domain LIBERO performance (ProWAM 98.8%, ProWAM-Base 99.2%). On the zero-shot LIBERO-Plus split, ProWAM reaches 85.8%, compared with Image-WAM at 83.1%, Being-H0.7 at 82.1%, ABot-M0 at 80.5%, π0.5 at 77.4% and Fast-WAM at 51.5%. Removing sub-goal pretraining (ProWAM-Base) drops LIBERO-Plus to 80.3%, a 5.5 point loss, which the authors attribute to the value of sub-goal pretraining for out-of-distribution robustness.

• RoboTwin: Under scene randomization ProWAM achieves 75.7% against 55.7% for the strongest prior WAM (MV-WAM), a +20.0 absolute gain. Dropping video pretraining (ProWAM-Base) reduces randomized success from 75.7% to 67.4% while clean-scene performance stays comparable (85.8% versus 87.6%), indicating that video pretraining mainly improves robustness to visual domain shift.

• RoboCasa365: Table 3 ranks ProWAM 4th overall with a 48.1% overall success rate, 76.0% on Atomic-Seen, 46.5% on Composite-Seen, and 18.2% on the challenging Composite-Unseen split — ahead of ABot-M0.6 (46.6% overall, 7.9% Composite-Unseen). The abstract also states the rank as 4th; note that the introduction and Section 4.2 text describe the result as "second overall," which conflicts with the table.

• Zero-shot real-world: On 4 physical tasks (pick-and-place, long-horizon multi-object assembly, contact-rich placement, spatially grounded placement) trained on DROID, ProWAM (5B) reaches 70.0% average success versus DreamZero (14B) at 55.0% and π0.5 (3B) at 50.0% — a +15.0 absolute and +27.3% relative gain. On spatial placement it reaches 100.0% success.

• Efficiency: For a 16-action chunk on a single NVIDIA H100, naive joint denoising costs 115.5 TFLOPs and 581 ms; the KV-cache strategy reduces this to 18.7 TFLOPs (84% reduction) and 313 ms (46% reduction) while achieving 3.3 percentage points higher OOD success in that evaluation.

• Sub-goal count k: One model trained with k ∈ {1, …, 4} supports all four settings at inference. In-domain and OOD performance is stable across k ∈ [1, 3], peaking at 75.7% OOD at k = 4. Going from k = 1 to k = 4 adds only 3.5 TFLOPs. Even at k = 1, ProWAM reaches 73.8% OOD at 15.2 TFLOPs, versus FastWAM's 6.3% at 13.2 TFLOPs.

• Sub-goal horizon: Shorter-horizon schedules help. OOD success rises monotonically from 68.4% (Far, r ∈ [0.5, 0.9]) to 70.2% (Mid, r ∈ [0.3, 0.7]) to 75.7% (Near, r ∈ [0.1, 0.5]), a +7.3 improvement, while in-domain performance stays around 86%. The Near schedule is the default.

• What sub-goals look like: Visualized sub-goals at r ∈ {0.3, 0.5, 0.7} advance with progress and re-anchor to each new observation. Single-pass generation introduces minor texture artifacts but preserves scene structure and task-relevant geometry.

• Failure analysis (Cube Stacking): In a successful rollout the sub-goal shows a centered placement with a clear contact boundary; in the failure case the cubes appear blurred together, the contact geometry is ambiguous, and the gripper releases off-center. The authors read this as evidence that execution tolerates texture imperfections but is sensitive to errors in task-relevant geometry.

• Pilot study (Appendix A): Adding final-goal prediction (FGP) to an existing WAM improves results in most evaluated settings, with larger gains under reduced training-data fractions ({0.3, 0.5, 1.0}).

Methodology in Plain English

ProWAM starts from a pretrained image-to-video diffusion model (Wan2.2-TI2V-5B) and adds a small action-prediction "expert" alongside the large video expert, in a mixture-of-transformers design where both experts use shared attention heads.

The key idea is how the future is represented. Instead of generating an entire video of what happens next, the model imagines a handful of still frames at chosen points along the remaining trajectory. Each imagined frame is tagged with a progress value — a number between 0 (right now) and 1 (task finished) — so the model learns to ask itself "what should things look like when I'm 30% of the way there?" At training time the number of sub-goals k is randomized, which lets one trained model accept different numbers of sub-goals at test time.

Architecturally, sub-goal slots are appended to the sequence of video latent frames and processed by the same weights. Their progress value is converted into a modulation offset that is added to the usual diffusion timestep modulation, so each slot knows which stage of the task it represents. An attention mask enforces that information flows one way: sub-goals and video frames influence the action expert, but action tokens never influence the visual pathway, protecting the video model's pretrained generative prior.

Training has two stages. Stage 1 trains only the video side on large-scale action-free video (including robot and egocentric data) to learn visual dynamics and progress-conditioned sub-goal prediction. Stage 2 jointly fine-tunes both experts on robot demonstrations so actions become grounded in the imagined sub-goals. The training objective is a weighted sum of three flow-matching velocity-regression losses: one for the near-term video transitions, one for the sub-goals, and one for the action chunks, with the video and action branches using independently sampled diffusion timesteps.

At inference, the model runs the expensive video expert once per replanning step, produces the sub-goal features, caches their key–value representations, and then runs only the lightweight action denoiser repeatedly against that fixed cache. When a new observation arrives, the sub-goals are recomputed and the cache refreshed.

Why This Matters

Impact on research. ProWAM argues that the field's binary framing — either generate a full future video (accurate but slow) or generate none (fast but blind over long horizons) — is a false choice. It shows that a sparse, progress-indexed set of visual waypoints is enough to anchor long-horizon control, and that this representation can be learned from action-free video at scale. It also demonstrates that sub-goal prediction can be folded into the world model itself rather than living in a separate planner module, avoiding the error propagation and multi-stage complexity of decoupled hierarchical pipelines.

Real-world applications:

• Industrial and warehouse robotics: the zero-shot Franka results cover pick-and-place and multi-object assembly, the kind of repetitive part-handling where retraining per environment is expensive.
• Contact-rich manipulation: the paper explicitly benchmarks contact-rich placement and analyzes a "Cube Stacking" failure case involving contact geometry, relevant to insertion and packing tasks.
• Household and service robots: benchmarked on RoboCasa365, a long-horizon household-task suite, where held-out task compositions (Composite-Unseen) are the hardest test.
• Cross-embodiment / cross-site deployment: training on DROID and testing zero-shot on new physical scenes targets the practical case where a robot must work in an environment it has never seen demonstrations for.

Industry relevance. The efficiency numbers are the operational story: 18.7 TFLOPs and 313 ms versus 115.5 TFLOPs and 581 ms for naive joint denoising, achieved with a 5B backbone that outperforms a 14B baseline (DreamZero) in the reported zero-shot comparison. Caching also means the expensive video expert runs once per replanning step rather than once per denoising step, which matters for real-time closed-loop control on a single GPU.

Future Directions

• Diagnosing and fixing geometric failures. The Cube Stacking analysis suggests the policy is sensitive to errors in task-relevant geometry in the imagined sub-goal, even while tolerating texture blur. Whether sharper or more geometrically faithful sub-goal generation closes this gap is left open.
• Progress schedules beyond fixed windows. The Near/Mid/Far comparison shows shorter horizons win, but the paper does not explore adaptive or learned schedules that vary r per task or per replanning step.
• Scaling k and K_max. The model was trained with k ∈ {1, …, 4} and peaks at k = 4; performance beyond that, and the interaction between larger k and very long horizons, is not reported.
• Narrowing the gap on the hardest held-out compositions. ProWAM's 18.2% on RoboCasa365 Composite-Unseen versus 36.7% for the top-ranked method indicates substantial room remains on held-out task compositions.
• Broadening real-world evaluation. The real-world result covers 4 tasks on Franka with DROID training data; generalization to other embodiments, task families and longer horizons is not reported.

Target Audience

Researchers and engineers working on robot learning, embodied AI, and vision-language-action models who already understand diffusion or flow-matching generative models — particularly those interested in how video-generation priors can be repurposed for control. It is also relevant to practitioners evaluating whether generative world models are computationally viable for real-time closed-loop manipulation, since the paper reports concrete latency and FLOP comparisons. Readers looking for introductory material on imitation learning or diffusion policies will find the architectural details dense and should start with the preliminaries in Section 2.

Authors’ abstract

World action models (WAMs) have emerged as a promising paradigm for robotic control by jointly predicting future visual dynamics and actions from an initial observation and instruction. However, existing WAMs struggle with long-horizon prediction, as generating dense video rollouts is highly inefficient. Some recent WAMs address this by predicting a single future frame without generating the full video, but this approach neglects how to progress toward the goal. We present ProWAM, a progressive world action model that jointly predicts actions and an ordered sequence of sparse visual sub-goals, providing explicit visual guidance to anchor action generation throughout task execution. This design scales naturally, as sub-goal prediction can be learned from large-scale action-free videos, allowing the video backbone to offload complex visual planning from the action policy. For efficient action generation, ProWAM executes a single video-backbone forward pass to cache sparse sub-goal features, eliminating iterative full-video generation and requiring only lightweight action denoising during replanning. Across extensive evaluations, ProWAM achieves superior out-of-distribution robustness. On simulation benchmarks, it sets new state-of-the-art results on LIBERO-Plus (85.8%) and randomized RoboTwin (75.7%), outperforming the strongest baseline with relative gains of up to +35.9%. On RoboCasa365, ProWAM achieves a 48.1% success rate and 18.2% on the challenging Composite-Unseen split, ranking 4th overall. Crucially, in zero-shot real-world experiments, ProWAM achieves 70.0% success, outperforming the strongest baseline by +15.0 (from 55.0% to 70.0%, a +27.3% relative gain) in novel scenes. These results demonstrate the value of progress-indexed visual foresight for closed-loop control. Our program is in https://sii-ferenas.github.io/ProWAM-page.