Rethinking Representations for World-Action Modeling

What it claims

On RoboTwin 2.0, frozen DINOv3 features beat Video-VAE latents as the prediction space for a world-action model, even though their pixel-reconstruction SSIM is lower. This holds without generative video pre-training. Preparing those features for diffusion ("Feature Calibration") adds 12.20 points of average success. ReWAM builds on both findings. It is a 3.58B-parameter Mixture-of-Transformers model in which only the action loss trains the representation bottleneck.

On RoboDojo with embodied pre-training, ReWAM is competitive but not the top method:

The summary does not say this.

How it works

The architecture follows Fast-WAM. It has two branches, a world DiT and an action DiT, each 30 layers deep, with hidden widths of 2048 and 1024.

• Attention: future-state tokens and action tokens both attend to current-state features, but not to each other. The world branch uses causal temporal attention.
• Training: both branches use linear flow matching.
• Inference: the world branch encodes the current state once to condition the action branch. It never generates future states, so deployment carries no video rollout.

The encoder has three stages:

1. Feature Calibration (follows RAE). - Takes frozen DINOv3 ViT-L/16 features from blocks 12, 14, …, 24. - Averages them and adds the spatial mean of block 24 to every patch. - Uses a shifted noise schedule with α = √(D/4096), where D is the size of the prediction target, so larger targets get noisier inputs.
2. Temporal Representation Bottleneck (TRB). - Each latent state comes from a pair of frames. - Non-overlapping 2×2×2 spatiotemporal blocks are flattened per camera and concatenated across views. - Self-attention plus a linear projection reduce them to 128-dim tokens. - Outputs are normalised with running per-channel EMA statistics, frozen at inference.
3. Action-Grounded Representation Shaping (AGRS). - The world loss is stop-gradiented at both the current and future states. - The action loss flows through the world branch's current-state features into the TRB. - The world branch therefore learns from both losses, but the TRB is shaped only by the action loss.

flowchart LR
  O[Multi-view frames] --> D[Frozen DINOv3 + calibration] --> T[TRB]
  T --> Z0[current state z0]
  T --> ZF[future states]
  Z0 --> W[World DiT]
  Z0 --> A[Action DiT]
  ZF -. stop-grad .-> W
  W -->|world loss, stops at z0| X[ ]
  A -->|action loss| Z0
  Z0 -->|only action grads| T

Evidence

RoboTwin 2.0

• Setup: 50 tasks, trained on 2,500 clean plus 25,000 random-setting trajectories for 5 epochs, evaluated on 100 episodes per task per setting.
• ReWAM (no embodied pre-training) scores 93.6 on clean and 93.6 on random.
• The best compared method without pre-training is AHA-WAM at 93.4 / 92.2.
• With pre-training, LingBot-VA averages 92.2 and Fast-WAM 91.8.
• The margin is therefore +0.2 points on clean and +1.4 on random.
• Representation ablations used a one-epoch budget, not five.

RoboDojo

• Setup: 35 tasks, 3,500 trajectories, 50k steps.
• Without pre-training, ReWAM scores 9.42 / 5.82%, ahead of the next-best X-WAM at 7.69 / 3.83%.
• With the same 600-hour pre-training, it beats Fast-WAM (10.95 / 6.80%).
• Per category, ReWAM leads on Generalization (14.65 / 10.50%) and Long-Horizon success rate (16.25%).
• It trails on Memory (9.05 vs HyVLA-0.5's 13.37) and Precision (12.32 vs X-VLA's 18.32).
• It is near zero on Open (0.25 / 0.25%).

Embodied pre-training: 75k steps on 600 hours of self-collected real-robot demonstrations. The environments and cameras differ from the downstream real-robot setup.

How to use it

• Code: https://github.com/hustvl/ReWAM. The page states no licence and does not say whether weights or the 600-hour dataset are released.
• Inputs:
• head camera at 256×320 and wrist cameras at 128×160;
• 14-dim joint-and-gripper actions and proprioception, z-score normalised;
• UMT5 instruction embeddings.
• Training: 32 H20 GPUs, AdamW, batch size 512, learning rate 1e-4 cosine-decayed to 1e-6.
• Real-robot setup: a bimanual Piper with 10k steps of task-specific post-training.

Limits

• The fetched text is truncated. It omits:
• the real-robot results (three tasks: Collect Objects, Fold Clothes, Unpack Lunchbox; 20 trials each);
• the ablation table isolating TRB and AGRS;
• the numbers behind the DINO-vs-VAE comparison.
• RoboTwin gains over prior work are small (0.2–1.4 points).
• Absolute RoboDojo success rates stay below 9% for every method.
• The pre-training data is self-collected, and the page does not say whether it is available.