Skip to content

RepWAM: World Action Modeling with Representation Visual-Action Tokenizers

RepWAM (Fudan / Robbyant / HKUST) is a representation-centric world action model that argues WAM performance is bottlenecked by the latent space current WAMs inherit from reconstruction-oriented video generators (WAN2.2 VAE, etc.). It introduces RepViTok — a two-stage tokenizer that first trains a ViT video autoencoder with a Perception Encoder feature-alignment loss, then trains a latent action tokenizer as an inverse-dynamics + forward-dynamics pair whose “actions” are transport-plus-residual operators between visual latents. A causal DiT is then pretrained from scratch to jointly model future visual states and their latent actions via flow matching, and adapted to real robot data for control. Reports 89.3 / 88.4 on RoboTwin 2.0 Easy / Hard averages at 5B, plus 30% → 50% real-robot success on a Franka pick-fruit task from swapping WAN2.2 VAE for RepViTok in a fixed 1.3B WAM.

  • Reconstruction-oriented video tokenizers (WAN2.2 VAE) leave WAMs with a latent space dominated by appearance rather than task-relevant semantics, and this is diagnosable both as worse open-loop action score (OLS 13.68/11.21 vs 18.82/14.15 for RepViTok) and as worse closed-loop control (20% vs 30% PickFruit success at fixed 1.3B WAM) [Table 2].
  • Aligning the video-autoencoder latent to a frozen Perception Encoder via temporally-pooled MSE — while keeping standard reconstruction + perceptual + adversarial losses — gives a tokenizer that is competitive with WAN2.2 on ImageNet/UCF101 reconstruction (rFID 0.80/0.23 vs 0.50/0.20, rFVD 1.09/0.16 vs 4.28/0.68 on UCF101 256/512) and simultaneously beats it on downstream WAM gFVD by 9.5% / 13.2% on seen/unseen splits [Table 5, Table 2].
  • Latent actions should be learned inside the aligned semantic visual latent space, not directly from pixels. Freezing the visual tokenizer and training a latent-action tokenizer as an IDM+FDM pair — with the FDM factorized as a transport operator M (soft optical-flow-in-token-space) plus a residual δ — captures manipulation-centric transitions that map more cleanly to robot actions than LAPA-style pixel-derived codes [§3.2, Fig. 4].
  • Two-stage pretraining (visual-and-latent-action pretraining, then robot-action adaptation) beats both no-latent-action baseline and Joint-Pred (single head predicting video and action latents jointly): gFVD 48.23/58.83 vs 61.01/72.91 (w/o) vs 94.25/98.77 (Joint Pred), and PickFruit 50% vs 30% vs 20% [Table 4].
  • Semantic latent alignment reduces reliance on video CFG at inference — the RepViTok-based WAM peaks at video-CFG scale 1.0 (i.e. no CFG extrapolation) on averaged success, so the unconditional video branch can be dropped, saving latency and activation memory [§4.4, Fig. 5].
  • Scaling RepWAM from 1.3B to 5B lifts long-horizon push-drawer and fine-grained insert-tube tasks more than short-horizon pick-fruit (which is bottlenecked by grasping perception rather than model capacity), while both sizes are best-or-tied-best against π0.5 and LingBot-VA on all three real tasks [Fig. 2].
  • On RoboTwin 2.0’s 50-task Easy/Hard averages, RepWAM-5B (from-scratch, no video-generation pretraining) beats π0.5 and Motus and closes most of the gap to LingBot-VA (which uses WAN video-generation pretraining) — the remaining gap is attributed to the pretrained backbone rather than the recipe, supported by the ablation swapping WAN2.2 VAE → RepViTok inside a fixed 1.3B WAM (78.0/76.0 → 86.6/83.1 Easy/Hard average) [Table 1, Table 3].

