PaperlayerView exact revision on arXiv

Source-linked AI summary

Trace-WAM: Learning to Act via 3D Motion Prediction for World Action Models

Xu Hu, Changjie Wu, Linzhe Shi, Hang Zhang, Chen Yang, Lingjun Zhang, Guowen Zhang, Mu Xu, Junran Peng, Qing Li, Zhaoxiang Zhang

arXiv:2610.36840v1cs.RO

TL;DR

World Action Models couple future visual prediction with action generation, but scene-wide RGB prediction also models appearance and background beyond the motion most relevant to manipulation. Trace-WAM adds future 3D motion prediction as auxiliary supervision while retaining a current-observation action path. It reports strong results across manipulation benchmarks, including gains over Fast-WAM in few-shot and distribution-shift settings.

  • Problem

    World Action Models combine future visual prediction with action generation, but future RGB must model the whole scene while manipulation depends primarily on how the robot and task-relevant objects evolve.

    Source: §1 · p. 2

  • Method

    Trace-WAM jointly trains video, future 3D motion, and action experts, using grid-anchored point trajectories with image-plane displacement and depth change as auxiliary supervision.

    Source: §1 · p. 2

  • Results

    Across RoboCasa-GR1, LIBERO, and LIBERO-Plus, Trace-WAM reports strong in-distribution and out-of-distribution results; on zero-shot LIBERO-Plus, it exceeds Fast-WAM by 18.3 percentage points.

    Source: §1 · p. 2, §4.2 · p. 6, §5 · p. 8

  • Takeaways & Limitations

    The results support future 3D motion prediction as an effective training signal for WAM policy learning, with action generation at deployment using current observations, language, and proprioception.

    Source: §1 · p. 2, §5 · p. 8

Abstract

from arXiv · show

World Action Models (WAMs) learn representations for robot control by coupling future visual prediction with action generation. Successful manipulation depends primarily on how the robot and task-relevant objects evolve, whereas future-video prediction must also reconstruct the appearance and background of the entire scene. We therefore propose Trace-WAM, which augments video-action learning with future 3D motion prediction as an auxiliary objective. This objective supervises temporal point correspondence together with image-plane displacement and depth change, directly describing the interaction dynamics and spatial changes involved in manipulation. To this end, we adopt a three-branch Mixture-of-Transformer (Tri-MoT) architecture to jointly learn video prediction, motion prediction, and action generation. The motion expert predicts future trajectories from current visual features, allowing motion supervision to shape the representations used for action generation. Because the action expert does not consume predicted trajectories, inference retains the current-observation policy path: a single visual forward pass followed by action denoising. Experiments on RoboCasa-GR1, LIBERO, and LIBERO-Plus demonstrate strong in-distribution performance and out-of-distribution generalization, with particularly clear gains over Fast-WAM under limited demonstrations and distribution shifts.

1. Introduction

Trace-WAM addresses the mismatch between scene-wide future-video prediction and manipulation-relevant motion by adding future 3D motion supervision to video-action learning. It jointly trains video, motion, and action experts, while the reported evaluations show strong performance across benchmarks and clearer gains in few-shot and distribution-shift settings.

1. Introduction

  • Future RGB prediction models scene-wide texture, illumination, and background, whereas manipulation depends primarily on how the robot and task-relevant objects change over time.The motivation is to retain rich visual context while adding a target that emphasizes interaction dynamics.Source: §1 · p. 2
  • Among moving tracks, 79.3% in LIBERO and 32.9% in Piper show depth changes of at least 5% of initial depth; retaining depth raises LIBERO-Plus success from 82.4% to 84.5%.The tracks move at least 2% of the image diagonal over 32 steps.Source: §1 · p. 2
  • Trace-WAM jointly learns future video prediction, future 3D motion prediction, and action generation, representing motion as grid-anchored trajectories with image-plane displacement and depth change.The trajectories describe future motion relative to each point's initial position.Source: §1 · p. 2
  • During training, video, motion, and action objectives jointly shape action-generation representations; at inference, action denoising uses cached visual features without future-video or motion prediction.The policy path processes current RGB observations once and also uses language and proprioception.Source: §1 · p. 2, Figure 3 · p. 3
  • Trace-WAM shows strong in-distribution and out-of-distribution performance, with particularly clear gains over Fast-WAM under limited demonstrations and distribution shifts.Evaluations cover RoboCasa-GR1, LIBERO, and LIBERO-Plus; real-robot results show higher success across three physical manipulation tasks.Source: §1 · p. 2

