Box-WAM: Box Trajectories as Compact State Representations for World Action Models

Qinxi Yu; Wenzhou Ding; Gloria X. Wang; Albert J. Zhai; Shenlong Wang

cs.ROOct 13, 2026

Box-WAM: Box Trajectories as Compact State Representations for World Action Models

Authors: Qinxi Yu, Wenzhou Ding, Gloria X. Wang, Albert J. Zhai, Shenlong Wang

Organizations: University of Illinois Urbana-Champaign

Abstract

World action models (WAMs) are a category of policy models that ex- plicitly predict both actions and future world states. They have been shown to outperform action-only models across various robotic manipulation tasks. How- ever, representing the future with full-frame video prediction can be computation- ally expensive during both training and inference. This raises a central question: what compact state representation preserves the planning and supervision bene- fits of future prediction? In this paper, we propose BoxWAM, which replaces dense future frames with oriented bounding box trajectories of task-relevant ob- jects. By explicitly imagining object-centric future motion, our method provides a lightweight prediction target that captures the key physical changes induced by robot actions. We show that box trajectory prediction substantially reduces com- putational cost while improving action prediction performance, suggesting that object-centric compact states can serve as an efficient alternative to video-based world action modeling for robotic manipulation. Result videos can be found on https://box-wam.github.io

Figures & tables

Figure 2: Overview of BoxWAM. Observation images are encoded into visual context tokens, which condition a rectified-flow transformer that jointly denoises future state tokens and action tokens. An enforced causal attention allows action tokens to condition on imagined object motion while preventing future state tokens from depending on future action tokens. The decoded future state forms a compact box trajectory, and the decoded action tokens form the robot action chunk.

Panel of Figure 3 — Figure 3: Simulation rollout visualizations on RoboTwin2.0. We visualize representative rollouts from eight simulation tasks with predicted future box trajectories overlaid on the scene. BoxWAM predicts compact object-centric future motion across diverse manipulation behaviors.

Method	Obs.	Avg.	Adjust Bottle	Beat Hammer	Click Alarm	Grab Roller	Handover Mic	Lift Pot	Open Laptop	Pick Dual Bottles	Turn Switch
Clean setting
DP	1 RGB	55.4	97.0	42.0	61.0	98.0	53.0	39.0	49.0	24.0	36.0
ACT	3 RGB	60.4	97.0	56.0	32.0	94.0	85.0	88.0	56.0	31.0	5.0
DP3	1 RGBD	81.2	99.0	72.0	77.0	98.0	100.0	97.0	82.0	60.0	46.0
$\pi_{0}$	3 RGB	71.4	90.0	43.0	63.0	96.0	98.0	84.0	85.0	57.0	27.0
Flow (action-only)	1 RGB	73.2	99.7	63.0	90.3	99.7	58.3	69.3	81.0	47.7	50.0

Table 1: Single-task RoboTwin2.0 success rates (%). BoxWAM and Flow (action-only) are evaluated over three seeds with 100 trials per seed for each task and setting. Other baselines are from the RoboTwin2.0 leaderboard. Gray rows indicate pretrained models.

Panel of Figure 4 — Figure 4: Real-world rollout visualizations. We visualize representative rollouts from three real-world manipulation tasks with predicted future box trajectories overlaid on the scene-camera view. The first two columns show early and late coffee-making stages; the remaining columns show watering and toy placement. Each row shows a different frame, illustrating how BoxWAM predicts compact object-motion plans for action generation under real-world perception noise, actuation error, and object pose variation.

Method	Clean	Randomized	Latency
$\pi_{0.5}$	52.6	26.9	55.3
$\pi_{0.5}$ + BoxWAM	61.5	34.0	57.7

Table 2: Multi-task fine-tuning with 2,500 demonstrations and joint-position control. Both models are trained for 30,000 steps for a fair comparison. Models are evaluated under clean and randomized settings. We report success rate (%). The inference latency (ms) is reported based on default 10 denoising steps from

\pi_{0.5}

Method	Make Coffee		Water Plant		Place Pokemon		Average
	SR	Progress	SR	Progress	SR	Progress
Diffusion Policy	10.0	34.0	8.0	9.4	34.0	47.0	23.7
Flow (action-only)	24.0	48.0	48.0	54.7	60.0	72.0	51.1
BoxWAM	26.0	54.4	66.0	70.7	70.0	78.0	60.9

Table 3: Real-world manipulation results. We report both final success rate and progress rate because full task success can be sparse. Each method is evaluated under the same execution protocol and randomized object pose.

