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.

• The paper introduces QWM, a model-based action-selection wrapper that searches over policy-proposed futures while updating policies and critics exclusively from real environment transitions.
• QWM improves sample efficiency and task performance across Robomimic and LIBERO manipulation benchmarks, outperforming its EXPO and RLPD backbones and reported model-based baselines under the evaluated protocols.
• Ablations show that action-conditioned Q-values, search during both data collection and evaluation, depth-two planning, conservative future-value weighting, and compact beam pruning are important for balancing lookahead benefits against model and critic errors.

Problem formulation and central claim

“Q-Learning With World Models” (2608.17163) addresses a specific tension in model-based reinforcement learning: learned dynamics models can improve decision quality and sample efficiency, but incorporating model-generated trajectories into policy or value training exposes the learner to compounding model error. The problem is especially acute in sparse-reward robotic manipulation, where long horizons, high-dimensional observations, and limited online interaction make both model accuracy and exploration difficult.

The paper’s central claim is that a world model need not be used as a source of synthetic training data. Instead, it can be used at decision time to search over imagined futures proposed by an existing Q-learning policy. The resulting framework, Q-Learning with World Models (QWM), retains the standard off-policy training pipeline: the policy and critic are updated exclusively from real environment transitions, while the world model is used to rank candidate actions during online data collection and evaluation. This separation is intended to preserve the empirical stability of model-free Q-learning while adding the predictive benefits of short-horizon planning.

The claim is stronger than a conventional best-of-$N$ action-selection procedure. Methods such as IDQL evaluate independently sampled actions using $Q(s,a)$ and execute the highest-valued candidate. QWM instead evaluates candidate root actions according to their predicted downstream consequences. Each candidate is expanded through a learned transition model, additional actions are sampled at predicted states, and the resulting intermediate and terminal values are aggregated recursively. The action-selection mechanism therefore uses the Q-function both as an action-conditioned value estimator and as a heuristic for prioritizing imagined branches.

QWM methodology

QWM constructs an alternating state-action tree from the current state. At each state node, the policy samples $N$ candidate actions. For each action, the world model generates $K$ possible next states, and this expansion continues to depth $D$. In principle, the tree contains $N^D$ action paths, so the implementation retains only the top $J$ partial paths according to accumulated discounted Q-values. This beam-like pruning makes the search computationally tractable and biases expansion toward branches that already appear promising under the critic.

The paper uses two complementary value estimators. The first evaluates candidate actions directly with the learned Q-function. This estimator is relatively low variance and does not depend on rollout depth, but it does not exploit the information contained in imagined future states. The second estimator recursively combines predicted rewards and values at future model states. It benefits from lookahead but is vulnerable to both dynamics-model error and critic error propagated through the imagined trajectory. QWM averages the two estimators, although the formulation permits a weighted combination.

At the root, each sampled action receives a tree-search score combining its immediate Q-value with the value of predicted successor states. The selected action is then obtained through a maximum or softmax operation over the root candidates. Crucially, search is applied in two distinct phases. During online sampling, it changes which real transitions enter the replay buffer. During evaluation, it improves execution without changing the learned policy or critic. The paper argues that these effects are complementary: sampling-time search can improve the quality of subsequent training data, whereas evaluation-time search can improve behavior even when the policy itself has not changed.

The implementation is deliberately modular. QWM is instantiated on top of EXPO and RLPD, two off-policy methods with different policy and critic parameterizations. EXPO combines a supervised base flow policy with a Q-optimized edit policy, while RLPD uses a Gaussian actor, high update-to-data ratios, offline-online replay mixing, and an ensemble of critics. QWM leaves the underlying optimization procedures unchanged, which supports the authors’ assertion that it is a wrapper around Q-learning rather than a replacement for a particular RL algorithm.

World-model construction

For low-dimensional Robomimic experiments, the world model is a deterministic residual MLP dynamics model. It predicts a state increment from the current state and 7-DoF action, and is pretrained on demonstration transitions using mean-squared error. The model is then held fixed while online policy and critic learning proceeds. This design simplifies the interpretation of the experiments: the world model does not co-evolve with the online policy, and its predictions do not enter TD targets.

For pixel-based LIBERO experiments, the authors adapt Wan2.2-TI2V-5B (Wan et al., 26 Mar 2025) into an action-conditioned video model. An MLP action encoder produces action tokens that condition the diffusion transformer together with task text. The VAE and text encoder remain frozen, while the action encoder and diffusion transformer are fine-tuned on demonstration videos and aligned action sequences. During search, the model generates short five-frame, $128 \times 128$ clips, but only the predicted next frame is used before recursively conditioning another generation. Inference uses one denoising step, an aggressive computational compromise that makes repeated tree-search queries feasible but also constrains prediction fidelity.

The practical implementation exposes an important distinction between the conceptual method and the actual experiments. Although the general formulation includes a learned reward model, the reported manipulation tasks are predominantly sparse-reward settings in which the authors do not learn a reward model. The search therefore relies heavily on Q-values at intermediate and leaf nodes. This makes QWM less dependent on accurate dense reward prediction, but it also means that the contribution of explicit model-based reward accumulation is not fully isolated in the main sparse-reward evaluations.

Empirical evaluation

The experiments cover four Robomimic tasks—Lift, Can, Square, and Tool Hang—and five LIBERO tasks. These environments span grasping, placement, precision insertion, and long-horizon assembly. The evaluation compares QWM with strong model-free Q-learning methods, including EXPO, RLPD, IDQL, QSM, DSRL, QAM, and FQL, as well as model-based methods TD-MPC2 and EfficientZero V2.

Against model-free baselines, QWM is reported to achieve the strongest performance across all evaluated tasks and to improve sample efficiency consistently. The relevant comparison is not merely between a model-free and model-based algorithm: all methods use online environment interaction and Q-learning-based critics, while QWM adds imagined lookahead to action selection. The result therefore supports the narrower conclusion that predictive test-time search can provide gains complementary to Q-function learning.

The improvements over the base algorithms are also consistent. QWM improves EXPO on Tool Hang, Square, and Can, with particularly pronounced gains on the harder Tool Hang and Square tasks. It likewise improves RLPD on Square, Can, and Lift. The paper further reports that QWM built on EXPO is more sample-efficient than QWM built on RLPD, which is consistent with the quality of the underlying base learner affecting the value of test-time search. This comparison suggests that QWM does not eliminate the importance of the policy and critic being searched; it amplifies a sufficiently capable Q-learning backbone.

The comparison with model-based baselines is presented as especially favorable. Under both sparse and dense reward variants, QWM achieves consistently strong performance across the Robomimic manipulation tasks, whereas TD-MPC2 and EfficientZero V2 obtain nonzero success only on Lift within the reported training horizon. This is a substantial qualitative result, although the supplied paper content does not provide the underlying numerical success-rate values or confidence intervals. The comparison should therefore be interpreted as evidence of a large performance separation under the authors’ protocol, not as a precisely quantified effect size.

The paper attributes this separation to the location of model use in the learning pipeline. TD-MPC2 and EfficientZero V2 use model predictions to construct planning or learning targets, whereas QWM uses the model only to select actions at runtime. The implication is that QWM can exploit model prediction without allowing model errors to enter the critic’s training distribution as synthetic transitions or search-derived targets. This is a plausible mechanistic explanation, but it does not mean that QWM is insensitive to model error: erroneous predictions can still cause the agent to execute poor actions, collect inferior data, or select systematically biased branches.

Pixel-based control and ablations

QWM also improves EXPO in pixel-based LIBERO experiments, although computational constraints prevent the authors from using world-model search during evaluation. The world model is used only during online data collection. Under this restricted setting, QWM learns faster and reaches stronger late-stage performance on Tasks 60, 79, and 29; on Task 28, both methods eventually approach near-perfect success, but QWM reaches high performance earlier. The results indicate that the framework transfers beyond low-dimensional state representations despite the substantially more difficult visual dynamics-modeling problem.

The limitation of evaluation-time search in the pixel setting is important for interpreting these results. The full method is defined to use search both during sampling and evaluation, and the authors’ own ablations show that disabling evaluation-time search reduces performance. Consequently, the visual experiments demonstrate robustness of the sampling mechanism, but they do not yet establish the full benefit of QWM under high-dimensional observations.

The ablations identify several conditions under which search is effective. Applying search during both online sampling and evaluation produces the most consistent learning improvements. Sampling-only search improves the replay distribution, while evaluation-only search improves action execution without altering training data. The combined result is therefore not redundant: it couples better data acquisition with better deployment-time decisions.

