DRILL: Training World Models to Improve Policies, Not Predict Pixels

AI summary

Overview

Research area: Robot learning / vision-language-action (VLA) policy post-training, video world models, and reinforcement learning (cs.CV, arXiv:2610.30814v1). Author affiliations span Hunyuan, Tencent and Peking University.

Technical level: Advanced. The paper combines a bilevel (inner/outer loop) training framework, meta-gradient estimation via truncated Hessian-vector products, and convergence analysis, alongside a large empirical study across five simulators. The core idea is accessible, but the machinery (GRPO, BPTT truncation, meta-gradients) assumes familiarity with modern RL for foundation models.

Scope: The paper argues that video world models used for VLA reinforcement learning should be trained by the downstream return of the policy that learns inside them, rather than by one-step pixel prediction, and validates this with two method variants (DRILL-VWI and DRILL-IMG) on five manipulation simulators and two VLA backbones.

What This Paper Is About

Recent systems post-train VLA robot policies with reinforcement learning inside a learned video world model instead of a real robot or simulator, which saves enormous amounts of environment interaction. The world model in these pipelines (e.g., WMPO-style recipes) is still trained the way a video predictor is trained: minimize next-frame pixel error. The authors show this is the wrong target, because a model that predicts pixels well is not necessarily the model that produces the most useful policy-gradient signal. DRILL replaces the pixel-fitting objective with one that directly rewards model changes which lead to better policies after the policy has learned inside the model.

Key Contributions

1. Formalizing the objective mismatch. The paper identifies and formalizes that VLA world models are trained for one-step pixel prediction but used to generate policy-learning signals, and shows empirically that pixel fidelity and policy utility can point in opposite directions.

2. The DRILL bilevel framework with two instantiations. DRILL trains the world model to maximize the return, in the target simulator, of the policy produced by an inner GRPO update on imagined rollouts. It is instantiated as DRILL-VWI (a signed, visitation- and advantage-weighted surrogate requiring no extra simulator rollouts) and DRILL-IMG (a meta-gradient estimator differentiating through the inner update with truncated Hessian-vector products and a small simulator batch).

3. Theory connecting the two variants. Proposition 1 shows DRILL-VWI recovers the local one-step component of the DRILL-IMG meta-gradient under score-compatibility and small-residual assumptions (with explicit residual terms for compatibility, trajectory locality, and demo-anchor transfer). Theorem 1 gives a non-asymptotic convergence rate for DRILL-IMG under truncated BPTT, with a truncation-bias term that decays geometrically in the truncation length K.

4. A new diagnostic, the Policy-Utility Residual (PUR). PUR is a method-agnostic measure of the angular gap between the pixel-MSE gradient and the policy-improvement direction; lower PUR is shown to track higher downstream success across the objectives evaluated.

Main Findings

• Better policies despite worse pixels. On the π₀-3B backbone, DRILL-IMG raises mean success rate by 13.1 points over WMPO and 7.9 points over RLVR-World averaged across the five benchmarks, while its world model has 17% higher one-step pixel MSE on the same training data. DRILL-IMG's reported MSE is 9.8 ± 0.4 (×10⁻³) versus 8.4 ± 0.3 for WMPO and 7.2 ± 0.3 for RLVR-World.

• Per-benchmark gains (π₀-3B, Table 1). DRILL-IMG scores 81.3 ± 1.3 on LIBERO-Long, 83.1 ± 1.4 on CALVIN, 55.7 ± 2.0 on RoboCasa, 79.4 ± 1.2 on Meta-World MT50, and 63.1 ± 2.6 on SimplerEnv, versus WMPO's 64.2 / 71.8 / 38.5 / 67.4 / 55.1. Adding RLVR-World pretraining to the world model pushes these to 83.9 / 85.4 / 59.1 / 81.8 / 67.3, a further 3.2-point average improvement.

• The advantage over WMPO is not a compute effect. A compute-matched control shows extra WMPO updates move LIBERO-Long success from 64.2 ± 1.6 to at most 65.1 ± 1.9, far below DRILL-IMG's 81.3 ± 1.3.

• Model-free PPO with the same simulator budget is far weaker. Sim-PPO, which receives the same simulator episode budget as DRILL-IMG but no imagined rollouts, reaches 28.5 / 32.1 / 12.4 / 31.6 / 22.8 across the five benchmarks, well below the world-model methods.

• Horizon determines whether DRILL-IMG beats DRILL-VWI. At the default H = 50, DRILL-VWI recovers 11.9 of the 17.1-point gap between WMPO and DRILL-IMG on LIBERO-Long. The IMG-over-VWI margin grows from 1.8 points at H = 20 to 5.2 at H = 50, 7.5 at H = 100, and 8.1 at H = 200, but shrinks to 0.6 at H = 10 and a non-significant −0.2 ± 1.1 at H = 5. On SimplerEnv, whose effective horizon is below H = 20, DRILL-VWI slightly exceeds DRILL-IMG (64.0 ± 1.9 versus 63.1 ± 2.6).