Table 7

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Task	# Subgoals	Ordered subgoals
Make Coffee	5	grasp mug; place mug; grasp K-Cup; place K-Cup; close lid
Water Plant	3	grasp pitcher; water plant; place pitcher
Place Pokemon	2	grasp toy; place toy

Appendix

Table 6: Real-world progress subgoals. Progress is measured by the number of completed ordered subgoals in each rollout.

Method	Cameras	Clean Avg.	Randomized Avg.
$\pi_{0.5}$	3 RGB	52.6	26.9
$\pi_{0.5}$ + BoxWAM	3 RGB	61.5	34.0

Appendix

Table 7: Multi-task simulation summary. We report average success rate (%) over RoboTwin2.0 tasks.

Task	$\pi_{0.5}$	$\pi_{0.5}$ + BoxWAM
adjust_bottle	68.0	94.0
beat_block_hammer	81.0	86.0
blocks_ranking_rgb	6.0	19.0
blocks_ranking_size	28.0	46.0
click_alarmclock	89.0	93.0
click_bell	78.0	73.0

Appendix

Table 8: Per-task multi-task simulation results under the clean setting trained with 2500 demos and joint based control. We report success rate (%).

Task	$\pi_{0.5}$	$\pi_{0.5}$ + BoxWAM
adjust_bottle	15.0	53.0
beat_block_hammer	20.0	9.0
blocks_ranking_rgb	6.0	9.0
blocks_ranking_size	10.0	19.0
click_alarmclock	48.0	62.0
click_bell	51.0	50.0

Appendix

Table 9: Per-task multi-task simulation results under the randomized setting trained with 2500 demos and joint based control. We report success rate (%).

World Action Models

World Action Models connect predictions of how scenes change under action to executable policies for robots, vehicles, and other embodied agents.

See all 218 papers

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding
Jun 6, 2026Ziang Li, Dongzhou Cheng, Yibin Wang +5Robot Policy Learning World Models for Robotics
Wuhan University · Shanghai Innovation Institute · Southeast University +2
MiniWAM: Learning Compact Future Targets for Efficient World-Action Modeling
Oct 8, 2026Jie Chen, Ruofei Bai, Yuxin Cai +5World Models for Robotics World Action Models
Department of Mechanical Engineering, National University of Singapore, Singapore · A*STAR Institute of Advanced Intelligence and Computing (A*STAR IAIC), Singapore · Nanyang Technological University, Singapore
TrajWAM: An Efficient Trajectory-Guided World Action Model for Dynamic Robotic Manipulation
Oct 19, 2026Haohan Geng, Sha Zhang, Weishuo Sun +7Dynamic Robotic Manipulation World Action Models
University of Science and Technology of China. · The Chinese University of Hong Kong. · University of Macau. +1

Jun 6, 2026J/K move · Enter open · S save

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding

Ziang Li, Dongzhou Cheng, Yibin Wang, Shiyue Wang+4

Wuhan University · Shanghai Innovation Institute · Southeast University+2

Robot Policy Learning World Models for Robotics Visuomotor Policy Learning Efficient Inference for World Action Models World Action Models

Open paper AskPDF arXiv

World Action Models (WAMs) extend robot policy learning by incorporating future prediction as an additional training objective, encouraging the policy to encode task-relevant temporal structure in its representations. Current WAMs often rely on large-scale generative architectures that incur high training costs and inference latency, making them difficult to deploy as efficient closed-loop policies. We propose Light-WAM, a lightweight World Action Model for efficient robot manipulation. Specifically, it is built with a compact video backbone and performs future-video supervision in a downsampled latent space, reducing the cost of video co-training while retaining its benefits for representation learning. For action prediction, Light-WAM introduces the StateFusionActionExpert, which reads adapted states from multiple backbone layers, fuses them through learned-query pooling, and directly predicts action chunks in a single forward pass. This design provides an efficient interface between video backbone representations and robot actions, avoiding the need for heavy generative action experts. Experiments demonstrate that Light-WAM maintains strong performance on LIBERO and achieves usable multi-task performance on RoboTwin 2.0, while using only 0.44B trainable parameters. It also achieves 72.03ms inference latency with 4.1GiB peak GPU memory and improved training throughput.

Abstract

World Action Models

Explore similar work

Light-WAM: Efficient World Action Models with State-Fusion Action Decoding

MiniWAM: Learning Compact Future Targets for Efficient World-Action Modeling

TrajWAM: An Efficient Trajectory-Guided World Action Model for Dynamic Robotic Manipulation