2. Related Work

Prior work connects video prediction and action generation in World Action Models, while auxiliary objectives extend supervision beyond RGB. Motion-based policy learning spans image-plane motion, 3D motion, and representation supervision; Trace-WAM directly predicts future 3D point trajectories alongside video and actions.

2.1. World Action Models

  • World Action Models integrate video priors into policy learning by conditioning actions on predictive video features or jointly modeling future observations and action sequences.Related work also adds semantic, depth, geometry, motion, affordance, interaction-configuration, and goal-region supervision; Trace-WAM adds future 3D motion prediction alongside RGB and action generation.Source: §2.1 · p. 3

2.2. Motion for Policy Learning

  • Motion representations connect action generation to robot and object movement, with depth adding displacement along the viewing direction relevant to reaching, grasping, and transport.Prior approaches use optical flow or 2D point-motion prediction, while others incorporate depth and 3D motion for spatial scene dynamics.Source: §2.2 · p. 3
  • Unlike tracker-feature alignment, Trace-WAM directly predicts numerical future 3D point-trajectory targets alongside RGB prediction and action generation; deployed actions use current visual context, language, and proprioception.The motion objective shapes video-action representations during training without requiring future-video or motion prediction at deployment.Source: §2.2 · p. 3

3. Method

Trace-WAM augments video-action policy learning with future 3D motion prediction, using motion supervision to shape action representations while keeping deployed action generation independent of predicted motion.

3. Method

  • Trace-WAM jointly learns future video prediction, future 3D motion prediction, and action generation, using motion as auxiliary supervision for video-action representations.The deployed policy generates actions from current observations.Source: §3 · p. 3

3.1. Problem Formulation

  • The policy executes an action chunk from current observations, while future RGB and motion labels are used only during training.The robot executes the first few actions before observing again and replanning.Source: §3.1 · p. 4

3.2. 3D Motion Representation

  • A fixed grid of main-camera anchors tracks image-plane displacement and depth change from each anchor’s initial position, with trajectories represented by compact B-spline coefficients.Trajectories are expressed in a unified reference-camera frame, and the reference configuration uses 10 control points per anchor.Source: §3.2 · p. 4

3.3. Model Architecture

  • Three Transformer experts process video, motion, and action tokens; structured attention lets motion and action use current video features without making the deployed policy depend on prediction branches.Current video tokens attend only to one another, while future-video, motion, and action tokens also attend to current video.Source: §3.3 · p. 4

3.4. Training Objective

  • Flow matching trains video and action velocity predictions, while motion supervision balances moving and static anchors and penalizes decoded residual motion at static anchors.The static-motion penalty encourages stationary points to remain stationary after decoding.Source: §3.4 · p. 5
  • Training first warms up with video and motion objectives, then adds the action objective for joint optimization; the VAE and text encoder remain frozen.The three experts and proprioceptive-token projection are trainable.Source: §3.4 · p. 5

4. Experiments

Experiments test Trace-WAM on RoboCasa-GR1, LIBERO, and zero-shot LIBERO-Plus, then examine motion-supervision design choices and physical-robot performance. Results show strong benchmark performance, gains under limited demonstrations and distribution shifts, and higher success across three real-world tasks.

4.1. Experimental Setup

  • Evaluations cover RoboCasa-GR1, LIBERO, and zero-shot transfer to LIBERO-Plus across seven perturbations, with no augmented LIBERO-Plus data used for fine-tuning.RoboCasa-GR1 includes full-data and 10% few-shot settings; training uses the Fast-WAM codebase and pretrained Wan2.2-5B video backbone.Source: §4.1 · pp. 5–6

