PaperlayerView exact revision on arXiv

Source-linked AI summary

Q-Learning With World Models

Perry Dong, Yueru Jia, Chelsea Finn, Dorsa Sadigh

arXiv:2608.17163v1cs.LGcs.AI

TL;DR

Online RL needs greater sample efficiency, but conventional model-based approaches can compound world-model bias in complex robotics. QWM adds test-time search over imagined trajectories to standard Q-learning while training on real transitions, and outperforms strong model-free and model-based baselines across manipulation tasks.

  • Problem

    World models have improved sample efficiency mainly in supervised policy learning, leaving open how to use them for online RL without compounding bias in complex robotics.

  • Method

    QWM uses a learned world model for test-time search over imagined trajectories, selecting high-value actions on top of a Q-function during online rollouts and evaluation.

  • Results

    QWM significantly outperforms strong model-free and model-based baselines across Robomimic and LIBERO in both sample efficiency and performance.

  • Takeaways & Limitations

    The results indicate that world models can consistently improve Q-learning across underlying algorithms while policy and value learning remain grounded in real interactions.

  • Takeaways & Limitations

    QWM adds substantial rollout and evaluation computation and depends on an expensive, challenging-to-train world model.

Abstract

from arXiv · show

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

QWM uses learned world models for test-time search over imagined trajectories on top of standard Q-learning, improving action selection without training policies or value functions primarily on imagined rollouts. By grounding learning in real environment transitions, it seeks sample-efficiency gains while avoiding compounding model bias in complex robotics settings.

  • Motivation: Off-policy RL’s improved sample efficiency has enabled RL fine-tuning of Vision-Language-Action models into reliable, high-performing policies.Further sample-efficiency gains become increasingly valuable as RL is applied to more complex scenarios.
  • Problem: Prior model-based RL methods optimize policies or value functions inside imperfect world models, causing compounded bias that worsens with task horizon and visual complexity.This makes scaling to large, high-dimensional problems such as real-world robotics difficult.
  • QWM framework: QWM places a learned world model on top of standard Q-learning to search imagined future trajectories and improve action selection at decision time.The framework uses this test-time search rather than directly training the policy or value function on imagined rollouts.
  • Action selection: QWM uses the learned Q-function to evaluate policy-proposed candidate actions and select the highest-value action before execution.World-model tree search aggregates values from intermediate and leaf imagined states to score initial actions.
  • Training and deployment: Test-time search operates during online rollouts and evaluation, while the policy and value function remain trained online on real environment transitions.This preserves predictive-model sample-efficiency benefits while avoiding the bias accumulation associated with primarily imagined-rollout training.

2 RELATED WORK

Related work improves online RL with prior data, learns or plans with environment models, and applies world models to VLA policy learning. Test-time scaling, including best-of-N sampling, has also improved generative models and policy learning across domains.

  • Reinforcement learning with prior data: Prior-data methods improve online RL performance and sample efficiency through exploration–exploitation balancing, value-estimate calibration, or retaining offline data during fine-tuning.The cited approaches use offline datasets to accelerate online learning.
  • Model-based RL: Model-based RL learns environment dynamics and uses them for policy or value learning by generating synthetic transitions or optimizing against imagined trajectories.Some methods mix short branched rollouts with real data, while others plan over imagined rollouts using trajectory optimizers such as cross-entropy or MPPI.
  • World models for VLAs: World-model methods for VLAs alternate real-rollout world-model fine-tuning with synthetic policy-training data, or apply on-policy RL directly to imagined rollouts.WoVR addresses reward hallucination and out-of-distribution drift in the resulting policies.
  • Test-time scaling and best-of-N inference: Test-time scaling, including best-of-N sampling, has improved performance across autoregressive language models, diffusion models, and policy learning.The passage identifies test-time scaling as a substantial body of work spanning these domains.

3 PRELIMINARIES

