Q-Learning With World Models — part 2/4

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

In this section, we present the key component of QWM to leverage world models to perform test-time search for better performance for Q-Learning. We instantiate QWM as a tree search over actions on top of standard Q-learning We first describe the approach for tree search with world models, split into three parts: constructing the search tree ( Section 4.1 ), aggregating the value of expanded nodes into a value for each proposed action ( Section 4.2 ), and searching with the Q-function ( Section 4.3 ); we then discuss implementation details ( Section 4.4 ).

### 4.1 Tree Search With the World Model

Our goal is to construct a tree such that searching over it yields good actions and adequate coverage of possible futures, and is general enough to be compatible with different problem settings. We construct the tree governed by three quantities: the number of actions sampled at each state N N , the number of next states sampled per action K K , and the search depth D D .

Concretely, at each decision step, we set the root of the tree to the current state s 0 s_{0} where depth 0 0 represents the root of the tree. Nodes alternate between state layers and action layers, with one state layer followed by one action layer constituting a single level of depth. At a state node at level d d , the tree is expanded by drawing N N candidate actions from the policy, { a d n } n = 1 N ∼ π θ ( ⋅ ∣ s d ) \{a_{d}^{n}\}_{n=1}^{N}\sim\pi_{\theta}(\cdot\mid s_{d}) . At each action node, the tree is expanded and form depth d + 1 d+1 by querying the world model K K times to obtain K K predicted future states, { s d + 1 k } k = 1 K ∼ M ψ ​ ( s d , a d n ) \{s^{k}_{d+1}\}_{k=1}^{K}\sim M_{\psi}(s_{d},a_{d}^{n}) .

We repeat this expansion recursively to depth D D ( Figure 2 ). At the final layer, N leaf N_{\text{leaf}} actions are sampled from the policy. The tree is intended to give a short-horizon prediction of the consequences of an action, not to plan exhaustively – this keeps rollout lengths short enough to prevent the world model’s own prediction error to accumulate.

### 4.2 Aggregating the Value

Rather than scoring only the leaves and propagating value to the initial action, we observe that the value of every state-action pair in the tree is useful for action selection. The value of a node can either be expressed as a function of state-action value Q ϕ Q_{\phi} or a function of state value V ϕ V_{\phi} , each with complementary trade-offs. As such, we express the value of a node as a combination of these two estimators of the same underlying quantity.

State-action value estimator. At depth d d , we can score each of the N N sampled actions a d n ∼ π θ ( ⋅ ∣ s d ) a_{d}^{n}\sim\pi_{\theta}(\cdot\mid s_{d}) directly with the critic, which already estimates the return of the full remaining trajectory: V Q ​ ( d ∣ s d ) = agg n ∈ [ N ] Q ϕ ​ ( s d , a d n ) . V_{Q}\big(d\mid s_{d}\big)\;=\;\operatorname*{agg}_{n\in[N]}\;Q_{\phi}\big(s_{d},\,a_{d}^{n}\big). (5) This estimator is low-variance and independent of the remaining search depth and error from the model, but relies entirely on Q ϕ Q_{\phi} being learned accurately.

State value estimator. We can alternatively estimate using the world model’s predicted per-step reward r ψ ​ ( s d , a d n ) r_{\psi}(s_{d},a_{d}^{n}) and a discounted value of the next state: V r ​ ( d ∣ s d ) = agg n ∈ [ N ] [ agg k ∈ [ K ] [ r ψ ​ ( s d , a d n ) + λ ​ V ​ ( d + 1 ∣ s d + 1 n , k ) ] ] , V_{r}\big(d\mid s_{d}\big)\;=\;\operatorname*{agg}_{n\in[N]}\left[\;\operatorname*{agg}_{k\in[K]}\Big[\,r_{\psi}\big(s_{d},a_{d}^{n}\big)\;+\;\lambda\,V\big(d+1\mid s_{d+1}^{n,k}\big)\Big]\right], (6) where s d + 1 n , k = M ψ ​ ( s d , a d n ) s_{d+1}^{n,k}=M_{\psi}\big(s_{d},a_{d}^{n}\big) and r ψ ​ ( s d , a d n ) r_{\psi}(s_{d},a_{d}^{n}) is the learned reward model. At the final layer, N leaf N_{\text{leaf}} actions are sampled from the policy and there is no rollout left to recurse into, so the leaf value reduces to state value which can be computed as an aggregate over Q-values for actions in that state, V ⁡ ( D ∣ s D ) = agg n ∈ [ N leaf ] Q ϕ ​ ( s D , a D n ) . V\big(D\mid s_{D}\big)\;=\;\operatorname*{agg}_{n\in[N_{\text{leaf}}]}\;Q_{\phi}\big(s_{D},\,a_{D}^{n}\big). (7) The state value estimator exploits the full imagined rollout but compounds world-model error with depth.