4.2. Main Results

  • 98.7% average success across LIBERO’s four suites compared with 97.6% for Fast-WAM and 98.5% for Joint-WAM; zero-shot LIBERO-Plus success was 84.5%, 18.3 points above Fast-WAM and 5.8 above Joint-WAM.Trace-WAM improved over Fast-WAM across all seven LIBERO-Plus perturbation types, including Robot at 87.4% and Background at 76.1%.Source: Table 2 · p. 6, §4.2 · p. 6, Table 3 · p. 7
  • 58.6% average success on RoboCasa-GR1 unseen-appearance, unseen-combination, and unseen-object tests exceeded DIAL, Fast-WAM, and Joint-WAM by 12.4, 16.3, and 8.6 points, respectively.In RoboCasa-GR1, few-shot and full-data success rates were 63.6% and 80.3%, outperforming Fast-WAM and Joint-WAM in both settings.Source: §4.2 · p. 6

4.3. Ablation Study and Analysis

  • 84.5% average LIBERO-Plus success for Trace-WAM exceeded 80.4% without moving/static flow-matching balance and 83.6% without the static-motion penalty.The results identify moving/static supervision balance as important, with static-motion regularization providing a further improvement.Source: §4.3 · p. 6, Table 4 · p. 7
  • 84.5% average success with auxiliary motion supervision exceeded Motion Inject at 81.6% and bidirectional video–motion attention at 80.3%.The comparison favors auxiliary motion supervision without action conditioning on future motion latents or bidirectional video–motion coupling.Source: §4.3 · p. 6, Table 4 · p. 7
  • 2.1 points: adding depth improved average success over 2D motion, with gains of 5.8 and 6.0 points under Background and Noise perturbations; B-spline targets added 1.1 points over raw motion.The cited comparisons report 82.4% for 2D Motion and 83.4% for Raw Motion; the authors also report that B-spline targets lower motion-prediction errors across all three horizons.Source: Table 4 · p. 7, §4.3 · pp. 7–8
  • At 32 steps, B-spline Motion achieved 1.43-pixel image-plane ADE and 2.768-mm depth MAE, versus 1.87 pixels and 3.427 mm for Raw Motion.The comparison aggregates 140 rollout episodes at each horizon; lower error is better.Source: Table 5 · p. 7, §4.3 · pp. 7–8

4.4. Real-world Experiments

  • 85.0% average success across three physical tasks compared with Fast-WAM’s 63.3%; Place and Close Drawer improved from 55.0% to 95.0%.Wipe Blackboard improved from 55.0% to 75.0%, and Place Bread in Bowl from 80.0% to 85.0%.Source: §4.4 · p. 8

5. Conclusion

Trace-WAM uses future 3D motion prediction as auxiliary supervision for video–action representation learning, while the deployed policy generates actions from current observations. It achieves strong in-distribution and out-of-distribution performance and improves success across three real-world manipulation tasks.

5. Conclusion

  • Trace-WAM jointly optimizes video, future 3D motion, and action objectives during training, while deployment generates actions from current observations, language, and proprioception without future-video or motion prediction.The paper concludes that future 3D motion prediction is an effective training signal for robust action representations in WAMs.Source: §5 · p. 8
  • Across LIBERO, zero-shot LIBERO-Plus, and RoboCasa-GR1, Trace-WAM achieves strong in-distribution and out-of-distribution performance, with clear gains in few-shot and OOD settings.Real-world experiments on a Piper robot also demonstrate improved success across three manipulation tasks.Source: §5 · p. 8
  • Auxiliary future 3D motion prediction focuses video–action representation learning on scene dynamics most relevant to manipulation.The deployed policy does not run future-video or motion prediction.Source: §5 · p. 8

Related papers

Loading 2610.36840v1…