QWM: Q-Learning with World Models

Overview

World models are emerging as one of the most promising directions in physical AI. At the same time, reinforcement learning (RL) fine-tuning is pushing frontier models beyond what pretraining alone can achieve. We propose QWM, a framework that uses a world model to perform test-time scaling on top of Q-learning.

Our key idea: use the world model to perform test-time scaling on top of Q-learning.

At each step, the policy proposes actions.
The world model imagines what could happen after each one.
The value function evaluates those futures.
The agent executes the highest-value action.

We use this search in two places. During online RL, QWM selects better actions, helping the agent collect higher-quality experience for learning. During evaluation, QWM improves execution directly by searching over possible futures before committing to an action.

Instead of asking “Which action looks best now?”, QWM asks: “Which action leads to the best future?”

Unlike prior model-based RL approaches, QWM doesn’t learn from imagined trajectories. This keeps world-model errors out of the learning loop, using imagination only to choose better actions from real experience.

Imagine

Use the world model to imagine what could happen after each candidate action.

Evaluate

Use the Q-function to score those imagined futures and pick the highest-value action.

Learn

Use the real world to learn — the policy and value function train only on real experience.

Method

QWM augments standard Q-learning with world-model-guided test-time search. At each decision step, the current state forms the root of a search tree: the policy proposes candidate actions, the world model predicts their possible future states, and this process is recursively expanded over a short horizon. The learned Q-function is used to evaluate candidate state-action pairs and guide the search toward high-value futures.

Construct a search tree: sample candidate actions from the policy at each state and use the learned world model to predict the resulting future states, recursively expanding the tree to a fixed depth.
Search and aggregate with the Q-function: evaluate state-action pairs throughout the imagined tree with the learned Q-function, retain promising branches, and aggregate intermediate and leaf values to produce a tree-search score for each root action.
Select actions without training on imagined data: use the resulting tree-search scores for action selection during both online data collection and evaluation, while training the policy and critic only on real environment transitions.

QWM tree search and value aggregation details — **Figure 2: Tree search and value aggregation in QWM.** From the current state, the policy proposes candidate actions and the world model predicts their future states. QWM recursively expands promising branches, evaluates imagined state-action pairs with the learned Q-function, and aggregates values from intermediate and leaf nodes to score and select the root action.

Results

QWM is evaluated across state-based Robomimic and pixel-based LIBERO settings, including comparisons to model-free baselines, model-based baselines, and the base methods it augments.

QWM compared with model-free baselines on Robomimic — **Figure 3: Compared with model-free baselines.** Success rates of QWM and model-free baselines in the online setting. QWM outperforms strong model-free baselines in sample efficiency.

QWM compared with TD-MPC2 and EfficientZero V2 — **Figure 4: Compared with model-based RL methods.** Success rates of QWM, TD-MPC2, and EZ-V2 across Lift, Can, Square, and Tool Hang under sparse and dense reward settings.

QWM compared with EXPO base method — **Figure 5: Effect of world-model tree search.** Online success rates of QWM and its EXPO base method on Tool Hang, Square, and Can. QWM improves sample efficiency across all three tasks, helping the agent collect more high-quality experience for learning.

RLPD tree-search comparison — **Figure 6: RLPD tree search.** QWM can also be instantiated on top of any RL algorithm with a Q-function. We show an example with RLPD, a sample-efficient RL algorithm that uses Gaussian policies. QWM improves the sample efficiency of RLPD even more across all three tasks.

QWM pixel-based LIBERO evaluation — **Figure 7: Pixel-based evaluation on LIBERO.** Online success rates of QWM and EXPO across five LIBERO tasks. QWM achieves clear gains on Tasks 60 and 79, improves learning speed on Task 28, and reaches comparable or stronger final performance across the remaining tasks.

Visualization

We compare action-conditioned next-step generation from the learned world model against ground-truth LIBERO trajectories. Each generated next step is stitched together into a full video, with ground truth on the left and the corresponding QWM prediction on the right.