Combined node value. The two estimators have complementary failure modes: V Q V_{Q} is depth-independent but blind to the rollout, while V r V_{r} uses the rollout but accumulates simulation error. We take advantage of both estimators to get a combined node value: V ⁡ ( d ∣ s d ) = 1 2 ​ ( V Q ​ ( d ∣ s d ) + V r ​ ( d ∣ s d ) ) V\big(d\mid s_{d}\big)\;=\;\tfrac{1}{2}\Big(V_{Q}\big(d\mid s_{d}\big)+V_{r}\big(d\mid s_{d}\big)\Big) (8) which can be extended to a weighted average α ​ V Q + ( 1 − α ) ​ V r \alpha V_{Q}+(1-\alpha)V_{r} , where α ∈ [ 0 , 1 ] \alpha\in[0,1] .

#### Root value.

The value of each originally sampled root action a 0 n a_{0}^{n} follows the same computation at d = 0 d=0 : Q t ​ s ​ ( s 0 , a 0 n ) = 1 2 ​ ( Q ϕ ​ ( s 0 , a 0 n ) + [ r ψ ​ ( s 0 , a 0 n ) + agg k ∈ [ K ] λ ​ V ​ ( 1 ∣ s 1 k ) ] ) , s 1 k = M ψ ​ ( s 0 , a 0 n ) . Q_{ts}(s_{0},a_{0}^{n})\;=\;\tfrac{1}{2}\left(Q_{\phi}(s_{0},a_{0}^{n})\;+\;\Big[r_{\psi}(s_{0},a_{0}^{n})+\operatorname*{agg}_{k\in[K]}\lambda\,V\big(1\mid s_{1}^{k}\big)\Big]\right),\qquad s_{1}^{k}=M_{\psi}(s_{0},a_{0}^{n}). (9) The action ultimately selected is a max or softmax over the original sampled actions, weighted by their tree-search values Q t ​ s ​ ( s 0 , a 0 n ) Q_{ts}(s_{0},a_{0}^{n}) . We use this search procedure in two places: to select actions during training, and to act at evaluation time.

### 4.3 Searching with the Q-function.

Because the tree grows very large very quickly with depth, due to computational constraints we in practice design a heuristic to select only a subset of nodes to expand and prune the rest using the Q-function as a search policy. Let J J be the number of paths that will be expanded. Let ℬ d \mathcal{B}_{d} denote the surviving set of J J partial paths after pruning at depth d d (with ℬ 0 = { a 0 n } n = 1 N \mathcal{B}_{0}=\{a_{0}^{n}\}^{N}_{n=1} , the single root path under consideration). Each surviving path b ∈ ℬ d b\in\mathcal{B}_{d} carries an accumulated discounted score Σ ⁡ ( b , d ) = ∑ d ′ = 1 d λ d ′ ​ Q ϕ ​ ( s d ′ b , a d ′ b ) , \Sigma(b,d)\;=\;\sum_{d^{\prime}=1}^{d}\lambda^{d^{\prime}}\,Q_{\phi}\big(s_{d^{\prime}}^{b},a_{d^{\prime}}^{b}\big), (10) where ( s d ′ b , a d ′ b ) \big(s_{d^{\prime}}^{b},a_{d^{\prime}}^{b}\big) is the state-action pair visited by path b b at depth d ′ d^{\prime} .