The paper formulates reinforcement learning as discounted-return maximization in an MDP and studies how learned world models can improve online Q-learning decision-making. QWM is built on off-policy Q-learning, including EXPO and RLPD implementations.

  • MDP formulation: RL seeks a policy that maximizes expected discounted return in an MDP defined by states, actions, rewards, initial-state distribution, discount factor, and transition probabilities.The MDP is specified by {S, A, ρ, r, γ, T}, with γ ∈ [0, 1).
  • World-model-assisted online RL: A transition dataset pretrains the world model, while online environment tuples are appended to replay buffer D for policy and Q-function updates.The dataset may be collected offline or accumulated online.
  • Off-policy Q-learning: The study focuses on off-policy RL, where Q-functions estimate discounted returns and are trained with temporal-difference learning.The central question is how to leverage a world model to improve Q-learning performance.
  • EXPO: EXPO combines a supervised base flow policy with a Q-optimized edit policy, then selects the value-maximizing candidate among N sampled base and edited actions.The edit ˆa_t is clipped to [−β, β] and added to the base action without backpropagating through the base policy.
  • RLPD: RLPD combines offline and online replay data, high update-to-data ratios, an ensemble of critics, and a Gaussian actor for sample-efficient off-policy RL.It uses symmetric sampling with equal proportions from offline and online buffers and stabilizes critic updates with random critic subsets and layer normalization.

4 METHOD

QWM performs test-time tree search over imagined trajectories on top of standard Q-learning, using world-model predictions and value aggregation to select actions during training and evaluation. The method controls search growth through bounded expansion and combines depth-independent Q-values with rollout-based estimates to balance learned-value accuracy against model-error accumulation.

  • 4.1 Tree Search: QWM constructs an action tree from the current state using N sampled actions, K predicted next states per action, and search depth D.The tree alternates state and action layers and recursively expands to depth D.
  • 4.1 Tree Search: Short-horizon imagined search predicts action consequences while limiting rollout length to prevent accumulation of world-model prediction error.The final layer samples N_leaf actions without further recursion.
  • 4.2 Value Aggregation: Node values combine Qϕ-based state-action estimates with world-model rollout-based state values, using a weighted average αVQ + (1 − α)Vr.The Qϕ estimator is low-variance and depth-independent, whereas the rollout-based estimator uses imagined trajectories but compounds model error with depth.
  • 4.3 Search with the Q-function: The selected action is a max or softmax over the originally sampled root actions, weighted by their tree-search values Qts(s0, an0), for both training and evaluation.The same search procedure selects actions during online training and evaluation.
  • 4.3 Search with the Q-function: To control rapid tree growth, QWM retains only J partial paths at each depth and prunes candidates using the Q-function as a search policy.When J is below the maximum number of nodes, candidates are discarded at every intermediate depth according to their cumulative scores.
  • World-model instantiation: For sparse-reward settings, QWM omits a learned reward model and uses a deterministic residual dynamics model for low-dimensional observations.The residual model is a three-layer MLP with hidden dimension 256.

5 EXPERIMENTS

Experiments evaluate QWM on challenging sparse-reward robotic manipulation tasks, comparing it with model-free and model-based RL, its base algorithms, pixel-based settings, and key search components. QWM achieves stronger performance or learning efficiency across these evaluations, with depth 2 and λ = 0.2 among the strongest search settings.

  • Environments: QWM is evaluated on 7-DoF robotic manipulation tasks from Robomimic and LIBERO under sparse task-completion rewards.Robomimic includes Lift, Can, Square, and Tool Hang.
  • Model-free comparisons: QWM achieves the strongest performance across all evaluated model-free tasks, improving on methods that execute a single policy-sampled action.The compared baselines are RLPD, DSRL, QSM, QAM, and FQL.
  • Model-based comparisons: QWM achieves consistently stronger performance than TD-MPC2 and EfficientZero V2 across evaluated manipulation tasks and reward settings.The comparison includes both sparse- and dense-reward variants of the model-based algorithms.
  • Base-algorithm comparisons: QWM consistently improves learning efficiency over EXPO and RLPD, with particularly pronounced gains on harder tasks such as Tool Hang.Tree search improves the sample efficiency of RLPD across Square, Can, and Lift, while QWM built on EXPO remains more sample-efficient overall.
  • Pixel-based evaluation: In pixel-based LIBERO evaluation, QWM improves overall performance over EXPO, with gains most evident on Tasks 60, 79, and 29.On Task 28, both methods eventually attain near-perfect success; evaluation-time world-model search was omitted due to computational constraints.
  • Ablations: Applying search during both online sampling and evaluation provides the strongest and most consistent learning efficiency, while depth 2 and λ = 0.2 are especially effective settings.Performance is relatively insensitive to the number of expanded nodes J, and the best number of action candidates N depends on the environment.

