# UniZero: Generalized and Efficient Planning with Scalable Latent World Models

arXiv: https://arxiv.org/abs/2406.10667 · HTML: https://mlsanity.com/paper/2406.10667/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 64 of its 82 references that are in the map.
- First posted on arXiv 15 Jun 2024.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 64)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling](https://mlsanity.com/paper/1412.3555/) (2014)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs](https://mlsanity.com/paper/2310.01801/) (2023)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)

## Built on by (most cited first, 13 of 13)

- [Raw2Drive: Reinforcement Learning with Aligned World Models for End-to-End Autonomous Driving (in CARLA v2)](https://mlsanity.com/paper/2505.16394/) (2025)
- [How to Synthesize Text Data without Model Collapse?](https://mlsanity.com/paper/2412.14689/) (2024)
- [DriveMamba: Task-Centric Scalable State Space Model for Efficient End-to-End Autonomous Driving](https://mlsanity.com/paper/2602.13301/) (2026)
- [First Order Model-Based RL through Decoupled Backpropagation](https://mlsanity.com/paper/2509.00215/) (2025)
- [WARPD: World model Assisted Reactive Policy Diffusion](https://mlsanity.com/paper/2410.14040/) (2024)
- [Security Considerations for Multi-agent Systems](https://mlsanity.com/paper/2603.09002/) (2026)
- [EgoFSD: Ego-Centric Fully Sparse Paradigm with Uncertainty Denoising and Iterative Refinement for Efficient End-to-End Self-Driving](https://mlsanity.com/paper/2409.09777/) (2024)
- [One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning](https://mlsanity.com/paper/2509.07945/) (2025)
- [TransZero: Parallel Tree Expansion in MuZero using Transformer Networks](https://mlsanity.com/paper/2509.11233/) (2025)
- [Action from Adjacent Set in Physical Space Outperforms the Best Prediction in World Models](https://mlsanity.com/paper/2607.23602/) (2026)
- [Coupled Local and Global World Models for Efficient First Order RL](https://mlsanity.com/paper/2602.06219/) (2026)
- [PriorZero: Bridging Language Priors and World Models for Decision Making](https://mlsanity.com/paper/2605.12289/) (2026)
- [Learning to Focus: Prioritizing Informative Histories with Structured Attention Mechanisms in Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2511.06946/) (2025)

## Abstract (arXiv)

Learning predictive world models is crucial for enhancing the planning capabilities of reinforcement learning (RL) agents. Recently, MuZero-style algorithms, leveraging the value equivalence principle and Monte Carlo Tree Search (MCTS), have achieved superhuman performance in various domains. However, these methods struggle to scale in heterogeneous scenarios with diverse dependencies and task variability. To overcome these limitations, we introduce UniZero, a novel approach that employs a modular transformer-based world model to effectively learn a shared latent space. By concurrently predicting latent dynamics and decision-oriented quantities conditioned on the learned latent history, UniZero enables joint optimization of the long-horizon world model and policy, facilitating broader and more efficient planning in the latent space. We show that UniZero significantly outperforms existing baselines in benchmarks that require long-term memory. Additionally, UniZero demonstrates superior scalability in multitask learning experiments conducted on Atari benchmarks. In standard single-task RL settings, such as Atari and DMControl, UniZero matches or even surpasses the performance of current state-of-the-art methods. Finally, extensive ablation studies and visual analyses validate the effectiveness and scalability of UniZero's design choices. Our code is available at textcolormagentahttps://github.com/opendilab/LightZero.