• DRILL-IMG costs more. It uses about 1.4× more compute per outer step and needs B_real = 30 simulator rollouts per outer step, with K = 2 truncated BPTT.

• Every component of the VWI weighting matters. Averaged over LIBERO, CALVIN, and RoboCasa: visitation-only 64.2 mean success, signed-advantage-only 66.8, confidence-only 60.1, all three combined 69.2, versus 58.2 for WMPO. The advantage exponent peaks at α = 1 (67.9 at α = 2). A minimal value-aware baseline, VA-WM, reaches 62.4 — 6.8 below DRILL-VWI and 11.0 below DRILL-IMG.

• Gains scale with world-model size and transfer across world-model backbones. On LIBERO-Long the DRILL-IMG margin over WMPO is 7.9 points at 300M parameters, 13.1 at 1.2B, and 16.4 at 7B. Ctrl-World reaches 81.3 ± 1.3 SR and a matched Cosmos-Predict2-style model reaches 79.6 ± 1.7.

• The gain comes from tracking the current policy, not from fixed reweighting. Freezing the VWI signed coaching multipliers from a GRPO-warmstart policy reaches only 64.9 ± 1.5 mean SR, 4.3 points below DRILL-VWI.

• Results replicate on a second backbone. On OpenVLA-7B over LIBERO-Long, CALVIN, and RoboCasa the ranking DRILL-IMG > DRILL-VWI > RLVR-World > WMPO is preserved (except a statistical tie on RoboCasa, 50.8 ± 2.3 versus 48.7 ± 2.1, p = 0.12). DRILL-IMG reaches 68.9 mean SR on OpenVLA versus 53.8 for WMPO, a 15.1-point gain.

• Truncated meta-gradients are reasonably faithful. On a 2D gridworld and Meta-World MT10, where the exact meta-gradient is tractable, the K = 2 truncated Hessian-vector product has cosine similarity 0.74 ± 0.04 and 0.71 ± 0.05 to the untruncated target.

• Greater robustness to perturbation. Under LIBERO zero-shot perturbations, DRILL-IMG shows an 18.6% average relative success-rate drop versus 30.2% for WMPO and 25.6% for RLVR-World, with pose perturbation the one exception where DRILL-VWI slightly wins.

• PUR tracks downstream success. Across the PUR-reporting headline rows the Spearman correlation with downstream success rate is ρ_S = −0.89 (95% bootstrap CI [−0.97, −0.71]). Two held-out variants give 1.5 points mean absolute error when final success rate is predicted from checkpoint-level PUR at step 5k. DRILL-IMG's final PUR is 0.36 ± 0.04, versus 1.42 ± 0.08 for WMPO and 0.97 ± 0.06 for RLVR-World.

• Imagination becomes more task-relevant, not prettier. On LIBERO-Long, WMPO plateaus at 0.61 ± 0.02 TRPS (task-relevant prediction score) despite lower pixel MSE, while DRILL-IMG reaches 0.81 ± 0.03. In a forced-choice human study with 120 trials and six blinded raters, visual-quality preferences are not significant, but 62% prefer DRILL-IMG over WMPO when task completion is shown (p < 0.001).

Methodology in Plain English

The setup. A VLA policy takes images and a language instruction and outputs action chunks. A video world model predicts the next frame given the observation history and an action. A frozen vision-language reward model scores frames between 0 and 1. The standard pipeline alternates between fitting the world model to demonstration pixels and running GRPO policy updates inside the model's imagined rollouts.

The reframing. Instead of asking "does the world model predict the next frame accurately?", DRILL asks "does the world model, after the policy learns inside it, produce a policy that does better in the real simulator?" Formally, this is a bilevel problem: an inner loop does one GRPO step using the world model as the environment, and an outer loop scores the resulting policy in the target simulator and assigns credit back to the world model.

The cheap version (DRILL-VWI). Because the expensive outer gradient is hard to compute, DRILL-VWI reweights the ordinary pixel-prediction loss with three factors already available from the imagined GRPO loop: how often the policy visits a state inside the model, the group-relative advantage of the action (signed, so positive-advantage transitions are sharpened and negative-advantage ones can be deliberately under-fitted), and a confidence term that downweights states where the model is unsure. Weights are stopped and clipped, so this is a local update direction rather than a normalized likelihood objective. It requires no extra simulator interaction compared with WMPO.

