PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jie Chen, Ruofei Bai, Yuxin Cai, Yifeng Zhang, Chengyang He, Jun Li, Wei-Yun Yau, Guillaume Sartoretti

arXiv:2610.12194v1cs.RO

TL;DR

Dense native visual futures make WAM training costly, motivating compact future targets. MiniWAM learns PRISM targets from current–future transitions and jointly predicts them with actions; it outperforms native future-feature prediction across two feature families while reducing training cost.

  • Problem

    WAMs commonly predict high-dimensional native visual features, creating substantial training cost, and it remains open whether compact control-relevant future representations can directly replace those targets.

    Source: §1 · pp. 1–2

  • Method

    PRISM learns compact current–future transition targets through inverse-dynamics supervision and feature reconstruction, then MiniWAM freezes the encoder and jointly predicts those targets and actions from current observations.

    Source: §1 · p. 2

  • Results

    Across WAN2.1 VAE and DINOv3 features, MiniWAM outperforms native future-feature prediction with 65× fewer future tokens and reduces WAM training cost by up to 8×.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Strong world–action modeling does not require predicting the native visual future; compact learned representations can provide stronger and substantially more efficient targets for predictive co-training.

    Source: §1 · p. 3

    Simulation-benchmark performance does not establish real-world safety or reliability; physical-robot deployment requires task-specific safety validation, operational constraints, and human oversight.

    Source: ETHICS STATEMENT · p. 10

Abstract

from arXiv · show

World modeling has emerged as an effective co-training objective for robot policies, giving rise to World Action Models (WAMs) that jointly predict actions and future states. However, most WAMs predict future states in the native representation space of pretrained visual backbones, resulting in high-dimensional targets with substantial training cost. We introduce MiniWAM, which instead predicts compact future representations learned from privileged current-future transitions. To construct these targets, we propose Predictive Representations via Inverse Spatiotemporal Modeling (PRISM), which combines inverse-dynamics supervision with feature reconstruction to emphasize control-relevant transition information while preserving useful future-state information. With the learned PRISM encoder frozen, MiniWAM is trained to jointly predict the resulting targets and robot actions from current observations. With 65$\times$ fewer native future feature tokens, MiniWAM consistently outperforms native future-feature prediction with both DINOv3 and WAN2.1 VAE features, while achieving up to an 8$\times$ speedup in world-action training. At 0.25B parameters, MiniWAM is already competitive with substantially larger WAMs on LIBERO, LIBERO-Plus, and RoboTwin 2.0 simulation benchmarks. Representation analyses further show that PRISM contributes behavioral structure beyond feature reconstruction alone. These results demonstrate that effective world-action modeling does not require predicting native visual futures, and that compact predictive representations provide a strong and substantially more efficient target for policy learning. The project page is available at: https://j1dan.github.io/MiniWAM.

1 INTRODUCTION

World Action Models jointly predict actions and future visual states, but dense native visual targets make training costly and may preserve more detail than action learning needs. MiniWAM replaces those targets with compact PRISM representations and reports better performance with lower training cost.

1 INTRODUCTION

  • Dense, high-dimensional future visual representations impose substantial computational cost in WAMs, although preserving every native-feature detail may not be necessary for action learning.The paper frames the open question as whether compact, control-relevant future representations can replace dense visual prediction as WAM targets.Source: §1 · pp. 1–2
  • PRISM compresses current–future visual transitions, using inverse-dynamics supervision for action-informative transition content and reconstruction to preserve useful future-state information.MiniWAM freezes the learned encoder and jointly predicts its targets and robot actions from current observations.Source: §1 · p. 2
  • MiniWAM changes the world-model branch’s prediction target instead of merely reducing the cost of predicting an unchanged native visual future.The paper reports that adding inverse-dynamics supervision to reconstruction improves performance across bottleneck capacities, and concludes that compact learned targets can be stronger and more efficient.Source: §1 · pp. 2–3
  • 65× fewer future tokens accompany better performance than native future-feature prediction across VAE and DINOv3 features, with WAM training cost reduced by up to 8×.The comparisons use WAN2.1 VAE and DINOv3 features.Source: §1 · p. 2