6 DISCUSSION

QWM uses world models for test-time search over actions on top of Q-learning during online rollouts and evaluation, rather than directly optimizing the policy. This approach retains real-data grounding but introduces computational and training costs that remain important challenges.

  • Method: QWM leverages world models through test-time action search layered on Q-learning during online rollouts and evaluation.The world model is used purely at test time rather than to optimize the policy directly.
  • Limitations: Tree search adds non-trivial computational overhead at rollout and evaluation time compared with a standard policy forward pass.This overhead may be prohibitive under tight latency requirements.
  • Limitations: QWM depends on learning a world model, which is expensive and often challenging to train.Reducing search overhead and obtaining high quality world models remain important future directions.
  • Limitations: Grounding the policy and Q-function in real data prevents bias from the model from compounding.This grounding is presented alongside the remaining need to improve search efficiency and world-model quality.

A AUTHOR CONTRIBUTIONS

PD conceived and led the project, developed the algorithm and prototype, conducted and analyzed experiments, and wrote the paper. YJ maintained the codebase and produced figures and videos, while CF and DS advised the research and edited the paper.

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

B ADDITIONAL EXPERIMENTS

QWM’s test-time search benefits specifically from using a Q-function rather than a state value function. Across both EXPO and standalone settings, searching with Q substantially outperforms V-based search.

  • Q versus V search: Searching with a Q-function substantially outperforms searching with a state value function V.The V variant fits V, forms an advantage estimate, and optimizes the policy by advantage-weighted regression.
  • Q versus V search: QWM-EXPO’s online success rates are substantially better than those of the V-search counterpart.The comparison uses an edit policy learned through a state value function V.
  • Q versus V search: QWM-RLPD’s online success rates are substantially better than those of the standalone V-search counterpart.The V counterpart uses a Gaussian policy learned from V in the style of AWR.

C EXPERIMENT DETAILS · C.1 HYPERPARAMETERS

QWM and model-free baselines use a shared online RL protocol for state-based Robomimic, while QWM preserves base algorithm optimization and adds world-model-guided tree search for action selection. The implementation distinguishes critic discounting from search discounting and specifies candidate-action, prediction, pruning, and aggregation mechanisms.

  • C.1 HYPERPARAMETERS: QWM and model-free baselines share demonstrations, online interaction budgets, optimization settings, and a common online RL protocol on state-based Robomimic.Training starts after 5,000 environment steps, with offline data ratio 0.5, no additional offline pretraining, UTD ratio 20, and batch size 256.
  • C.1 HYPERPARAMETERS: Q-learning-based methods use critic TD discount γ = 0.99 and target network update coefficient τ = 0.005.EfficientZero V2 instead follows the authors’ official protocol with recommended UTD ratio 0.4.
  • C.1 HYPERPARAMETERS: The critic discount γ applies only to the Bellman TD backup, whereas tree-search discount λ controls deeper future-value contributions.The two discounts serve distinct roles in QWM’s implementation.
  • C.1 HYPERPARAMETERS: QWM is instantiated with EXPO and RLPD while leaving each algorithm’s original policy and critic optimization procedures unchanged.World-model-guided tree search is additionally used for action selection during online interaction and evaluation.
  • C.1 HYPERPARAMETERS: At each state node, the policy proposes N candidate actions, each action receives K world-model predictions, and at most J partial paths are retained during expansion.The pruning procedure controls search-tree size.
  • C.1 HYPERPARAMETERS: QWM uses max aggregation for both intermediate-node values and leaf values during recursive tree search.Intermediate-node aggregation combines candidate-action values with their predicted future-state values.

