Off-policy reinforcement learning (RL) has become increasingly sample-efficient, enabling applications such as RL fine-tuning of Vision-Language-Action models into reliable, high-performing policies. World models offer a further lever for sample efficiency, as they predict state changes rather than actions alone, but their success has largely been confined to supervised policy learning. Prior model-based RL methods often optimize the policy or value function directly on imagined rollouts, which is prone to compounding bias and struggles to scale to large, high-dimensional problems such as real-world robotics, a problem that worsens with task horizon and visual complexity. In this work, we instead ask whether we can leverage world models directly on top of standard Q-learning to improve performance, while remaining trained and grounded in the real, online setting. We propose QWM, a framework that leverages world models to perform test-time search over imagined trajectories on top of Q-learning to select high-value actions during both online rollouts and evaluation. Since the policy and value function are trained only on real transitions, QWM avoids compounding model bias while still gaining the sample-efficiency benefits of predictive search. On challenging manipulation benchmarks Robomimic and LIBERO, QWM significantly outperforms strong prior state-of-the-art methods on both sample efficiency and performance.
Figures & tables
Figure 1: Overview of QWM . Left: World-model tree search. From the current state, the policy proposes candidate actions and the world model imagines future states. The value of each initial action is computed by aggregating value of intermediate and leaf nodes. Right: Online action selection. At each environment step, QWM uses tree-search scores to select among policy-proposed actions before executing the selected action in the environment. Tree search is used for both online sampling and evaluation, while the policy and critic are trained only on real environment transitions.
Figure 2: Tree Search and Value Aggregation in QWM . From the current state s0 , the policy proposes N candidate actions. The world model predicts the corresponding future states, and this expansion repeats to depth D . Rather than evaluating only the leaf nodes, values are aggregated from intermediate nodes, where deeper imagined branches are discounted according to the tree-search discount λ .
Figure 3
Figure 3: Compared with model-based RL methods. Success rates of QWM , TD-MPC2, and EZ-V2 across Lift, Can, Square, and Tool Hang under sparse and dense reward settings. QWM achieves strong performance across challenging manipulation tasks.
Figure 4: Effect of world-model tree search on EXPO. Online success rates of EXPO and QWM on Tool Hang, Square, and Can. Tree search consistently improves sample efficiency, with larger gains on Tool Hang and Square.
Figure 5: Effect of world-model tree search on RLPD. Online success rates of RLPD, RLPD with tree search, and QWM on Square, Can, and Lift. Tree search improves the sample efficiency of RLPD across all three tasks, while QWM built on EXPO remains more sample-efficient overall.
Figure 6: Pixel-based evaluation on LIBERO. Online success rates of QWM and EXPO across five LIBERO tasks. QWM achieves clear gains on Tasks 60 and 79, improves learning speed on Task 28, and reaches comparable or stronger final performance across the remaining tasks.
Figure 7: Ablation of world-model lookahead on state-based tasks. We compare applying world-model-guided action selection during online sampling, evaluation, or both stages. The sampling + evaluation variant achieves the most consistent improvement across tasks, showing the complementary benefits of improving online experience collection and test-time action selection.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Search with V vs Q on top of EXPO. Online success rates of QWM -EXPO against a search with V variant where the edit policy is learned via a state value function V. Search with Q does substantially better than the V counterpart.
Figure 9: Search with V vs Q as a standalone method. Online success rates of QWM -RLPD against search with V variant with a Gaussian policy learned from V in the style of AWR. Search with Q does substantially better than the V counterpart.
Hyperparameter
Value
Start training steps
5,000
Offline data ratio
0.5
Offline pretraining steps
0
UTD ratio
20
Batch size
256
Critic TD discount γ
0.99
Appendix
Table 1: Shared optimization hyperparameters for QWM and model-free baselines on Robomimic.
Hyperparameter
Value
Candidate actions N
8
World-model samples per action K
8
Intermediate-node aggregation
Max
Leaf-value aggregation
Mean
Search during online sampling
Yes
Search during evaluation
Yes
Appendix
Table 2: Shared tree-search settings for state-based Robomimic experiments.
Figure 10: Qualitative comparison of world-model predictions on LIBERO. We compare the current observation, ground-truth future observations, and action-conditioned predictions over short rollout horizons. The predicted observations capture the major robot and object changes relevant to tree-search action evaluation, while prediction errors become more visible at later rollout steps.
Planning with a learned latent world model is a promising route to control from raw pixels, but a strong world model alone is not enough. We show this experimentally: even with a perfect world model (operationalized by replacing the learned forward predictor with an idealized rollout of the true environment dynamics), a finite-budget sample-based planner still fails on some tasks, indicating that the bottleneck can lie in search rather than in world-model accuracy. Motivated by this gap, we propose IMWM (Intuition Model + World Model), which pairs the world model with an intuition model trained from demonstrations to recognize promising actions. The two models collaborate through three lightweight components: (i) Retrieval Initialization, which initializes the planner's action proposal from a retrieved demonstration; (ii) Hybrid Cost, which combines the intuition score with the world-model rollout cost; and (iii) a Reliability Gate, which adjusts how much the planner trusts intuition in each setting. Across four pixel-based goal-reaching tasks (Two-Room, Reacher, Push-T, and OGBench-Cube), IMWM has higher mean success than the world-model-only planner on all four, with the largest gains on Two-Room (99.2%, +11.5 percentage points) and OGBench-Cube (94.7%, +28.5 percentage points).
Baoqi Gao, Ruize Han, Miao Wang +1
Beihang University · Shenzhen University of Advanced Technology
This paper presents the World-Action Model (WAM), an action-regularized world model that jointly reasons over future visual observations and the actions that drive state transitions. Unlike conventional world models trained solely via image prediction, WAM incorporates an inverse dynamics objective into DreamerV2 that predicts actions from latent state transitions, encouraging the learned representations to capture action-relevant structure critical for downstream control. We evaluate WAM on enhancing policy learning across eight manipulation tasks from the CALVIN benchmark. We first pretrain a diffusion policy via behavioral cloning on world model latents, then refine it with model-based PPO inside the frozen world model. Without modifying the policy architecture or training procedure, WAM improves average behavioral cloning success from 59.4% to 71.2% over DreamerV2 and DiWA baselines. After PPO fine-tuning, WAM achieves 92.8% average success versus 79.8% for the baseline, with two tasks reaching 100%, using 8.7x fewer training steps.
Yuci Han, Alper Yilmaz
Photogrammetry and Computer Vision Lab, The Ohio State University, Columbus, OH 43210, USA.
Video world models are emerging as a scalable alternative for evaluating generalist robot policies, bypassing the physical constraints and engineering burdens of real-world deployment. However, evaluating policies with video world models remains challenging, as world-model errors can make generated rollouts unreliable and slow inference limits large-scale throughput. We introduce RoboWorld, an automated evaluation pipeline that pairs a fast autoregressive video world model with a task-progress-aware vision-language model scoring. To enable reliable long-horizon autoregressive world-model rollouts, we propose Step Forcing, which combines anchored and one-step self-forwarded contexts to reduce train-test mismatch while preserving action-observation dynamics. Together, these components enable RoboWorld to align strongly with real-world robot evaluation across tasks and environments, achieving Pearson's r = 0.989 and Spearman's ρ = 0.970.