2 RELATED WORK

Prior work combines future prediction with robot-policy training, but dense visual targets make WAMs costly. Existing efficiency and compact-latent approaches leave open whether a learned control-relevant transition representation can directly replace native future targets.

2 RELATED WORK

  • WAMs jointly model future states and robot actions to capture environment dynamics for control, but dense spatiotemporal visual targets substantially increase policy-training cost.The cited passage describes strong manipulation performance from this formulation.Source: §2 · p. 3
  • Efficiency work reduces generation cost, model or token budgets, or world–action interaction; compact-latent approaches instead use future features to shape policy-internal states or conditioning.These approaches suggest predictive information need not require explicit high-dimensional rollouts.Source: §2 · p. 3
  • PRISM differs by learning a control-relevant transition representation as the world-prediction target, jointly compressing spatial and temporal dimensions.The passage contrasts this with methods that retain native visual targets or use compact representations mainly as internal states or conditioning signals.Source: §2 · p. 3

3 MINI WORLD ACTION MODEL WITH PREDICTIVE REPRESENTATIONS

PRISM learns compact representations from privileged current–future visual features, and MiniWAM freezes that encoder to jointly predict the representations and actions from current observations. The method applies this target construction to both generative and discriminative visual features.

3.1 OVERVIEW

  • PRISM maps current and future visual features into a compact transition representation that replaces native future features as the WAM prediction target.The encoder uses privileged current–future features during target construction; future observations are unavailable at deployment.Source: §3.1 · p. 3

3.2 STAGE 1: LEARNING THE PREDICTIVE REPRESENTATION

  • Stage 1 trains the transition encoder with inverse-dynamics supervision and feature reconstruction, respectively emphasizing control-relevant dynamics and useful future-state information.The encoder processes current and future features together, and the decoder is used only during Stage 1.Source: §3.1 · p. 3, §3.2 · pp. 4–5

3.3 STAGE 2: WORLD–ACTION MODELING WITH PRISM TARGETS

  • Stage 2 freezes the PRISM encoder and jointly predicts its compact targets and action trajectories from current-observation tokens, with world and action predictions refined bidirectionally.At deployment, MiniWAM does not use the encoder or future observations and jointly denoises world and action variables.Source: §3.1 · p. 3, §3.3 · p. 5

3.4 VISUAL REPRESENTATION INSTANTIATIONS

  • The PRISM formulation is tested with frozen WAN2.1 VAE features and frozen DINOv3 patch features, while keeping the bottleneck formulation unchanged.This varies the underlying visual feature space to test compact predictive targets across generative and discriminative representations.Source: §3.4 · p. 5

4 EXPERIMENTS

The experiments compare PRISM targets with native future-feature prediction, assess MiniWAM against robot policies, and analyze PRISM’s efficiency and learned representation. PRISM improves control and training efficiency in the reported comparisons, while representation analyses find stronger transition alignment and behavioral structure.

4.1 EXPERIMENTAL SETUP

  • The evaluation covers controlled world-modeling comparisons, robustness under distribution shifts, multi-task bimanual manipulation, and training efficiency.LIBERO is used for controlled comparisons, LIBERO-PLUS for perturbed conditions, and ROBOTWIN 2.0 for bimanual manipulation and efficiency measurements.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 PRISM VERSUS DIRECT FUTURE-FEATURE PREDICTION

  • PRISM reaches 94.8% with VAE features and 97.2% with DINOv3 on LIBERO, improving over native future-feature co-training in both families.Native co-training raises the VAE-based policy from 62.3% to 92.8% and the DINOv3-based policy from 89.0% to 93.2%; on LIBERO-PLUS-2K, PRISM improves over native co-training by 6.3 and 13.2 percentage points, respectively.Source: §4.2 · p. 6, Table 1 · p. 6
  • At horizon 32, estimated total training speedup reaches 8.46× on ROBOTWIN 2.0, while Stage-2 training alone reaches 10.48×.PRISM reduces predictive co-training cost across both benchmarks and visual feature spaces, with larger gains under heavier workloads; Table 2 estimates total speedup including Stage 1.Source: §4.2 · pp. 6–7, Table 2 · p. 7

