BoxWAM: Box Trajectories as Compact State Representations for World Action Models

What is the minimal future a robot needs to imagine?

World action models (WAMs) are policy models that explicitly predict both actions and future world states, and have been shown to outperform action-only models across robotic manipulation tasks. However, representing the future with full-frame video prediction is computationally expensive during both training and inference. This raises a central question: what is the minimal state representation that preserves the planning and supervision benefits of future prediction? We propose BoxWAM, which replaces dense future frames with oriented bounding-box trajectories of task-relevant objects. By explicitly imagining object-centric future motion, our method provides a lightweight prediction target that captures the key physical changes induced by robot actions. Box-trajectory prediction substantially reduces computational cost while improving action-prediction performance, suggesting that object-centric compact states are an efficient alternative to video-based world action modeling for manipulation.

Each object's future is the image-plane location of the 8 corners of its minimum-volume oriented 3D bounding box: a 16-dimensional vector per object that keeps coarse shape and motion while discarding background texture, lighting, and shadows.

Inference pipeline

Predicted boxes across all 50 RoboTwin2.0 tasks

One successful clean-setting rollout per task from the multi-task π0.5 + BoxWAM model (50 tasks × 50 demonstrations), with the predicted box trajectory drawn on every frame. Played at 2× speed.

BoxWAM improves action prediction at near-zero extra cost

Across simulation and real robots, predicting compact box futures beats action-only imitation, while adding only a few milliseconds of inference latency compared with the tens of milliseconds that dense video or latent futures cost.

Real-world manipulation

Single-task simulation on RoboTwin2.0

Success rate (%) per task, 50 demonstrations per task. BoxWAM and Flow (action-only): 3 seeds × 100 trials. Other baselines from the RoboTwin2.0 leaderboard; π0 is pretrained. Scroll sideways for all nine tasks.

BoxWAM has the best randomized-setting average, and the best clean-setting average among RGB-only methods (only depth-based DP3 scores higher).

Multi-task generalist on RoboTwin2.0

2,500 demonstrations (50 tasks × 50) · joint-position control

Comparison with existing world action models

27,500 demonstrations · end-effector pose control

A separate, larger protocol from the multi-task table above: π0.5 + BoxWAM is fine-tuned on 27,500 RoboTwin2.0 demonstrations with end-effector pose actions for 35k steps (predict 50 actions, execute 36), and compared against existing WAMs. Latency on an RTX Pro 6000.

BoxWAM stays within a few milliseconds of π0.5 while closing most of the gap to much larger, slower WAMs. For reference, a differential-IK baseline (Mink + M2T2) given ground-truth task plans, poses, depth and masks reaches 26.78% Easy / 26.30% Hard.

Why boxes? Future-state representation ablation

Box futures add just +4.1 ms over action-only and outperform every other compact state in both settings. Dense VAE scores highest on clean but costs ~82 ms more, and needs 15k training steps vs. 6k for the compact variants. "≈ box": same latency as the box variant.

Attention ablation

BoxWAM works best when actions can attend to imagined future boxes, but boxes cannot attend to future actions. In other words, the box plan leads. All runtimes measured on an RTX Pro 6000 Blackwell GPU.