Source-linked AI summary
Box-WAM: Box Trajectories as Compact State Representations for World Action Models
Qinxi Yu, Wenzhou Ding, Gloria X. Wang, Albert J. Zhai, Shenlong Wang
TL;DR
WAMs predict actions and future states, but dense future frames are costly and compact state choices remain an open question. BoxWAM predicts task-relevant box trajectories alongside actions; across simulation and real-world experiments, it improves over action-only policies with minimal inference overhead.
Problem
Dense future representations require predicting irrelevant visual details, while latent states can be difficult to learn reliably and are not interpretable.
Source: §1 · p. 2Method
BoxWAM jointly predicts actions and trajectories of task-relevant objects’ oriented 3D boxes, represented by their eight image-plane corners.
Source: §1 · p. 2, §3.1 · p. 4Results
Across simulation and real-world experiments BoxWAM improves over action-only policies; its real-world average across success and progress metrics is 60.85%, versus 51.1% for the action-only model and 23.7% for Diffusion Policy.
Source: §4.2 · p. 7, §5 · p. 9Takeaways & Limitations
Compact object-centric futures provide an efficient and interpretable representation for manipulation, with minimal inference overhead reported across the experiments.
Source: §5 · p. 9Reliable object tracking is needed for ground-truth supervision, and performance may degrade under severe occlusion, highly deformable objects, or tasks where boxes are insufficiently informative.
Source: §6 · p. 9
Abstract
from arXiv · showhide
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
1 Introduction
The introduction frames future-state representation as an open WAM design question: dense futures are costly, while BoxWAM uses compact object-centric box trajectories to guide action prediction.
1 Introduction
- Dense future representations can model irrelevant appearance details, while latent representations may be difficult to learn reliably and are not interpretable.The introduction asks what future state a robot should imagine to improve action prediction.Source: §1 · p. 2
- BoxWAM represents each manipulated object with an oriented 3D box’s eight image-plane corners and jointly predicts box and action tokens.Structured attention lets actions use imagined object motion without allowing box predictions to depend on future actions.Source: §1 · p. 2
- The paper evaluates task-specific imitation and multi-task generalist policies in simulation, plus three real-world tasks, against action-only policies and existing WAMs.It also ablates state representation and state-action coupling.Source: §1 · p. 2
- Across simulation and real-world evaluations, BoxWAM improves action prediction while reducing future-modeling cost.The evaluation includes RoboTwin2.0 bimanual tasks and three real-world single-arm tasks.Source: §1 · p. 2
2 Related Work
Prior work studies world models, world action models, and structured representations, but the paper positions box trajectories as future states inside a WAM for manipulation.
2 Related Work
- World action models jointly predict actions and future states, using the future as an intermediate structure for action generation rather than candidate-action evaluation.This differs from action-conditioned world models that predict future states before a planner selects actions.Source: §2 · p. 2
- BoxWAM uses explicit object-centric box trajectories and imagined future boxes to improve action prediction.Its attention design lets actions condition on imagined object motion while keeping future-state predictions independent of future actions.Source: §2 · p. 3, Figure 2 · p. 3
- Structured representations such as keypoints and object boxes have often served modular pipelines or task conditioning, rather than learned future states inside WAMs.The paper studies box trajectories in that WAM setting.Source: §2 · p. 3
3 Method
BoxWAM replaces dense future-state prediction with object-centric box trajectories and jointly predicts those trajectories and robot actions. Its state-action design lets imagined object motion guide action generation while keeping future-state predictions independent of future actions.
3.1 Problem Formulation and Representation
- BoxWAM jointly predicts future state-action trajectories conditioned on the current image, proprioceptive state, and task condition.The task condition is empty for single-task policies and a language instruction for multi-task policies.Source: §3.1 · p. 3
- Each future state is a 16-dimensional vector of image-plane coordinates for the eight projected corners of an object's minimum-volume oriented 3D bounding box.This retains coarse object shape and motion while discarding pixel-level details such as texture, lighting, and shadows.Source: §3.1 · pp. 3–4
- Training selects the currently manipulated object, or the next object contacted in the demonstration when none is active; inference predicts its box trajectory without future contact order or ground-truth boxes.Real-world box supervision is reconstructed from RGB-D recordings, known meshes, and camera calibration, with rigid pose propagation during grasping.Source: §3.1 · p. 4
- Actions are trained as normalized relative joint-position commands and denormalized at inference to recover absolute target joint positions.The resulting target joint positions are sent to the controller.Source: §3.1 · p. 4
3.2 BoxWAM
- A transformer-based rectified-flow policy jointly generates future box and action tokens, with actions attending to future boxes but boxes barred from attending to future actions.The attention structure makes predicted box trajectories a task-centric object-motion plan for action generation.Source: §3.2 · p. 4
- The training objective jointly models future box trajectories and action chunks with equally weighted state and action losses.The model learns a velocity field from noise to the joint box-action target.Source: §3.2 · pp. 4–5
- At inference, BoxWAM initializes box and action tokens with Gaussian noise and updates them together using Euler integration of the learned rectified-flow equation.The horizon is 15 for task-specific simulation and 50 for multi-task and real-world settings.Source: §3.2 · p. 5
- For π0.5 integration, BoxWAM adds 50 future-box tokens and two linear projection layers, using the same state-action attention mask.Both policies are fully finetuned from the same checkpoint on 2,500 demonstrations for 30,000 steps; BoxWAM also receives box-trajectory supervision.Source: §3.2 · p. 5
4 Experiments
Experiments test BoxWAM in task-specific and multi-task RoboTwin2.0 settings, three real-world tasks, and ablations of future-state representations, token interactions, and runtime. Reported results favor box trajectories across randomized simulation, real-world execution, and several efficiency comparisons.
4.1 Experimental Setup
- The setup tests task-specific and 50-task RoboTwin2.0 learning, plus coffee making, plant watering, and toy placement on a real xArm7.Simulation includes clean and randomized settings; real-world evaluation uses 50 demonstrations and 50 randomized test scenarios per task.Source: §4.1 · p. 6
4.2 Results
- BoxWAM had the highest average success rate in randomized task-specific evaluation and outperformed all clean-setting baselines except depth-based DP3 on average.Table 1 covers RoboTwin2.0 single-task success rates.Source: Table 1 · p. 7, §4.2 · p. 7
- 52.6% to 61.5% clean success and 26.9% to 34.0% randomized success followed adding box-trajectory prediction to π0.5 in multi-task training.Both models used 2,500 demonstrations and 30,000 training steps; improvements were not uniform across tasks.Source: Table 2 · p. 7, §4.2 · p. 7
- In real-world evaluation under the same execution protocol and randomized object pose, BoxWAM averaged 60.85% across success and progress metrics, versus 51.1% for the action-only model and 23.7% for Diffusion Policy.Source: §4.2 · p. 7, Table 3 · p. 8
4.3 Ablations and Runtime Analysis
- Box trajectories achieved the highest randomized success rate and outperformed other compact state targets in both settings, while dense visual features did not consistently improve success.The ablation compares future-state targets and reports success and available latency measurements.Source: Table 4 · p. 8, §4.3 · p. 8
- BoxWAM performed best when actions could attend to future-state tokens but future-state tokens could not attend to future actions.Separate prediction was also competitive, suggesting a benefit from auxiliary supervision.Source: Table 5 · p. 8, §4.3 · p. 8
- Runtime comparisons show an efficiency advantage for compact box trajectories over full video-style future prediction, with lightweight overhead relative to action-only inference.The reported comparisons use an RTX Pro 6000 Blackwell GPU.Source: §4.3 · pp. 8–9
- 87.94% Easy and 87.62% Hard success at 64.9 ms with 3.35B parameters were reported for π0.5+BoxWAM; the IK baseline achieved 26.78% Easy and 26.30% Hard.The comparison used a much smaller training budget for BoxWAM; the IK baseline used ground-truth task plans, poses, depth, and masks without task-specific tuning.Source: Figure 6 · p. 9, §4.3 · p. 9
5 Conclusion
BoxWAM uses compact, object-centric box trajectories as future states for action generation. Across simulation and real-world experiments, it improves over action-only policies with minimal inference overhead.
- Across simulation and real-world experiments, BoxWAM improves over action-only policies while adding minimal inference overhead.Source: §5 · p. 9
- BoxWAM uses compact, object-centric box trajectories as future states for action generation.Source: §5 · p. 9
- Predicting task-relevant object motion instead of dense visual futures provides an efficient and interpretable state representation for manipulation.Source: §5 · p. 9
- Ablations show box trajectories are most effective as intermediate object-motion plans that guide action prediction.Source: §5 · p. 9
6 Limitations
BoxWAM depends on reliable object tracking to provide ground-truth supervision. Its performance may degrade in settings where tracking is difficult or boxes do not adequately represent the objects.
- Reliable object tracking is required to construct ground-truth supervision for BoxWAM.Source: §6 · p. 9
- Severe occlusion and highly deformable objects are settings where the approach may degrade.Source: §6 · p. 9
- The approach may also degrade when bounding boxes are insufficiently informative for a task.Source: §6 · p. 9
A Real-World OBB Ground-Truth Construction
Real-world BoxWAM supervision uses image-plane oriented bounding-box trajectories reconstructed from RGB-D observations, calibration, meshes, and robot proprioception. The pipeline estimates pose before grasp, propagates it through grasp using end-effector motion, then projects box corners into the image.
- Real-world supervision reconstructs oriented bounding-box trajectories from RGB-D observations, camera calibration, object meshes, and robot proprioception.Source: A Real-World OBB Ground-Truth… · p. 15
- Before grasp, the pipeline uses a prompted object mask and FoundationPose with the object mesh to estimate its 6-DoF pose.Source: A Real-World OBB Ground-Truth… · p. 15
- At a switch frame shortly after grasp, it computes a fixed object-in-gripper transform and uses end-effector kinematics to propagate object pose while grasped.Source: A Real-World OBB Ground-Truth… · p. 15
- The annotation pipeline assumes known meshes, calibrated RGB-D, and an approximately fixed object-to-gripper transform; slip, release motion, or initial-pose errors can systematically affect targets.Source: A Real-World OBB Ground-Truth… · p. 15
- Given the fused pose, mesh extents, and camera intrinsics, the pipeline projects the eight 3D box corners into the camera image.Source: A Real-World OBB Ground-Truth… · p. 15
- The projected 16-dimensional pixel-coordinate vector is used as the per-frame box state for real-world training.Source: A Real-World OBB Ground-Truth… · p. 15
B Real-World Progress Metrics
Real-world evaluation covers Make Coffee, Water Plant, and Place Pokémon. It reports final task success and progress, with progress based on completion of ordered subgoals.
- Real-world policies are evaluated on Make Coffee, Water Plant, and Place Pokémon.Source: B Real-World Progress Metrics · p. 15
- The evaluation reports final task success rate and progress rate over ordered task-specific subgoals.Source: B Real-World Progress Metrics · p. 15
- Progress is the fraction of ordered subgoals completed per rollout, averaged across rollouts and expressed as a percentage.Source: B Real-World Progress Metrics · p. 15, Table 6 · p. 15
C Additional Multi-Task Simulation Results
This section reports multi-task simulation success rates as an average across RoboTwin2.0 tasks and as per-task results under clean and randomized settings.
C Additional Multi-Task Simulation Results
- Models trained on 50 RoboTwin2.0 tasks with 50 clean-setting demonstrations per task were evaluated for 100 trials under both clean and randomized settings.All evaluations used seed 0.Source: C Additional Multi-Task Simulation… · p. 15
- Table 7 summarizes average success rate (%) across RoboTwin2.0 tasks in the multi-task simulation setting.Source: Table 7 · p. 16
- Tables 8 and 9 report per-task success rates (%) for multi-task simulation under clean and randomized settings, respectively.Both tables use 2500 demonstrations and joint-based control.Source: Table 8 · p. 16, Table 9 · p. 17
Related papers
This paper cites
- OpenVLA: An Open-Source Vision-Language-Action Model2024
- $π_0$: A Vision-Language-Action Flow Model for General Robot Control2024
- Diffusion Policy: Visuomotor Policy Learning via Action Diffusion2023
- RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control2023
- Planning-oriented Autonomous Driving2022
- Recurrent World Models Facilitate Policy Evolution2018