4.3 COMPARISON WITH STATE-OF-THE-ART ROBOT POLICIES

  • The 0.25B MiniWAM reaches 97.2% success on LIBERO, comparable to 6B-scale Fast-WAM and DreamWAM, and reaches 46.1% on ROBOTWIN 2.0, exceeding Fast-WAM and AHA-WAM.On LIBERO-PLUS, it achieves 73.6% success, outperforming Fast-WAM and approaching DreamWAM while using roughly an order of magnitude fewer world–action model parameters and no pretrained video-generation backbone.Source: §4.3 · pp. 7–8, Table 3 · p. 7

4.4 ABLATION STUDIES

  • Adding inverse-dynamics supervision to reconstruction improves success by roughly 4.6 percentage points across bottleneck capacities, while success decreases from 74.1% to 61.7% as the bottleneck shrinks.The combined PRISM objective performs best across capacities and consistently outperforms reconstruction alone.Source: §4.4 · p. 8, Table 4 · p. 8

4.5 SPATIOTEMPORAL BOTTLENECK ANALYSIS

  • On wrist observations, PRISM raises within-task CKA from 0.249 to 0.705 and deconfounded CKA from 0.071 to 0.393, strengthening cross-backbone alignment beyond raw transitions and a random encoder.The analysis uses unseen LIBERO-90 tasks; global CKA provides little separation on the main view.Source: §4.5 · pp. 8–9, Table 5 · p. 8
  • In leave-task-out retrieval, main-view Wan–DINO neighborhood agreement rises from 0.17 to 0.51, while wrist-view agreement rises from 0.11 to 0.33 and action RMSE falls from 0.37 to 0.27.A separate DINOv3-only ablation reports action RMSE reductions from 0.414 to 0.356 on the main view and from 0.326 to 0.266 on the wrist view, supporting behavioral structure beyond reconstruction alone.Source: §4.5 · p. 9, Table 6 · p. 9

5 DISCUSSION

MiniWAM replaces dense native-future prediction with compact PRISM targets learned through inverse dynamics and reconstruction. The discussion identifies transfer and encoder adaptation as future work and cautions that simulation results do not establish real-world safety or reliability.

5 DISCUSSION

  • MiniWAM uses compact PRISM targets learned from current–future transitions through inverse dynamics and reconstruction, instead of dense native-future prediction.The paper reports improved policy performance with reduced training cost across VAE and DINOv3 features.Source: §5 · p. 9
  • PRISM is learned separately for each downstream setting; shared pretraining on broader robot datasets is proposed for studying transfer across tasks, environments, and embodiments.The two-stage method also freezes PRISM before world–action training; joint or partial Stage-2 adaptation is proposed as future work.Source: §5 · pp. 9–10
  • The authors report that generative AI assisted with experimental and analysis protocols, code, interpretation, and manuscript refinement, while they retain responsibility for the work.The paper describes its methods and protocols and says implementation, configurations, checkpoints, and evaluation scripts will be released upon acceptance.Source: AI USE STATEMENT · p. 10, REPRODUCIBILITY STATEMENT · p. 10
  • Benchmark performance on existing datasets and simulation does not establish safety or reliability in real-world environments.Physical-robot deployment requires task-specific safety validation, operational constraints, and human oversight.Source: ETHICS STATEMENT · p. 10

A MINIWAM AND PRISM ARCHITECTURE

The architecture compresses current and future visual features into a spatiotemporal PRISM bottleneck, then uses world and action experts to predict compact targets and actions. The appendix specifies the feature pipelines, compression counts, decoder, and attention arrangement.

