Q-Learning With World Models — part 3/4

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

To better understand the significance of different pieces of QWM , we ablate over three key components: (1) the importance of using QWM during online RL sampling, evaluation, or both, (2) the hyperparameters for constructing the tree, and (3) the number of expanded nodes during search to constrain the tree size. We present additional experiments comparing against performing search with only a state value function Appendix B .

Sampling versus evaluation. We first ablate whether QWM is applied during online RL data collection, at evaluation time, or in both. The sampling only variant performs test-time search during online data collection but disables it at evaluation, while the evaluation only variant applies tree search only at evaluation time. The full sampling + evaluation variant uses test-time search in both. As shown in Figure 7 , applying test-time search at both stages provides the strongest and most consistent learning efficiency. The full variant improves faster and more smoothly and reaches high success rates earlier. Sampling-time search affects which transitions are collected and added to the replay buffer, allowing subsequent policy and critic updates to benefit from higher-value online experience. Evaluation-time search, in contrast, directly improves action selection without modifying the learned policy or training data. Performing test-time search at both stages therefore combines improved online experience collection with stronger execution-time action selection, which gives more consistent advantage over either single-stage variant.

Tree-search hyperparameters. We next ablate key hyperparameters for constructing the search tree: search depth D D , recursive discount factor λ \lambda , and the number of action candidates N N . In Figure 7 , we see that search depth has a clearer effect on learning. Depth 2 2 generally provides the strongest performance gain in our experiments, whereas depth 1 1 improves more slowly as a shallow search may not capture enough future consequence to reliably distinguish candidate actions. For tasks requiring longer term planning, higher depths can be beneficial, though making the depth too high could expose the search to additional world-model prediction error and uncertainty. Consequently, a moderate search depth provides a favorable trade-off between incorporating future information and maintaining reliable predictions.

The value aggregation discount factor λ \lambda , which controls how strongly future tree values contribute to the recursive value aggregation in Equation 8 is important for performance and setting λ \lambda too large or too small leads to lower performance, and relatively small to moderate values, particularly λ = 0.2 \lambda=0.2 for our settings, provide the strongest learning efficiency. A larger λ \lambda places greater weight on values propagated from deeper imagined states, making the value aggregation more sensitive to accumulated world-model and value-estimation errors, while an overly small λ \lambda limits the contribution of a deeper search. Overall, the results favor large weighting of short-term estimates with a moderately small future-value weighting.

We next examine the effect of the number of action candidates N N considered at each node of the search tree. Increasing N N allows the search to evaluate a broader set of possible actions before committing to the one that maximizes the recursive value estimate, which in principle should improve the quality of the selected action by reducing the chance that a good candidate is overlooked. Empirically, we find that the number of actions N N that performs the best is dependent on the environment. In general, performance improves as N N increases from very small values, since too few candidates can fail to include actions that meaningfully diverge from one another, limiting the benefit of search altogether. However, the gains from increasing N N diminish beyond a moderate number of candidates, and excessively large N N can even hurt performance, likely because evaluating many candidates amplifies the influence of world-model and value-estimation errors across a larger number of imagined rollouts, while also increasing computational cost with limited additional benefit. We find that a moderate number of action candidates offers the best trade-off between exploration of the action space and robustness to compounding prediction errors, and we adopt this setting as default in our main experiments.

Number of expanded nodes. Lastly, we ablate over the number of expanded nodes in our search heuristic. In Figure 7 , we observe that performance is relatively insensitive to the number of expanded nodes J J , as the heuristic accounts for which paths have the largest value to decide which nodes to keep expanding and which to prune. While this greedy approach does not guarantee finding the highest value paths of the full tree, it provides a good approximation. This suggests that retaining a small number of high-value branches is already sufficient to capture useful candidate futures.

## 6 Discussion

In this paper, we present QWM , a framework for leveraging world models via test-time search over actions on top of Q-learning to improve performance. Rather than using the world model to optimize the policy directly, as in conventional model-based RL, we use it purely at test time both during online rollouts and at evaluation. We show that QWM outperforms strong model-free baselines, while avoiding the compounding model bias that afflicts methods which optimize policies primarily inside a learned world model. Despite these results, QWM has limitations. The tree search introduces non-trivial computational overhead at rollout and evaluation time relative to a standard policy forward pass, which may be prohibitive under tight latency requirements. Additionally, QWM depends on learning a world model, which is an expensive and often challenging to train; while grounding the policy and Q-function in real data prevents bias from the model from compounding, reducing search overhead and obtaining high quality world models remain important future directions.