Two-stage pretraining then adaptation. Stage 1 (visual tokenizer): a ViT autoencoder with temporal-causal + full-spatial attention encodes the initial frame (patched) and subsequent tubelets (t×h×w) into video latents; the decoder is symmetric with an unpatchify head. Reconstruction, perceptual, adversarial, and — critically — feature-alignment losses (temporally-pooled MSE of a linear projection of the latents against a frozen Perception Encoder) are combined with weights (1, 0.1, 1, 0.5). Stage 2 (latent action tokenizer, visual tokenizer frozen): for consecutive latent frames z_t, z_{t+1}, an IDM compresses the transition to a small latent action a_t (channel bottleneck to prevent content leakage). An FDM then reconstructs z_{t+1} as M(a_t) · z_t + δ(a_t), where M acts as a soft transport operator (optical-flow analog in the semantic token space) and δ captures residual state changes. Trained with a forward next-latent prediction loss plus a backward consistency loss on the reversed pair.

The WAM itself is a causal diffusion transformer trained from scratch (1.3B and 5B). Video and latent-action tokens are grouped into chunks over a temporal window and prefixed by language and initial-context tokens; block-causal masking within a chunk lets visual and action tokens share attention weights but use modality-specific FFNs. A conditional flow-matching objective is applied jointly to visual and latent-action components. Muon optimizer, PLM text encoder via cross-attention. Data: RepViTok on Panda-70M; WAM pretraining on AgiBot (~G tokens); embodiment adaptation on a mix of AgiBot, RoboMIND, RoboCOIN, InternA1 with continuous motor commands.

  • RoboTwin 2.0 (50-task benchmark, Easy / Hard averages): RepWAM-5B 89.3 / 88.4 vs π0.5 82.7/76.8, Motus 88.7/87.0, LingBot-VA (WAN-pretrained) 92.9/91.6; RepWAM-1.3B 86.6/83.1 [Table 1].
  • Ablation with WAN2.2 VAE ↔ RepViTok swap at fixed 1.3B: Easy 78.0 → 86.6, Hard 76.0 → 83.1 [Table 3].
  • Franka dual-arm real-world (3 tasks, 40+ rollouts per size): pick-fruit both sizes tied top; long-horizon push-drawer RepWAM-5B new best; fine-grained insert-tube gains largest at 5B [Fig. 2].
  • Visual-tokenizer ablation on AgiBot (Seen / Unseen splits): RepViTok gFVD 61.01/72.91 (vs 67.42/83.98 WAN2.2), OLS 18.82/14.15 (vs 13.68/11.21), PickFruit 30% (vs 20%); ViTok (no alignment) sits between them [Table 2].
  • Two-stage vs Joint-Pred vs no-latent-action on same visual tokenizer: two-stage best on every metric (gFVD 48.23/58.83, PSNR 22.86/19.93, OLS 19.87/16.98, PickFruit 50%) [Table 4].
  • Video CFG averaged success peaks at scale 1.0 for RepViTok-based WAM (i.e., no CFG extrapolation needed), unlike LingBot-VA which benefits from CFG > 1 [Fig. 5].
  • Reconstruction: rFID 0.80/0.23 (256/512 ImageNet), rFVD 1.09/0.16 (256/512 UCF101) — competitive with WAN2.2 VAE while delivering downstream gains [Table 5].

RepWAM is the arxiv paper form of the RepViTok tokenizer already cited as the “semantic visual-action tokenizer” inside LingBot-VA 2.0: Native Video-Action Pretraining for Generalizable Robot Control — same Robbyant / Ant Group lab, and this paper is the standalone tokenizer-and-WAM ablation study that grounds LingBot-VA 2.0’s headline “native from-scratch pretraining” claim. It joins Diffusion Transformers with Representation Autoencoders and VideoRAE: Taming Video Foundation Models for Generative Modeling via Representation Autoencoders in the representation-autoencoder line by extending the “freeze a foundation-model encoder, align the latent to it” recipe from image generation and video generation into world-action modeling for robotics — where the alignment target is not a general image-encoder but a Perception Encoder, and the latent space is asked to host both future prediction and action inference. It also complements μ₀: A Scalable 3D Interaction-Trace World Model (µ₀ — frozen trace WM + small action expert, 3D-interaction interface) and ImageWAM: Do World Action Models Really Need Video Generation, or Just Image Editing? (KV-cache from image editor into flow-matching action expert) with a third structural bet: instead of factoring the world model and the action expert into separate frozen-and-small components, share one aligned semantic latent between them and train the WAM jointly from scratch. The reduced reliance on video CFG at inference is a small but concrete efficiency win that follows from the semantic alignment.