Year: 2026
Type: preprint
Abstract: 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.
Source: arXiv (Cornell University)
Institutions: Stanford University, Peking University
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: green
Awards: G1453181110, Track 4: Advanced CI Coordination Ecosystem: Monitoring and Measurement Services, G2154383771, G3562386289, Track 1: ACCESS Resource Allocations Marketplace and Platform Services (RAMPS), G4010157413, ACO: An Open CI Ecosystem to Advance Scientific Discovery (OpenCI), G4720003262, Category I: Crossing the Divide Between Today's Practice and Tomorrow's Science, Track 3: COre National Ecosystem for CyberinfrasTructure (CONECT), G5367127336, G5464919783, CAREER: Safe and Influencing Interactions for Human-Robot Systems, G6671297155, Track 2: Customized Multi-tier Assistance, Training, and Computational Help (MATCH) for End User ACCESS to CI, G8876996369
OpenAlex ID: W7203793813
Cites:
Cited by:
Related to: