Q-Learning With World Models — part 1/4

Cleaned-text code points [0, 15595). Bibliography omitted here; source.md includes it.

# Q-Learning With World Models

###### 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.

## 1 Introduction

Recent advances in off-policy reinforcement learning (RL) have substantially improved its sample efficiency, enabling applications such as RL fine-tuning of Vision-Language-Action (VLA) models into highly reliable, high-performing policies ( Intelligence et al., 2025 ; Dong et al., 2026a ) . As RL is applied to increasingly complex scenarios, further gains in sample efficiency become increasingly valuable, if not essential. World models offer a natural lever for closing this gap, as they have the ability to predict changes in state rather than actions alone. So far, however, their success has been largely confined to supervised policy learning. This raises a natural question: is there a simple way to leverage world models to improve the performance of online RL, without resorting to traditional model-based RL approaches that are prone to compounding the world model’s bias into the training process?

Prior work on model-based RL has explored learning a dynamics model and optimizing a policy or value function directly within it ( Hansen et al., 2024 ; Hafner et al., 2024 ) , but scaling this approach to large, high-dimensional problems such as real-world robotics remains difficult: policies trained primarily on rollouts inside an imperfect world model inherit and compound its biases, and this problem worsens as task horizon and visual complexity increase. Rather than using the world model as a substitute for real interaction, we ask whether it can instead sharpen an RL agent that remains trained and grounded in the real, online setting. In this work, we study Q-learning with a world model, where the model serves not as a replacement for environment interaction but as a way for improving performance.

Our key insight is that a learned world model can be used on top of standard Q-learning to enable test-time search over actions by imagining future trajectories, and that this improved action selection alone directly translates into higher performance for Q-learning methods. We propose Q-Learning with World Models ( QWM ), a framework that leverages world models for test-time search directly on top of standard Q-learning. Instead of using the state value function to conduct search as in typically done in model-based RL, the learned Q-function also enables more powerful search to evaluate candidate actions and select the highest-value one before it is executed. This test-time search over actions is applied both during online rollouts, to obtain better data for online RL, and at evaluation time, to act more effectively with the learned policy. Rather than sampling a single action from the policy or sampling multiple actions and selecting the one with the highest Q-value, QWM leverages the imagination capability of world models to select the best action for downstream performance, while avoiding the compounding bias and instability of policies trained inside a learned model.

Our main contribution is QWM , a simple and general framework for leveraging learned world models to improve Q-learning through test-time search. Because the underlying policy and value function are still trained online against real environment transitions, QWM avoids the bias accumulation that afflicts methods relying primarily on imagined rollouts, while still reaping the sample-efficiency benefits such a predictive model can offer at decision time. We evaluate QWM on challenging manipulation benchmarks, Robomimic and LIBERO, across both state-based and visual observation settings, and find that it significantly outperforms strong prior methods on both sample efficiency and performance.

## 2 Related Work

Reinforcement learning with prior data. To improve performance and sample efficiency of online RL, prior works have studied the problem of using an offline dataset to accelerate online learning. Popular approaches involve balancing exploration and exploitation ( Yang et al., 2023 ; Zhang et al., 2023 ; Mark et al., 2023 ) , calibrating value estimates ( Nakamoto et al., 2024 ; Dong et al., 2026c ) , or retaining offline data alongside newly collected online data during fine-tuning ( Vecerik et al., 2018 ; Nair et al., 2021 ; Ball et al., 2023 ; Dong et al., 2025a ; Dong et al., 2025b ) . Offline RL algorithms have also been applied directly to online fine-tuning ( Fujimoto et al., 2019 ; Fujimoto and Gu, 2021 ; Hansen-Estruch et al., 2023 ; Park et al., 2025 ; Dong et al., 2025c ; Dong et al., 2026b ; Dong et al., 2026e ) , where the prior data is first trained with offline RL before obtaining new samples online. Our work differs from these lines of research in that we focus on leveraging prior data through world models to improve online RL, rather than on how prior data is incorporated into the RL objective itself; this makes our approach applicable to any RL fine-tuning algorithm.

Model-based RL. Model-based RL methods learn a dynamics model of the environment and use it for policy or value learning. One family of methods generates additional synthetic transitions by rolling the learned model forward and treats them as if they were real experience ( Sutton, 1991 ; Gu et al., 2016 ; Kalweit and Boedecker, 2017 ; Kurutach et al., 2018 ; Kaiser et al., 2024 ) , or otherwise optimizes the policy directly against imagined trajectories ( Zhang et al., 2019 ; Hafner et al., 2020 ; Hafner et al., 2022 ; Hafner et al., 2024 ; Hafner et al., 2025 ) or through short branched rollouts mixed with real data ( Feinberg et al., 2018 ; Janner et al., 2021 ) . A second family instead uses the learned model only to plan over imagined rollouts with trajectory optimizers such as the cross-entropy method or MPPI ( Nagabandi et al., 2018 ; Ebert et al., 2018 ; Chua et al., 2018 ; Hansen et al., 2022 ; Hansen et al., 2024 ) ; some of these methods further learn a terminal value function to bound the planning horizon ( Hansen et al., 2022 ; Hansen et al., 2024 ; Kim et al., 2026 ) . A third family combines a learned (or known) model with Monte Carlo tree search guided by a learned value function, mainly for discrete-action, densely-simulated domains ( Silver et al., 2017 ; Schrittwieser et al., 2020 ; Ye et al., 2021 ; Schrittwieser et al., 2021 ) , with recent works extending it to continuous control ( Hubert et al., 2021 ; Wang et al., 2024 ) . These methods rely on some combination of having a known dynamics model, discrete actions, or bootstrapping policy and value targets from search statistics computed over model-imagined rollouts. Furthermore, many of these works rely on abundant simulated data, where model error and sample budget are far less of a concern than in real-world, sparse-reward robotic manipulation. Across all of these approaches, the policy or value function trained substantially on model-generated rollouts inherits and compounds the model’s own biases directly into training, an issue that grows with task horizon and visual complexity ( Janner et al., 2021 ; Chua et al., 2018 ; Kaiser et al., 2024 ; Hafner et al., 2020 ) . QWM sidesteps this issue by leveraging the model purely at test-time rather than a source of additional training data on top of standard Q-learning, and train the policy and critic on real data.

#### World models for VLAs.

A closely related line of work applies world models to VLA policies. VLAW ( Guo et al., 2026 ) and World-VLA-Loop ( Liu et al., 2026 ) alternate between fine-tuning world models on real rollouts and using it to generate synthetic training data for policy learning. WMPO ( Zhu et al., 2025 ) and World4RL ( Jiang et al., 2026a ) run on-policy RL directly on imagined rollouts, while WoVR ( Jiang et al., 2026b ) targets the resulting reward hallucination and OOD drift via keyframe-initialized rollouts and policy-aligned co-evolution. RISE ( Yang et al., 2026 ) , World-Gymnast ( Sharma et al., 2026 ) , and GigaBrain-0.5M ( Team et al., 2026 ) score imagined rollouts with a learned progress/value model or VLM and use the resulting rewards for policy-gradient updates. All of these use the world model primarily as a training data generator, which is prone to inheriting compounding bias. SAILOR ( Jain et al., 2025 ) learns a world model and reward model to guide imitation and recover from out-of-support states. QWM can be applied to any RL fine-tuning methods for VLAs with a Q-function, and instead of using the world model to generate data or guide imitation learning, we leverage world models to improve Q-learning at test time.

#### Test-time scaling and best-of- N N inference.

A substantial body of work has demonstrated that test-time scaling, for example via best-of- N N sampling, can significantly improve the performance of generative models across domains, including autoregressive language models ( Wang et al., 2022 ; Snell et al., 2024 ; Chow et al., 2024 ) , diffusion models ( Ma et al., 2025 ) , and policy learning ( Dong et al., 2025b ; Hansen-Estruch et al., 2023 ; Dong et al., 2026d ; Chen et al., 2023 ) . In most of these settings, candidates are scored independently by a verifier or value function, and only the single best sample is retained, without modeling how a candidate’s near-term consequences unfold. QWM can instead be viewed as a form of test-time scaling in which search evaluates candidates on their predicted downstream value, yielding a substantially better estimate of which actions are best.

## 3 Preliminaries

We consider a Markov decision process (MDP) specified by the tuple { 𝒮 , 𝒜 , ρ , r , γ , T } \{\mathcal{S},\mathcal{A},\rho,r,\gamma,T\} , where 𝒮 \mathcal{S} denotes the state space, 𝒜 \mathcal{A} denotes the action space, ρ ⁡ ( s ) \rho(s) is the distribution over initial states, r : 𝒮 × 𝒜 → ℝ r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R} is the reward function, γ ∈ [ 0 , 1 ) \gamma\in[0,1) is the discount factor, and T ⁡ ( s ′ | s , a ) T(s^{\prime}|s,a) is the transition probabilities governing the environment. The goal of RL is to find a policy π \pi that maximizes the expected discounted return 𝔼 π ​ [ ∑ t = 0 T γ t ​ r ​ ( s t , a t ) ] \mathbb{E}_{\pi}\left[\sum_{t=0}^{T}\gamma^{t}r(s_{t},a_{t})\right] .

In this paper, we study how a learned world model can be used to improve online RL fine-tuning at decision time. Concretely, we assume access to a dataset 𝒟 offline = { ( s , a , r , s ′ ) } \mathcal{D}_{\text{offline}}=\{(s,a,r,s^{\prime})\} of environment transitions – either collected offline or accumulated online – which is used to pretrain a world model M ψ ​ ( s ′ | s , a ) M_{\psi}(s^{\prime}|s,a) that predicts the next state given the current state and action. The world model is used at decision time. During online training, the agent collects tuples ( s t , a t , r t , s t + 1 ) (s_{t},a_{t},r_{t},s_{t+1}) through environment interaction; these are appended to a replay buffer 𝒟 \mathcal{D} and used to update the policy and Q-function toward higher returns. Our central question is how to leverage the world model to improve the performance of Q-learning.

We focus on off-policy RL, where the Q-function estimates the discounted return of the policy given a state and action, and is trained with TD learning: ℒ ⁡ ( ϕ ) = 𝔼 ( s t , a t , s t + 1 ) ∼ 𝒟 ​ [ ( r t + γ ​ Q ϕ ′ ​ ( s t + 1 , a ~ t + 1 ∗ ) − Q ϕ ​ ( s t , a t ) ) 2 ] , \mathcal{L}(\phi)=\mathbb{E}_{(s_{t},a_{t},s_{t+1})\sim\mathcal{D}}\left[\left(r_{t}+\gamma Q_{\phi^{\prime}}(s_{t+1},\tilde{a}^{*}_{t+1})-Q_{\phi}(s_{t},a_{t})\right)^{2}\right], (1) where Q ϕ ′ Q_{\phi^{\prime}} is a target network and a ~ t + 1 ∗ \tilde{a}^{*}_{t+1} is the next action selected by the RL policy. We implement QWM on top of two state-of-the-art off-policy Q-learning algorithms.

EXPO. The first is EXPO ( Dong et al., 2025b ) , a high-performing online RL method that couples a base policy with a lightweight edit policy. EXPO maintains two parameterized policies: a base flow policy—in our case, the base policy π base \pi_{\text{base}} , trained with a supervised loss—and an edit policy π edit \pi_{\text{edit}} , trained to maximize the Q-function: ℒ ⁡ ( π edit ) = − 𝔼 ( s t , a t ) ∼ 𝒟 , a ^ t ∼ π edit ​ [ Q ϕ ​ ( s t , a t + a ^ t ) − α ​ log ⁡ π edit ​ ( a ^ t ∣ s t , a t ) ] . \begin{split}\mathcal{L}(\pi_{\text{edit}})=-\mathbb{E}_{(s_{t},a_{t})\sim\mathcal{D},\;\hat{a}_{t}\sim\pi_{\text{edit}}}\bigl[Q_{\phi}(s_{t},\,a_{t}+\hat{a}_{t})-\alpha\log\pi_{\text{edit}}(\hat{a}_{t}\mid s_{t},a_{t})\bigr].\end{split} (2) The edit policy predicts an edit a ^ t \hat{a}_{t} , clipped to [ − β , β ] [-\beta,\beta] , that is added to the base action a t a_{t} to produce an edited action a ~ t = a t + a ^ t \tilde{a}_{t}=a_{t}+\hat{a}_{t} . This design avoids backpropagating gradients through the base policy while still grounding TD updates in near-optimal actions. The final inference and TD-backup policy is an on-the-fly (OTF) policy that selects the value-maximizing candidate among N N sampled base and edited actions: a ~ ∗ = arg ⁡ max a ∈ ⋃ i = 1 N { a i , a ~ i } ​ Q ϕ ​ ( s , a ) , \tilde{a}^{*}=\underset{a\;\in\;\bigcup_{i=1}^{N}\{a_{i},\,\tilde{a}_{i}\}}{\arg\max}\;Q_{\phi}(s,a), (3) where Q ϕ Q_{\phi} is trained with a standard TD objective.

RLPD. The second is RLPD ( Ball et al., 2023 ) , a sample-efficient off-policy RL algorithm that combines offline and online replay data with high update-to-data (UTD) ratios, an ensemble of critics, and a Gaussian actor. RLPD trains the actor π θ \pi_{\theta} with a SAC-style ( Haarnoja et al., 2018 ) entropy-regularized objective, ℒ ⁡ ( π θ ) = − 𝔼 s t ∼ 𝒟 , a t ∼ π θ ​ [ Q ϕ ​ ( s t , a t ) − α ​ log ⁡ π θ ​ ( a t ∣ s t ) ] , \mathcal{L}(\pi_{\theta})=-\mathbb{E}_{s_{t}\sim\mathcal{D},\;a_{t}\sim\pi_{\theta}}\bigl[Q_{\phi}(s_{t},a_{t})-\alpha\log\pi_{\theta}(a_{t}\mid s_{t})\bigr], (4) and stabilizes the high-UTD critic updates using an ensemble of M M critics with random subsets used for the target computation, together with layer normalization. At each update, transitions are drawn via symmetric sampling, with equal proportions from the offline buffer and the online replay buffer.

## 4 Method