## 7 Acknowledgments

This work was supported in part by an NSF CAREER award, NSF #1941722 the RAI Institute, ONR grant N00014-22-1-2293, and ONR grant N00014-22-1-2621. This work used the Delta system at the National Center for Supercomputing Applications [award OAC 2005572] through allocation CIS260152 from the Advanced Cyberinfrastructure Coordination Ecosystem: Services & Support (ACCESS) program, which is supported by U.S. National Science Foundation grants #2138259, #2138286, #2138307, #2137603, and #2138296.

## Appendix A Author Contributions

PD conceived the idea, devised the algorithm, and led the project. PD implemented the initial prototype, ran experiments, provided hands-on guidance on all experiments, and analyzed and interpreted the results. PD also wrote and positioned the paper. YJ led maintenance of the research codebase, conducted experiments, and produced the figures and videos. CF and DS contributed to the research design, advised the project, edited and positioned the paper.

## Appendix B Additional Experiments

To isolate the benefit of performing test-time search with a Q Q -function, we compare against a variant that learns only a state value function V V and searches with it. Here we fit V V , form an advantage estimate, and optimize the policy by advantage-weighted regression ( Peng et al., 2019 ) . We evaluate this variant both as the edit policy within EXPO and as a standalone method. Results are reported in Figure 8 and Figure 9 . Searching with V V substantially underperforms searching on top of Q-learning across all settings. We attribute this to two factors: the Q-function supplies an additional action-conditioned estimator that V V alone cannot provide, which yields a stronger estimator overall, and the underlying Q Q -learning base algorithm is itself stronger. These results motivate leveraging world models for Q-learning rather than performing search with V V as typically done with model-based RL.

## Appendix C Experiment Details

### C.1 Hyperparameters

For the state-based Robomimic experiments, QWM and all model-free baselines follow the same online RL protocol, including identical demonstration sources, online interaction budgets, and optimization settings. Unless otherwise specified, all methods start online training after 5,000 environment steps, use an offline data ratio of 0.5, do not perform additional offline pretraining, and use an update-to-data (UTD) ratio of 20 with batch size 256. For Q-learning-based methods, we use a critic TD discount γ = 0.99 \gamma=0.99 and a target network update coefficient τ = 0.005 \tau=0.005 . Importantly, γ \gamma is used only in the Bellman TD backup of the critic and is distinct from the tree-search discount λ \lambda , which controls the contribution of deeper future values during search. For EfficientZero V2, we follow the official training protocol provided by the authors and use its recommended UTD ratio of 0.4.

#### QWM -specific hyperparameters.

We instantiate QWM on top of two base RL algorithms, EXPO and RLPD. For both variants, the original policy and critic optimization procedures are kept unchanged, while world-model-guided tree search is additionally used for action selection during online interaction and evaluation.

At each state node, the policy proposes N N candidate actions, and each action is evaluated through K K world-model predictions. To control the size of the search tree, we retain at most J J partial paths during intermediate expansion, following the pruning procedure described in Section 4.4 . We use max aggregation for both intermediate-node values and leaf values. Specifically, intermediate-node aggregation combines the values of candidate actions and their predicted future states in the recursive value computation in Equation 8 , while leaf aggregation combines the N leaf N_{\text{leaf}} Q-values sampled at the final search depth in Equation 7 . The same max aggregation is used when combining world-model branches in the root action score. The tree-search settings are summarized in Table 2 . Here, λ \lambda is used only within the tree-search value recursion and is independent of the critic TD discount γ \gamma . We note that this λ \lambda removes the factor of 1 2 \frac{1}{2} in Equation 9 and Equation 8 , and to exactly match the equations the λ \lambda in the table should be multiplied by 1 2 \frac{1}{2} .

### C.2 Environments

We evaluate QWM on two widely used robotic manipulation benchmarks: Robomimic ( Mandlekar et al., 2021 ) and LIBERO ( Liu et al., 2023 ) . These benchmarks provide diverse manipulation tasks with human-collected demonstrations and are commonly used to evaluate offline-to-online robot learning algorithms.

