Think Like a World Model, Act Like a VLA: Distilling World-Model Representations into Compact Robot Policies

AI summary

Overview

• Research area: Robot learning — vision-language-action (VLA) models, world models, and representation-level knowledge distillation for robot manipulation.
• Technical level: Advanced (assumes familiarity with VLA architectures, flow-matching action heads, and distillation losses, though the central idea is simple enough to follow without that background).
• Scope: One sentence: The paper shows that a compact 0.8B VLA policy can inherit the physically grounded internal representations of a frozen world model through a single cached feature-alignment loss, improving accuracy on simulation and real-robot manipulation at zero added inference cost.

What This Paper Is About

VLA models map observations directly to actions and have no training objective that forces them to model how the world responds to those actions, so their robustness is bounded mainly by how much data they have seen. World models do carry that missing objective and produce more physically grounded features, but generating a future costs seconds per decision (the paper cites DreamZero at 3s per inference on an H100), which rules them out of a real control loop. The goal of this work is to separate what a world model knows from what it does: transfer the grounded internal features to a small policy during training, then discard the generative machinery entirely so the deployed policy is identical to the undistilled baseline.

Key Contributions

1. A distillation recipe that transfers world-model grounding at zero inference cost. One cosine alignment term against a cached teacher feature, no teacher forward pass during training, and no change to the deployed computation graph.
2. A 0.8B policy competitive with much larger ones. Consistent gains on a humanoid simulation benchmark and on real single-arm and bimanual hardware, with the undistilled policy of identical parameter count serving as the control.
3. Ablations showing the recipe is robust to its own design choices. The gain survives changes of student scale (0.8B, 1B, 4B), student backbone (Qwen3.5-VL, InternVL, Qwen3-VL), alignment layer, and teacher family, indicating a broad representational prior rather than a fragile alignment between two specific networks.
4. A multi-run evaluation protocol that crosses 2 evaluation seeds with 2 GPUs (an A100-80GB and an RTX 5090) and reports mean plus standard deviation, motivated by the observation that neither benchmark is deterministic.

Main Findings

• LIBERO: the distilled 0.8B student reaches 97.9% ± 0.5 four-suite average success, which is 2.6 points above the identical undistilled student (95.3% ± 0.7), with per-suite results of 99.3 ± 0.6 (spatial), 99.2 ± 0.7 (object), 99.4 ± 0.5 (goal), and 93.8 ± 1.1 (libero-10). Every other policy below 4B in the comparison table sits between 78.7% and 95.3%.
• RoboCasa-GR1 humanoid manipulation: the same 0.8B student improves from 48.2% ± 2.1 to 50.5% ± 2.3, a gain of 2.3 points, moving it past QwenFAST (39.0%), QwenPI (43.9%), QwenOFT (48.8%), and both Isaac-GR00T N1.5 (48.2%) and N1.6 (47.6%), and to within 4.3 points of the 4B QwenGR00T (54.8% ± 2.0), of which the student is a five-fold reduction in parameters.
• Real single-arm fruit pick-and-place (AgileX Nero): the distilled 0.8B policy reaches 93.3% (28/30), matching the 4B π-style policy (93.3%) and 3 trials above the undistilled control (83.3%).
• Real single-arm egg pick-and-place: every policy drops; the control falls to 46.7% (14/30), distillation recovers 4 of the 7 trials separating it from the 4B policy, reaching 60.0% (18/30) against 70.0% (21/30).
• Real bimanual fruit handover (TRIP-Bag, 7-DoF): the same ordering holds, with 40.0% for the control, 46.7% for the distilled student, and 53.3% for the 4B policy.
• Deployment cost is unchanged by construction: the deployed policy runs in 32 ms and 1.86 GB on a consumer RTX 5090, against π0.5 at 65 ms on the same card and DreamZero at 45.9 GB (which the 32 GB card cannot hold). At inference the student uses one backbone prefill and four flow steps.
• Teacher ablation on LIBERO: every teacher improves on the 95.3% undistilled control — V-JEPA2-AC by 1.2 points (96.9% ± 0.8), Fast-WAM by 1.6 points (96.9% ± 0.8), and Cosmos3-Nano by 2.6 points (97.9% ± 0.5). Because all three lift the same student, the paper attributes the transfer to world-model representations in general rather than to one particular teacher.
• Student scale and backbone ablation on RoboCasa-GR1 (single A100 run): Qwen3.5-VL 0.8B improves 50.3 → 52.8, InternVL 1B improves 50.9 → 53.3, and Qwen3-VL 4B improves 56.8 → 58.4.
• Alignment-layer ablation (quarter schedule, single A100 run): the final layer is best at 49.8%, half-depth is within half a point at 49.4%, two-thirds depth is weakest at 45.9%, and one-third depth gives 47.5%; the ordering is not monotonic in depth, and the authors read the objective as tolerant of where it is attached.
• Failure modes are placement errors, not recognition errors: the policy consistently finds and reaches for the target, but releases early on the way to the receptacle, closes on the receptacle's own geometry when the two overlap in the image, or slips on eggs during grasp and mid-carry (where it still tracks and re-reaches). The authors note this is precisely what a representational prior would be expected to help with least.