A MINIWAM AND PRISM ARCHITECTURE

  • The frozen visual inputs come from Wan2.1 VAE or DINOv3; DINO features are temporally compressed to four future maps before bottleneck processing.The DINO pipeline uses final-layer-normalized patch tokens and excludes class and register tokens.Source: A MINIWAM AND PRISM ARCHITECTURE · pp. 12–13
  • The default PRISM grid has 3 × 4 × 4 locations and 64 channels, yielding 48 world tokens per camera.Its dimensions denote temporal positions, spatial rows and columns, and channels per camera.Source: A MINIWAM AND PRISM ARCHITECTURE · p. 13
  • A 3D convolutional transition encoder compresses current and future features into a compact spatiotemporal grid, which is flattened into world tokens.The VAE and DINO bottlenecks both produce a (3, 4, 4) grid; a 1 × 1 × 1 head sets the feature dimension to 64.Source: A MINIWAM AND PRISM ARCHITECTURE · p. 13
  • PRISM replaces 3,136 native future feature-grid locations with 48 latent locations, about 65.3× fewer; future/latent transformer tokens fall from 784 to 48, about 16.3×.The token counts exclude current-observation conditioning, action tokens, and text context, so they do not imply the same reduction in total sequence length or wall-clock training cost.Source: A MINIWAM AND PRISM ARCHITECTURE · p. 13
  • The feature decoder reconstructs future features from the bottleneck and is discarded for policy inference.It reconstructs features rather than RGB; the Stage-2 target is the frozen encoder’s posterior mean.Source: A MINIWAM AND PRISM ARCHITECTURE · p. 13
  • The 0.25B Mixture-of-Transformers uses separate world and action experts with concatenated queries, keys, and values for mixed attention.Rows are queries and columns are keys in the attention masks; within each block attention is unrestricted.Source: A MINIWAM AND PRISM ARCHITECTURE · p. 13

B TRAINING AND INFERENCE DETAILS

The appendix records shared run settings and the flow-training schedule, along with benchmark-specific prediction and execution horizons. The PRISM encoder supplies Stage-2 targets, while the Stage-2 decoder is inactive.

B TRAINING AND INFERENCE DETAILS

  • LIBERO runs use 15- or 50-epoch schedules (16,260 or 54,200 optimizer updates), while RoboTwin runs use 6 or 25 epochs.The reported epoch counts are run settings; saved checkpoint configurations are the final provenance for each run.Source: B TRAINING AND INFERENCE DETAILS · p. 14
  • Both branches use independent Gaussian noise; Stage 1 activates action flow loss with clean world tokens, while Stage 2 jointly denoises actions and world tokens.Stage 2 uses the frozen Stage-1 encoder’s posterior mean as its target and initializes a new policy from the configured seed.Source: B TRAINING AND INFERENCE DETAILS · p. 14
  • PRISM runs use reconstruction weight λrec = 1 for both Wan-VAE and DINOv3, while the Stage-2 decoder is inactive with λrec = 0.The trained encoder alone supplies Stage-2 targets.Source: B TRAINING AND INFERENCE DETAILS · p. 14
  • LIBERO predicts 16-action chunks every 10 environment steps; RoboTwin Table 3 policies execute a complete 24-action chunk before replanning.Separate efficiency workloads use visual horizons of 16 and 32 with a fixed 16-step action horizon.Source: B TRAINING AND INFERENCE DETAILS · p. 14
  • The passage reports NumPy 1.26.4 and Einops 0.4.1, and states that LIBERO simulator dependencies include MuJoCo 3.3.2 and Robosuite 1.4.0.Source: B TRAINING AND INFERENCE DETAILS · p. 15

C DATASETS AND EVALUATION PROTOCOLS

The evaluation protocols specify the standard LIBERO trial setup, the perturbed LIBERO-PLUS population and text handling, and RoboTwin’s training split and leaderboard evaluation. The three benchmarks use distinct episode, perturbation, and execution procedures.

C.1 STANDARD LIBERO

  • LIBERO evaluation covers four suites and 2,000 episodes, with 50 episodes per task and the ith saved benchmark initial state used for trial i.Policies predict normalized 16-step action chunks, execute the first ten actions, then replan; success follows the environment’s task-success predicate.Source: §C.1 · p. 15