The faithful version (DRILL-IMG). DRILL-IMG actually differentiates through the inner GRPO step. That splits into two factors: a policy-gradient estimate in the target simulator (computed from a small batch of B_real = 30 rollouts), and a mixed derivative through the imagined update, computed as a Hessian-vector product with K = 2 steps of truncated backpropagation-through-time. The theory explains the truncation bias decays geometrically with K, so shallow truncation suffices in practice.

The diagnostic. PUR measures the angle between the direction that reduces pixel MSE and the direction that improves downstream return. Values near 0 mean pixel-fitting helps the policy, near 1 mean the two are orthogonal, and near 2 mean pixel-fitting actively fights policy improvement. It is estimated with finite differences in the simulator, deliberately independent of any particular training method, so it can be applied to WMPO, RLVR-World, and DRILL alike.

The experiments. Five public manipulation simulators (LIBERO-Long, CALVIN ABC→D, RoboCasa, Meta-World MT50, SimplerEnv), two VLA backbones (π₀-3B and OpenVLA-7B), and five baselines (Scratch-GRPO, WMPO, VLAW, RLVR-World, and Sim-PPO). Headline results use five seeds with Welch's t-tests and 95% CIs; ablations use three. The headline run used about 17k NVIDIA H20-hours, roughly 24 H20 GPUs for four weeks. No real-robot results are included.

Why This Matters

Impact on research. The paper challenges a default assumption in the fast-growing area of world-model-based RL for embodied agents: that a better video predictor is a better training environment. It supplies both a training objective that targets policy utility directly and a measurable diagnostic (PUR) for how misaligned any world-model objective is with policy improvement. The finding that RLVR-World-style prediction-correctness pretraining is complementary rather than competing — stacking gives a further 3.2 points on average — maps out how the two families of methods fit together.

Real-world applications:

• Robot manipulation: Faster, cheaper post-training of general-purpose manipulation policies without the millions of real or simulated environment steps typically required.
• Industrial automation and warehouse robotics: Reducing the amount of physical trial-and-error needed when adapting a policy to a new pick-and-place or assembly task.
• Simulation-light development: DRILL-VWI's zero extra simulator rollout requirement makes policy improvement feasible where a high-fidelity simulator is unavailable or slow.
• Data-efficient policy iteration for hardware teams: The small B_real = 30 rollout budget used by DRILL-IMG gives a template for trading a modest amount of real interaction for a much better learned environment.

Industry relevance. The work comes from Hunyuan, Tencent plus Peking University and targets foundation-scale systems — VLAs paired with video world models above 1B parameters. Its practical selling points are directly industrially relevant: a drop-in surrogate (DRILL-VWI) for existing WMPO-style pipelines, a compute-matched control showing the gains aren't just extra FLOPs, and documented scaling behavior from 300M to 7B world-model parameters.

Future Directions

• Real-robot validation. The authors explicitly flag that all evaluation is simulation-only and leave open sensitivity to real-world perception noise, actuation delay, safety constraints, and sim-to-real distribution shift.
• Cheaper meta-gradients. DRILL-IMG's truncated backpropagation and Hessian-vector products add compute and memory overhead; reducing this cost — or better characterizing when DRILL-VWI is a safe substitute — is named as a key next step.
• Curvature-level analysis of the RLVR stacking effect. The paper offers an interpretation (RLVR pretraining lowers the variance of the supervised prediction gradient

Authors’ abstract

World models are increasingly used to reduce environment interaction when post-training vision-language-action (VLA) policies with reinforcement learning. Yet they are typically trained to predict pixels one step ahead, even though their role here is not merely to render visually faithful videos but to produce learning signals from imagined rollouts that improve the policy. We show that this mismatch is substantial: lower pixel MSE can produce worse policy gradients, and a world model trained for policy utility can yield better downstream policies despite 17% higher pixel MSE on the next frame. We propose the Downstream-Return Imagination Learning Loop (DRILL), a bilevel framework that updates the world model to maximize the policy's return after an inner GRPO update on imagined rollouts. DRILL has two instantiations. DRILL-VWI is a closed-form surrogate, requiring no additional simulator rollouts, that reweights prediction errors by policy visitation, the group-relative GRPO advantage, and model confidence. DRILL-IMG is a full meta-gradient method that differentiates through the inner update using truncated Hessian-vector products. We show that DRILL-VWI recovers the local one-step component of the DRILL-IMG meta-gradient under score-compatibility and small-residual assumptions. Across five manipulation simulators and two VLA backbones, DRILL-IMG raises mean success rate by 13.1 points over WMPO and 7.9 points over RLVR-World, and complements RLVR pretraining rather than competing with it. More broadly, we believe video world models for VLA reinforcement learning can be productively trained by the policy gradients they produce, not only by pixel-level prediction.