Q-Learning With World Models

This presentation explains how Q-Learning With World Models (QWM) combines the stability of model-free Q-learning with the foresight of model-based planning. Rather than using learned dynamics models to generate synthetic training data—which risks compounding model errors—QWM uses the world model exclusively at decision time to evaluate and rank candidate actions by simulating their short-term consequences. The approach improves both sample efficiency and final performance in robotic manipulation tasks while keeping policy and critic training grounded entirely in real environment transitions.
Script
Most reinforcement learning agents face a painful trade-off: learn from real experience and stay safe but slow, or learn from imagined experience and risk catastrophic model errors. Q-Learning With World Models refuses this choice entirely.
The key insight is architectural. Instead of feeding model predictions into the learning pipeline where errors compound through backpropagation, QWM uses the world model only at decision time to search over futures proposed by an already-trained Q-function.
The search mechanism builds an alternating state-action tree. At each state, the policy samples candidate actions; for each action, the world model generates possible next states; actions are sampled again at those predicted states, and the process repeats to a fixed depth. A beam search retains only the most promising branches according to accumulated Q-values, making the procedure tractable while biasing exploration toward high-value futures.
Across Robomimic and LIBERO manipulation benchmarks, QWM built on top of EXPO and RLPD consistently outperforms its base algorithms and surpasses model-free methods like IDQL and QSM. More striking still, it substantially outperforms model-based baselines TD-MPC2 and EfficientZero V2, which achieve near-zero success on harder tasks within the same training horizon.
The ablations reveal where the method's power comes from. Search depth of two provides the strongest gains: depth one captures too little, while greater depth exposes the agent to compounding model and critic errors. Critically, variants using only a state-value function instead of action-conditioned Q-values substantially underperform, directly challenging the conventional model-based practice of MCTS-style planning with state values.
QWM demonstrates that you can have the stability of model-free learning and the foresight of planning in a single framework, but the unresolved question is computational cost. The method requires multiple model queries and Q-evaluations per action, a trade-off especially visible in pixel-based tasks where evaluation-time search becomes prohibitively expensive. If you want to explore how test-time search can amplify your own Q-learning agents, visit EmergentMind.com to dive deeper into this work and create your own research videos.