C.2 LIBERO-PLUS

  • LIBERO-PLUS reports the full 10,030-task perturbed population and a fixed 2,000-task ablation subset; clean model-facing instructions are encoded for evaluation.The subset covers all nonempty suite × perturbation category × difficulty strata; for non-language perturbations, clean instructions exclude filename-derived suffixes.Source: §C.2 · p. 15

C.3 ROBOTWIN 2.0

  • RoboTwin training uses demo-clean data from 50 tasks with 50 demonstrations each; leaderboard evaluation uses 100 valid episodes per task under demo-clean and demo-randomized conditions.The Table 3 policies execute all 24 predicted actions before replanning; efficiency profiling uses separate visual horizons of 16 and 32.Source: §C.3 · p. 15

C.4 PROVENANCE OF EXTERNALLY REPORTED BASELINES

Table 8 identifies sources for externally reported results marked in Table 3, whose values are success rates; unmarked entries are evaluations of released checkpoints. The section clarifies which baseline values are matched to the evaluated setting.

C.4 PROVENANCE OF EXTERNALLY REPORTED BASELINES

  • Table 8 supplies per-cell sources for Table 3 entries marked with †; unmarked entries are evaluations of released checkpoints.A dash indicates that the value is not used in the main comparison.Source: §C.4 · p. 16
  • Table 3 reports success rates (%), and the external ROBOTWIN values use the clean2clean/clean2random leaderboard track with clean-only training.This differs from results reported for some methods trained on clean and randomized demonstrations.Source: Table 8 · p. 16, §C.4 · p. 16
  • For WorldVLA, the matched LIBERO-Plus checkpoint scores 79.1 on standard LIBERO and 25.0 on LIBERO-PLUS, rather than the 512 × 512 model’s 81.8 LIBERO score.The section also uses UniVLA’s 42.9 from its archival CVPR 2026 LIBERO-Plus paper; AHA-WAM’s clean-trained leaderboard entry is 33.8, versus 92.8 after clean-plus-randomized training.Source: §C.4 · p. 16

D CONTROLLED BASELINES AND ABLATIONS

This section defines matched native-feature, PRISM, and no-co-training arms, then describes objective ablations and efficiency measurements. Pipeline measurements show faster Stage-2 updates with PRISM while retaining nearly all core throughput.

D CONTROLLED BASELINES AND ABLATIONS

  • Native and PRISM co-training use the same action objective, while native arms predict all future feature blocks or patch grids and PRISM predicts the frozen encoder’s posterior-mean grid.The no-co-training arm sets the direct world-flow coefficient to zero but retains sampled world tokens and joint attention, allowing action loss to update world-expert parameters.Source: D CONTROLLED BASELINES AND ABLATIONS · p. 16
  • The Stage-1 objective ablation compares reconstruction-only, inverse-dynamics-only, and combined supervision using frozen encoders in the same Stage-2 architecture and schedule.Capacity variants change only bottleneck grid or channel width, with matched data and optimizer budgets.Source: D CONTROLLED BASELINES AND ABLATIONS · p. 16
  • Efficiency measurements distinguish a training-core boundary from a pipeline boundary and use matched horizons, camera counts, and hardware for speedup comparisons.Throughput uses one RTX PRO 6000 Blackwell Max-Q GPU at global batch 256; Stage-1 cost is included in the additional accounting.Source: D CONTROLLED BASELINES AND ABLATIONS · pp. 16–17
  • Stage-2 speedups are 1.30 on LIBERO and 1.24 on ROBOTWIN.All speedups are computed before rounding.Source: D CONTROLLED BASELINES AND ABLATIONS · p. 17
  • 3.95×–9.60× pipeline Stage-2 speedups accompany retention of 96.9–99.2% of PRISM core throughput.Pipeline timing includes data loading, collation, host-to-device transfer, forward/backward passes, and optimizer updates; offline extraction is excluded.Source: Table 9 · p. 17, D CONTROLLED BASELINES AND ABLATIONS · p. 17

