Vision-Language Agents Can Learn in Their Own Latent Worlds
Authors: Haoxi Li, Sikai Bai, Jianfei Ma, Qinglin Hou, Xiaocheng Lu, Tao Han, Jie Zhang, Song Guo
Organizations: Department of Computer Science and Engineering, HKUST, Hong Kong, China · School of Computing, National University of Singapore, Singapore · Department of Computer Science, University of Southern California, Los Angeles, USA
Learning from experience is essential for building the next generation of agents that continually expand their capabilities. Pretrained vision-language models (VLMs) already internalize rich world knowledge, yet this knowledge is rarely leveraged to drive policy improvement through online interaction. Realizing this potential efficiently remains challenging, as policy updates shift the agent's interaction distribution, requiring the world model to adapt accordingly, while continually updating large pretrained models is computationally expensive. To address this challenge, we introduce Dwell, a model-based reinforcement learning (RL) framework that reframes a VLM as a recurrent latent transition operator, endowing it with the ability to predict, adjust, and act in latent space. Leveraging the VLM's own representations and its dual role as both world model and policy, the agent continually learns from real interactions to anticipate and evaluate action outcomes, and from latent imagination to improve its behavior. Across five interactive visual environments, Dwell achieves the strongest overall performance among trainable methods while substantially reducing the compute required for imagination, showing that effective world models for VLM agents need not reconstruct the future; they need to predict futures that are useful for action. Our code is publicly available at https://github.com/Tom111-lhx/Dwell
Figures & tables
Figure 1: Overview of Dwell . Left: The VLM learns a recurrent latent world by matching action-conditioned priors to observation-conditioned posteriors, while grounding the latent states with a frozen V-JEPA2 encoder. Right: The same VLM backbone serves as both world model and policy, alternating online interaction with latent planning for policy improvement.
Figure 2: Learning Efficiency. Success rate throughout training on Sokoban and FrozenLake . Left : comparison with model-based approaches. Right : comparison with model-free approaches.
Figure 3: Computational Efficiency. End-to-end GPU-hours required to reach the reported training step on FrozenLake , Sokoban , Navigation , and EB-ALFRED . GPU-hours include all computation incurred during online training, including world-model updates, imagined rollouts, and policy/value optimization. All methods are evaluated under the same environment-specific training schedule and hardware configuration.
Figure 4: Ablation Studies. Left: effect of removing the representation-grounding objective Lrep or the predictive dynamics objective Ldyn . Right: effect of freezing the world model after initialization instead of continually updating it with newly collected experience. Predictive dynamics provide the primary supervision for useful latent imagination, while representation grounding and continual world-model adaptation further improve policy learning.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
JWM=Ldyn+αLrep+βLrew+ηLreg.
Appendix
Algorithm 1 Dwell
Variant
Sokoban
FrozenLake
Navigation
Dwell (Shared)
0.53±0.012
0.51±0.021
0.39
Bridged Latent Interface
0.38±0.138
0.42±0.097
0.31±0.017
Appendix
Table 2: Effect of the shared latent interface. We compare the default Dwell , where predicted latent states are consumed directly by the policy, with a bridged variant that maps world-model states into a separate policy latent space through a learned adapter. Results are success rates averaged over three training seeds.
Environment
Observation
Actions
Horizon
Eval. episodes
Sokoban
top-down RGB, 6×6 , one box
4 moves
25/25
256
FrozenLake
top-down RGB, 4×4
4 moves
25/25
128
Navigation
first-person RGB, instruction
8 motions
30/10
120
PrimitiveSkill
RGB, instruction
3 parameterized skills
3/3 turns
64
EB-ALFRED
first-person RGB, instruction
162 skills
30/30
20
Appendix
Table 3: Environment settings. Horizon is the maximum number of training / evaluation decisions. PrimitiveSkill observations also include task-state text; each episode has three language turns with at most two primitive commands per turn (at most six in total), and its 64 evaluation episodes cover four tasks (16 each).
Real reward
Dwell in imagination
Environment
Success
Step
Label
Reward
Episode end
Sokoban
+10
−0.1
binary
1.11[p≥τ]−0.01
p≥τ
FrozenLake
+10
0
3-class
0.11[pgoal≥0.5]
goal or hole
Navigation
+10
0
binary
1.1p−0.1
cutoff at H
PrimitiveSkill
+1
0.1+Δs
binary, stage
1.1p−0.1
success or H
EB-ALFRED
+1
−0.1
binary
1.1p1[t=H]−0.1
cutoff at H
Appendix
Table 4: Reward assignment. p is the predicted success probability, τ a calibrated threshold, p a damped version of p , and t the imagination step. For PrimitiveSkill, the step reward is a 0.1 format reward plus Δs , the increase in the highest completed task stage, and the feedback head also predicts the stage.
Pretrained vision-language-action (VLA) policies show promising zero-shot generalization, but often fail under deployment-time distribution shift, leading to decreased robustness and inconsistent instruction following. While prior work commonly tackles this by finetuning on in-distribution data, it assumes demonstrations collected on tasks in the target environment. In this work, we propose DREAMSTEER, a deployment-time steering framework for pretrained VLAs without any finetuning or parameter modifications. The key insight in DREAMSTEER is to leverage a latent world model and a value model to steer pretrained VLA policies. During deployment, DREAMSTEER samples candidate action chunks from a VLA policy and predefined motion primitives, imagines their outcomes using an action-conditioned latent world model, and ranks the imagined trajectories with a language-conditioned value model. Across four real-world manipulation benchmarks with unseen objects, DREAMSTEER improves task success rate from 23.75% to 66.25% and instruction-following accuracy from 38.75% to 56.25% over the base VLA policy.
Hanchen Cui, Sergio Arnaud, Arjun Majumdar +5
1Fundamental AI Research (FAIR), Meta · University of Minnesota Twin Cities
Post-training VLA policies typically rely on supervised fine-tuning with costly expert demonstrations or reinforcement learning with expensive and potentially unstable real-world exploration. World models offer a promising alternative by evaluating candidate behaviors through imagined futures, yet effective post-training requires more than accurate prediction: imagination must be scheduled where it is useful, bounded within reliable horizons, and translated into trustworthy policy supervision. In robotic manipulation, the value of imagination varies substantially across execution stages, while extended rollouts can accumulate prediction errors and introduce unreliable learning signals. We introduce WISE (World-model-guided Imagination Scheduling for Efficient Post-training of Vision-Language-Action Models), a unified framework that coordinates when and how world-model imagination is used during policy refinement. WISE selectively invokes imagination at interaction-relevant states, performs bounded multi-view rollouts, evaluates candidate futures using progress and completion signals, and uses their relative outcomes to refine actions generated from real interaction contexts. Extensive experiments with both π0 and π0.5 demonstrate consistent improvements across diverse manipulation tasks while reducing GPU computation time by approximately 80% compared with full imagination. Real-world evaluations further show substantial gains in robustness and generalization under diverse real-world distribution shifts.
Chenhao Zhang, Hanyu Zhao, Hang Cheng +2
1Tsinghua University · 2Beijing Academy of Artificial Intelligence (BAAI)
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.