Robomimic. For state-based experiments, we evaluate QWM on four manipulation tasks from Robomimic: Lift, Can, Square, and Tool Hang. These tasks cover diverse manipulation capabilities, including object grasping and lifting, object relocation, precise insertion, and long-horizon assembly. We use the low-dimensional state observations provided by Robomimic, consisting of robot proprioception and object-related features, including end-effector pose, gripper states, and object states. The action space is a 7-DoF operational space control (OSC) command.

Following prior work, we use different demonstration settings for different tasks. Lift uses a 10-episode subset of the original dataset to create a more challenging evaluation setting, Can uses the multi-human (MH) dataset, and Square and Tool Hang use the standard proficient-human (PH) split. The PH dataset contains demonstrations collected by a single proficient teleoperator, while the MH dataset contains demonstrations collected by multiple teleoperators with varying levels of expertise. All experiments follow the standard Robomimic evaluation protocol and report online success rates.

LIBERO. For pixel-based experiments, we evaluate QWM on five tasks from the LIBERO benchmark. LIBERO is a language-conditioned lifelong robot learning benchmark designed to evaluate knowledge transfer and generalization in robotic manipulation. It contains diverse manipulation tasks with different object configurations, spatial arrangements, and task semantics.

We follow the task selection and evaluation protocol used in prior work and evaluate on five representative tasks (Task 60, Task 79, Task 29, Task 28, and Task 2). The agent receives RGB observations from the robot cameras and generates actions conditioned on visual observations and task instructions. We use the standard LIBERO simulation setup and report online success rates averaged over evaluation episodes.

### C.3 Additional World Model Details and Prediction Quality

We provide additional training and inference details for the world models used in our experiments. For the state-based experiments, the dynamics model is pretrained offline on demonstration transitions ( s , a , s ′ ) (s,a,s^{\prime}) using an MSE objective with Adam, a learning rate of 3 × 10 − 4 3\times 10^{-4} , batch size 256, and 100k gradient steps, and is kept fixed throughout subsequent policy and critic optimization. For the vision-based experiments, we fine-tune the action-conditioned Wan2.2-TI2V-5B world model on demonstration videos paired with aligned raw action sequences. Before training, video clips are encoded once by the frozen Wan2.2 VAE and text prompts are encoded by the frozen umT5-XXL encoder, allowing both representations to be cached and removed from the training loop. We then jointly fine-tune the diffusion transformer and action encoder using the standard Wan2.2 flow-matching objective, while keeping the VAE and text encoder frozen. Training uses AdamW with a learning rate of 1 × 10 − 5 1\times 10^{-5} , a per-GPU batch size of 1 on 8 GPUs, bfloat16 precision, gradient checkpointing, and 200k gradient steps. For LIBERO, we train on 5-frame 128 × 128 128\times 128 clips and model the agent-view and wrist-camera observations as two video streams. During tree search, however, each world-model query advances the rollout by only one step: we use the second frame of the generated clip as the predicted next observation, and then condition a new generation on this prediction for subsequent tree expansion. At inference time, we use 1 denoising step. To examine the next-step generation quality of the resulting world model, Figure 10 compares the generated next observation with the corresponding ground-truth next observation conditioned on the same current observation and action. The predictions capture the task-relevant changes in robot configuration and object motion at the next step, supporting their use as predicted future observations during tree search.

### C.4 Comparisons

We evaluate QWM against both state-of-the-art model-free algorithms and model-based algorithms, as while QWM uses a learned model, it does not use data from the model for training and is trained on top of standard Q-learning.

Model-free baselines. The first set of comparisons is against state-of-the-art model-free algorithms. As QWM learns with only data in the real world instead of from the model, we compare against other such methods to evaluate effectiveness.

EXPO ( Dong et al., 2025b ) . EXPO jointly learns an expressive base policy and a lightweight edit policy, where the edit policy transforms actions sampled from the base policy toward higher-value distributions.

IDQL ( Hansen-Estruch et al., 2023 ) . IDQL trains a diffusion policy through behavior modeling and performs implicit policy extraction using best-of- N N sampling, selecting the candidate action with the highest Q-value.