E REPRESENTATION-ANALYSIS PROTOCOLS

The representation analyses use a paired LIBERO-90 cohort and compare native and PRISM readouts across Wan and DINO backbones. They specify CKA conditioning and retrieval protocols, alongside limits on interpretation and uncertainty.

E.1 COHORT AND REPRESENTATION READOUTS

  • 4,500 transitions from 90 task-by-progress cells form the cohort, with a paired 1,800-transition subset for nested CKA analyses.The task IDs are disjoint from LIBERO-40 encoder training data, and main and wrist cameras are analyzed separately.Source: §E.1 · p. 17
  • Native Wan readouts contain 5 × 16 values and native DINO readouts 17 × 768, while each PRISM posterior-mean readout retains 64 × 3 × 4 × 4 values.The physical input window is matched, but native temporal resolution and feature dimension differ; pooled PRISM is a sensitivity analysis.Source: §E.1 · p. 17

E.2 CROSS-BACKBONE CKA

  • Cross-backbone CKA compares native transition readouts with trained full-grid posterior means using matched Wan and DINO inverse-dynamics-plus-reconstruction checkpoints.The comparison uses reconstruction weight λrec = 1 and does not mix in inverse-dynamics-only posterior banks.Source: §E.1 · p. 18, §E.2 · p. 18
  • The debiased linear CKA estimator uses zero-diagonal Gram matrices and a U-statistic, reducing high-dimensional estimator bias when representation dimensions exceed paired samples.Feature dimensions are not standardized because anisotropic rescaling changes geometry.Source: §E.2 · p. 18
  • CKA is evaluated globally, within task after task-wise mean removal, and after held-out residualization against task, progress, and current-only features.The deconfounding removes only structure linearly predictable by the specified nuisance model, not necessarily all current-scene information.Source: §E.2 · p. 18
  • Only the primary global before/after CKA difference has a hierarchical bootstrap interval; the other comparisons are paired point estimates and do not establish causal mediation or coordinate-wise equivalence.The analysis also compares trained encoders against random transition-encoder controls, but only one trained checkpoint pair and one random initialization are examined.Source: §E.2 · pp. 18–19

E.3 CROSS-BACKBONE NEIGHBORHOOD RETRIEVAL

  • Neighborhood retrieval uses 882 queries and 3,618 gallery transitions without shared source demonstrations, excludes same-task gallery items, and retrieves top-five neighbors across four readouts.Jaccard@5 measures overlap between Wan- and DINO-retrieved IDs; full-grid PRISM is primary and pooled PRISM is a sensitivity analysis.Source: §E.3 · p. 19
  • Behavioral retrieval metrics are action RMSE, end-effector displacement error in metres, and gripper-command agreement, averaged symmetrically over Wan- and DINO-retrieved neighbors.The DINO reconstruction-only versus inverse-dynamics-plus-reconstruction comparison uses identical transition IDs and retrieval rules.Source: §E.3 · p. 19

E.4 DINOV3 OBJECTIVE ABLATION

The ablation isolates inverse-dynamics supervision in matched DINOv3 bottlenecks, using the same leave-task-out retrieval protocol to compare neighborhood behavior.

E.4 DINOV3 OBJECTIVE ABLATION

  • The ablation compares DINOv3 bottlenecks with matched latent width 64 and posterior grid 3×4×4.Both objectives reconstruct native features; full PRISM additionally uses inverse dynamics with reconstruction weight 1.Source: §E.4 · p. 19
  • Both arms use the same 882 queries, 3,618 gallery transitions, and leave-task-out cosine top-5 retrieval protocol.The comparison uses identical transition IDs and retrieval rules.Source: §E.4 · p. 19
  • The DINO-only ablation metrics differ from Table 6, which compares raw transitions with PRISM and averages Wan and DINO retrieval metrics.Source: §E.4 · p. 19

Related papers

Loading 2610.12194v1…