At depth d + 1 d+1 , every surviving path b ∈ ℬ d b\in\mathcal{B}_{d} is expanded into N N candidate children { a d + 1 b , n } n = 1 N ∼ π θ ( ⋅ ∣ s d + 1 b ) \{a_{d+1}^{b,n}\}_{n=1}^{N}\sim\pi_{\theta}(\cdot\mid s_{d+1}^{b}) , each scored by its prospective cumulative sum: Σ ~ ​ ( b , n , d + 1 ) = Σ ⁡ ( b , d ) + λ d + 1 ​ Q ϕ ​ ( s d + 1 b , n , a d + 1 b , n ) . \tilde{\Sigma}(b,n,d+1)\;=\;\Sigma(b,d)\;+\;\lambda^{d+1}\,Q_{\phi}\big(s_{d+1}^{b,n},a_{d+1}^{b,n}\big). (11) The surviving set at depth d + 1 d+1 is the top J J over all candidates, ranked by Equation 11 : ℬ d + 1 = top ​ - J { Σ ~ ( b , n , d + 1 ) : b ∈ ℬ d , n ∈ [ N ] } . \mathcal{B}_{d+1}\;=\;\operatorname*{top\text{-}}J\Big\{\tilde{\Sigma}(b,n,d+1)\;:\;b\in\mathcal{B}_{d},\;n\in[N]\Big\}. (12) At depth D D , each surviving path b ∈ ℬ D − 1 b\in\mathcal{B}_{D-1} samples N leaf N_{\text{leaf}} final actions, whose values are aggregated per-beam: L ⁡ ( b ) = agg n ∈ [ N leaf ] Q ϕ ​ ( s D b , a D b , n ) . L(b)\;=\;\operatorname*{agg}_{n\in[N_{\text{leaf}}]}\;Q_{\phi}\big(s_{D}^{b},a_{D}^{b,n}\big). (13) The root’s tree-search value combines all surviving beams. Intuitively, Equation 12 selects survivors by the additive discounted sum Σ ~ \tilde{\Sigma} . When J J equals the max number of nodes, no candidates are discarded at Equation 12 ( ℬ d \mathcal{B}_{d} always contains all N d N^{d} paths). When J J is less than the max number of nodes, Equation 12 discards candidates at every intermediate depth based on Equation 11 .

### 4.4 Practical Implementation

We instantiate the action-conditioned world model in two forms depending on the observation modality. In practice, because we operate in sparse reward settings where rewards are zero except for the terminal state, we do not learn a reward model on top of the world model. For low-dimensional state observations, we use a deterministic residual dynamics model M ψ ​ ( s t , a t ) = s t + Δ ψ ​ ( s t , a t ) M_{\psi}(s_{t},a_{t})=s_{t}+\Delta_{\psi}(s_{t},a_{t}) implemented as a three-layer MLP with hidden dimension 256. The model takes the Robomimic low-dimensional state representation, including end-effector pose, gripper states, and object features, together with the 7-DoF action as input. It is pretrained offline on demonstration transitions using an MSE objective. For pixel observations, we adapt Wan2.2-TI2V-5B ( Wan et al., 2025 ) into an action-conditioned video world model by introducing an MLP-based action encoder. Specifically, each robot action is mapped into action tokens through a three-layer MLP, which are then concatenated with the text conditioning tokens and provided to the diffusion transformer as additional conditioning inputs. The action encoder and diffusion transformer are jointly fine-tuned while keeping the VAE and text encoder frozen. The model is trained on demonstration video clips with aligned action sequences using the standard Wan2.2 flow-matching objective. During inference, the world model predicts short-horizon future observations conditioned on candidate actions. For LIBERO, we generate 5-frame 128 × 128 128\times 128 video clips, and use the predicted next frame as the subsequent observation for iterative rollout and tree search.

## 5 Experiments

The goal of our experiments is to answer the following core questions:

( Q1) How does QWM perform compared to state-of-the-art model-free and model-based RL methods?

( Q2) How does QWM perform compared to the base method it is implemented on top of?

( Q3) Does QWM scale to a pixel-based setting?

( Q4) What components of QWM are most important for performance?

Environments. We evaluate QWM on challenging robotic manipulation tasks from Robomimic ( Mandlekar et al., 2021 ) and LIBERO ( Liu et al., 2023 ) , where a 7-DoF robot arm is required to complete diverse manipulation behaviors under sparse task-completion rewards. For Robomimic, we consider four tasks: Lift, Can, Square, and Tool Hang. Lift requires the robot to grasp and lift an object; Can requires the robot to grasp a cylindrical object and place it at a target location; Square involves precisely inserting a square nut onto a peg; and Tool Hang is a long-horizon multi-stage assembly task where the robot constructs a stand and hangs a tool. For LIBERO, we evaluate on five tasks following the experimental protocol of prior work ( Dong et al., 2026d ) .