Methodology in Plain English

The student is a QwenGR00T policy: a Qwen3-VL-family vision-language backbone paired with a GR00T-style flow-matching action expert that predicts a chunk of K future actions by velocity regression; this is trained with the ordinary action loss, supervised throughout by ground-truth demonstration actions, with no teacher action targets used (the teacher and student do not share an action parameterization, and a recipe depending on teacher actions could not be teacher-agnostic).

One term is added. For each camera view, the student's final-layer hidden states over that view's image tokens are mean-pooled into one vector per view; a two-layer MLP projector maps it to the teacher's feature width; and a cosine loss (with a stop-gradient on the teacher target) pulls that vector into alignment with the teacher's corresponding pooled feature. The total objective is the action loss plus λ_align times the alignment loss, with λ_align = 0.5 everywhere.

The efficiency comes from two choices. First, the teacher is frozen, so its targets depend only on the training frames: they are computed once ahead of time on four GPUs (about an hour for LIBERO), written to a memory-mapped cache keyed by (trajectory id, base index) with one pooled vector per camera view per row, and read back during training. No teacher weights are ever loaded during training, teacher and student can live in incompatible environments, and one cache can serve multiple students because the projector auto-sizes to whatever student/teacher width pair it is given. Second, the student is asked to agree in direction rather than reproduce exact values, so it keeps whatever additional structure its action objective demands and the two backbones need not share a feature space or dimensionality.

The main teacher is the understanding tower of Cosmos 3, an omnimodal mixture-of-transformers world model; the authors extract its Qwen3-VL-8B reasoner from the unified checkpoint, drop the other towers, and read image-token hidden states at layer 24, mean-pooled per view, giving 4096 dimensions per view. The tower is deterministic, with no diffusion timestep or sampling. Because the alignment head is auxiliary, it is discarded when training ends, and the deployed policy is architecturally identical to the undistilled baseline, down to the number of flow steps.

Training used 4× A100-80GB with DeepSpeed ZeRO-2 at effective batch size 256 under a cosine schedule. Evaluation crossed 2 seeds with 2 GPUs, reporting the mean over four runs with standard deviation; LIBERO spread is around one to one and a half points, while RoboCasa-GR1 can move by several points when the GPU changes.

Why This Matters