A search depth of two generally provides the strongest empirical gain in the reported ablation, while depth one is often insufficient to capture meaningful consequences of an action. Greater depth can help on tasks requiring longer-horizon planning, but it also increases exposure to model error. The authors report that relatively small to moderate values of the recursive discount, particularly $\lambda = 0.2$ in their settings, are most effective. Larger values make action selection more sensitive to accumulated prediction and value-estimation errors; values that are too small suppress the contribution of lookahead. The exact implementation details report a default $\lambda$ of $Q(s,a)$0 under a convention that absorbs a factor of one-half, highlighting a minor but relevant ambiguity between the formal equations and the implementation table.

Increasing the number of candidate actions initially improves performance, because a small candidate set may omit qualitatively distinct behaviors. Beyond a moderate value, however, gains diminish and performance can decline. The paper attributes this degradation to the increased opportunity for model and critic errors to influence selection, in addition to rising inference cost. The number of expanded paths $Q(s,a)$1 has comparatively little effect, suggesting that Q-guided pruning can preserve useful branches with a small beam. This finding supports the use of aggressively constrained search, but it also indicates that QWM’s search is not an exhaustive planner: its effectiveness depends on the critic’s ability to identify branches worth expanding.

The comparison between Q-based and V-based search is one of the paper’s strongest conceptual results. Variants that search using only a state-value function substantially underperform QWM across the evaluated settings. The authors identify two causes: $Q(s,a)$2 provides action-conditioned information that $Q(s,a)$3 cannot supply, and the underlying Q-learning algorithms are stronger than the corresponding value-based policy-learning alternatives. The result directly challenges the common model-based RL practice of using a state value function to guide imagined search. It supports the paper’s more specific design choice of placing search on top of an action-conditioned critic rather than treating the world model and value function as a conventional MCTS-style pair.

Limitations and open questions

QWM introduces nontrivial computational overhead during both online interaction and evaluation. Even with pruning, the method requires multiple policy evaluations, world-model queries, and Q-function evaluations for each executed action. This cost is particularly significant for the pixel-based implementation, where each query invokes a large video diffusion model. The reported LIBERO experiments avoid evaluation-time search for precisely this reason. Thus, the paper establishes an accuracy and sample-efficiency advantage under additional inference computation, but it does not characterize the trade-off in wall-clock time, energy, hardware utilization, or control latency.

The framework also depends on a high-quality action-conditioned world model. Real-data grounding prevents model error from directly contaminating policy and critic targets, but it does not remove model bias from the control loop. Search can still select actions based on hallucinated object motion, inaccurate contact dynamics, or critic estimates evaluated at model-generated states. The qualitative pixel predictions capture task-relevant changes at the next step, while errors become more visible at later rollout steps; this observation is consistent with the paper’s use of short horizons, but it leaves open how QWM behaves under severe distribution shift, stochastic contact dynamics, or tasks requiring reliable long-horizon planning.

The reported benchmarks are simulated manipulation environments initialized from demonstration data. The method’s central motivation concerns real-world robotics, yet the supplied experiments do not establish performance under sensor noise, actuator latency, unmodeled dynamics, partial observability, or hardware-specific safety constraints. In addition, the comparisons do not provide enough numerical detail in the supplied content to assess statistical significance, variance across random seeds, or compute-normalized performance. A remaining empirical question is whether QWM retains its advantage when baselines receive comparable inference-time compute and when model-based methods are tuned specifically for the same sparse-reward manipulation protocol.

Conclusion

QWM presents a clear architectural separation between learning and planning. Q-learning remains grounded in real transitions, while a learned world model is used for short-horizon test-time search over policy-proposed actions. Across Robomimic and LIBERO, the framework improves sample efficiency and success relative to both its EXPO and RLPD backbones, outperforms the reported model-free baselines, and achieves stronger results than TD-MPC2 and EfficientZero V2 under the evaluated protocols.

The paper’s principal technical contribution is not simply adding a dynamics model to Q-learning, but using the model to refine action selection without allowing imagined trajectories to define policy or critic targets. Its ablations further indicate that action-conditioned Q-values, moderate search depth, conservative future-value weighting, and search during both sampling and evaluation are central to performance. The unresolved issue is whether these gains remain favorable after accounting for the substantial computational cost and under the model uncertainty encountered in physical robotic systems.