Baselines. We compare against both state-of-the-art model-free algorithms and model-based algorithms. We refer to Section C.4 for details on algorithmic comparisons.

### 5.1 How Does QWM Perform Compared to State-of-the-Art Model-Free and Model-Based RL Methods?

Compared with model-free baselines. Success rates of QWM and model-free baselines in the online settings. QWM outperforms strong model-free baselines in sample efficiency.

Model-free. We compare against strong model-free RL baselines in Section 5.1 , all of which leverage online interaction to learn Q-functions. QWM achieves the strongest performance across all tasks. RLPD, DSRL, QSM, QAM, and FQL each sample a single action from the policy for execution in the environment. While this action is intended to approximately maximize the Q-function, optimization delays and errors prevent it from being truly optimal. QWM instead leverages the world model’s ability to predict future outcomes, performing test-time search over candidate actions and selecting the ones with the best predicted future returns. While IDQL selects actions by sampling multiple candidates and selecting the one with the highest Q-value, it considers only the immediate next action rather than a sequence of future actions, the latter carrying substantially more signal about long-term performance. These results demonstrate that the test-time search in QWM provides benefits complementary to Q-learning, enabling better action selection without substituting online interactions.

Model-based. We further compare QWM with state-of-the-art model-based RL methods TD-MPC2 and EfficientZero V2. We include both sparse and dense reward variants for the model-based algorithms, as many of these methods were designed for dense rewards, and struggles to learn as well in sparse reward settings. We present the results in Figure 3 . QWM achieves consistently stronger performance across the evaluated manipulation tasks, while TD-MPC2 and EfficientZero V2 only obtain non-zero success on Lift for the number of training steps reported. Different from conventional model-based RL methods that incorporate model predictions into policy optimization or learning targets, QWM uses the world model only at test-time search while keeping policy and value learning grounded on real environment transitions. This design allows QWM to leverage future predictions for improved action selection without relying on model-generated trajectories during learning, providing an effective way to combine the predictive capability of world models with online Q-learning while avoiding compounding bias from the world model.

### 5.2 How Does QWM Perform Compared to the Base Method It Is Implemented on Top of?

In this section, we evaluate the performance of QWM compared to the base algorithm it is implemented on. We implement QWM on top of EXPO and RLPD for the paper as representative sample efficient Q-learning algorithms, but QWM can be implemented on other model-free algorithms as well. We start by comparing QWM against EXPO in Figure 4 . QWM consistently improves learning efficiency across Tool Hang, Square, and Can, with particularly pronounced gains on harder tasks such as Tool Hang. Comparing with RLPD as the base algorithm, as shown in Figure 6, using QWM with RLPD consistently improves learning compared to RLPD alone with particularly clear gains in sample efficiency. These findings show that QWM provides consistent benefits across different underlying RL algorithms. At the same time, QWM will benefit from a more sample efficient base algorithm as shown by the results of QWM on top of EXPO compared to QWM on top of RLPD.

### 5.3 Does QWM Scale to a Pixel-Based Setting?

We further evaluate whether the performance gains of QWM extend to high-dimensional visual observations. Due to computational constraints, we use the world model only for sampling during online RL data collection, not for evaluation. As shown in Figure 6 , QWM achieves a clear overall improvement over EXPO. The gains are most evident on Tasks 60, 79, and 29, where QWM learns faster and reaches stronger late-stage performance. On Task 28, both methods eventually attain near-perfect success, but QWM reaches high performance earlier, demonstrating a clear improvement in sample efficiency. As shown in the ablations in Figure 4 , restricting the world model to sampling-only (without evaluation) reduces performance, suggesting the visual results can be even stronger given sufficient compute. Nonetheless, even in this sampling-only setting, QWM consistently improves over the base algorithm, showing that its benefits extend naturally to high-dimensional visual inputs, which are substantially harder to learn a dynamics model from.

### 5.4 What Components of QWM Are Most Important for Performance?