C.2 ENVIRONMENTS · C.3 ADDITIONAL WORLD MODEL DETAILS AND PREDICTION QUALITY

QWM is evaluated on state-based Robomimic and pixel-based LIBERO manipulation tasks, using benchmark-specific observations, actions, demonstrations, and evaluation protocols. Its state-based world model is pretrained offline on demonstration transitions with fixed optimization details before policy and critic training.

  • C.2 ENVIRONMENTS: QWM is evaluated on Robomimic and LIBERO, two widely used robotic manipulation benchmarks with diverse tasks and human-collected demonstrations.
  • C.2 ENVIRONMENTS: Robomimic state-based experiments cover Lift, Can, Square, and Tool Hang across grasping, relocation, insertion, and long-horizon assembly capabilities.
  • C.2 ENVIRONMENTS: Robomimic uses low-dimensional state observations containing robot proprioception and object-related features, with actions represented by a 7-DoF operational space control command.
  • C.2 ENVIRONMENTS: Demonstration settings vary by Robomimic task: Lift uses a 10episode subset, Can uses the multi-human dataset, and Square and Tool Hang use the proficient-human split.The proficient-human split comes from one teleoperator, whereas the multi-human dataset comes from multiple teleoperators with varying expertise.
  • C.2 ENVIRONMENTS: LIBERO pixel-based experiments evaluate five tasks in a language-conditioned lifelong learning benchmark designed to test knowledge transfer and manipulation generalization.The benchmark varies object configurations, spatial arrangements, and task semantics.
  • C.2 ENVIRONMENTS: The LIBERO agent receives RGB camera observations and conditions actions on visual observations and task instructions under the standard simulation protocol.Results are reported as online success rates averaged over evaluation episodes.
  • C.3 ADDITIONAL WORLD MODEL DETAILS AND PREDICTION QUALITY: For state-based experiments, the dynamics model is pretrained offline on demonstration transitions (s, a, s′) using an MSE objective with Adam and a learning rate of 3 × 10−4.Training uses batch size 256 and 100k gradient steps, and the model remains fixed during subsequent policy and critic optimization.

C.4 COMPARISONS

QWM is compared with model-free methods trained only on real-world data and model-based RL methods that use learned dynamics models directly for training. The comparison spans several Q-learning, diffusion-policy, and latent-world-model approaches.

  • Comparison categories: QWM is evaluated against state-of-the-art model-free and model-based algorithms while remaining trained on real transitions through standard Q-learning.Unlike model-based baselines, QWM uses its learned model for test-time search rather than training.
  • Model-free baselines: EXPO edits base-policy actions toward higher-value distributions, while IDQL selects the highest-Q candidate from best-of-N diffusion-policy samples.These methods represent alternative model-free approaches for improving action selection with learned value information.
  • Model-free baselines: RLPD improves sample efficiency with replay data, high update-to-data ratios, critic ensembles, and a Gaussian policy, whereas QSM guides diffusion denoising using Q-function gradients.Both methods use model-free reinforcement-learning or policy-optimization mechanisms rather than a learned dynamics model for training.
  • Model-free baselines: DSRL learns reinforcement-learning policies over diffusion noise seeds, while QAM propagates value gradients through diffusion trajectories using adjoint matching.Both methods adapt diffusion-based action generation using objectives derived from the learned Q-function.
  • Model-based baselines: TD-MPC2 plans with short-horizon MPPI in latent space, while EfficientZero V2 uses Sampling-based Gumbel tree search and reuses bootstrapped targets for policy and value training.These model-based baselines use learned world models directly during training, unlike QWM.
Loading 2608.17163v1…