• Impact on research: It reframes the world-model-versus-VLA tradeoff, showing that the grounding a world model earns through future prediction lives in its internal features and can be inherited without ever predicting a future. This connects representation-level distillation (in the lineage of FitNets and REPA) to robot policy learning, and it suggests that improvements in world models can be harvested by much smaller policies at no deployment cost.
• Real-world applications:
  • Warehouse and fulfillment picking, where single-arm pick-and-place reliability (93.3% here) directly maps to throughput.
  • Food handling and packing, where the egg task's contact-sensitivity (60.0% versus 46.7% for the control) illustrates both the value and the current limits of the approach.
  • Bimanual and humanoid manipulation, including timed handovers (46.7% on TRIP-Bag) and the 24 GR1 Fourier humanoid environments of RoboCasa-GR1.
  • Edge and on-robot deployment, where the 1.86 GB footprint and 32 ms latency fit hardware budgets that a 45.9 GB world model cannot.
• Industry relevance: Because the teacher is used only offline, the training loop costs the same as the baseline run, the projector is thrown away, and the deployed graph and flow-step count are unchanged, the recipe can be retrofitted into existing VLA pipelines without touching inference stacks. The teacher-agnostic result (three different teachers all improving the same student) lowers the risk of building on any single world model, and one cached teacher target set can be reused across multiple student models.

Future Directions

The paper does not present a future-work section; the following follow from its results and stated limitations.

1. Closing the remaining gap to larger policies: the distilled 0.8B student remains 4.3 points behind the 4B QwenGR00T on RoboCasa-GR1 after also reaching 50.5%, and stronger published entries at 4B or more (up to ACE-Ego-0 at 72.8%, which assembles roughly 600M frames through a five-stage pipeline) remain ahead. Whether the alignment term scales further as teachers improve is untested.
2. Addressing placement and contact failures: the dominant failure modes are releasing early near the receptacle, closing on the receptacle's geometry, and slipping on smooth objects, none of which is a scene-understanding failure. Since a representational prior is least expected to help here, complementary supervision (for example on contact or precise placement) is an open question.
3. Broadening evaluation beyond the reported tasks: real-robot validation covers three tasks on one single-arm platform and one bimanual platform, with roughly 30 minutes of teleoperated demonstration per single-arm task and 1 hour for the bimanual task, and 30 trials per cell. Whether the gain persists across more embodiments, lighting, viewpoints, and clutter is not established.
4. Understanding what is actually transferred: the layer ablation is inconclusive — the objective tolerates where it is attached, and the spread between the best and second-best layers is smaller than a quarter-length schedule can resolve — and the paper leaves open which component of world-model features (temporal structure, causal consequence, or something else) carries the benefit.

Target Audience

Robotics and embodied-AI researchers working on VLA policies, robot manipulation, and efficient deployment; practitioners who need capable policies inside tight latency and memory budgets; and students or engineers interested in knowledge distillation and representation alignment as a practical, low-cost alternative to scaling model size or training data. Readers without a robotics background will still follow the core argument, but the experimental tables assume familiarity with manipulation benchmarks and policy architectures.

Authors’ abstract

Vision-Language-Action (VLA) models map observations to actions with no objective that accounts for how the world responds, so their robustness is bounded primarily by data coverage. World models carry precisely that missing objective and are better grounded for it, yet rolling the future forward costs seconds per decision and rules them out of the control loop. We show the two can be separated. What a world model knows about physical scenes lives in its internal features; generating the future is merely the objective that produced them, so the grounding can be inherited while the generative machinery is left behind. We add one feature-alignment term to ordinary VLA training: a frozen world model is run over the training frames once and cached, and the student learns to agree with that cache. No teacher is loaded during training, the projector is discarded after it, and the deployed policy is identical to the undistilled baseline, running in 32ms and 1.86GB on a consumer RTX5090, so every gain is attributable to the representation rather than to added capacity or test-time compute. A 0.8B student reaches 97.9% on LIBERO, improves from 48.2% to 50.5% on RoboCasa-GR1 humanoid manipulation, and the same objective carries over to real hardware, on both a single-arm and a bimanual platform. The gain survives changes of student scale, backbone, alignment layer, and teacher, indicating a broad representational prior rather than a fragile alignment between two particular networks. Project page: https://thaw-vla.trung-dt.com/.