# Vector Quantized Models for Planning

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

## Key facts

- Cited by 48 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 64 references that are in the map.
- First posted on arXiv 8 Jun 2021.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Dota 2 with Large Scale Deep Reinforcement Learning](https://mlsanity.com/paper/1912.06680/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Learning Dexterous In-Hand Manipulation](https://mlsanity.com/paper/1808.00177/) (2018)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [dm_control: Software and Tasks for Continuous Control](https://mlsanity.com/paper/2006.12983/) (2020)
- [Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents](https://mlsanity.com/paper/1709.06009/) (2017)
- [Thinking Fast and Slow with Deep Learning and Tree Search](https://mlsanity.com/paper/1705.08439/) (2017)
- [DeepStack: Expert-Level Artificial Intelligence in No-Limit Poker](https://mlsanity.com/paper/1701.01724/) (2017)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [Action-Conditional Video Prediction using Deep Networks in Atari Games](https://mlsanity.com/paper/1507.08750/) (2015)
- [Deep Reinforcement Learning from Self-Play in Imperfect-Information Games](https://mlsanity.com/paper/1603.01121/) (2016)
- [Model-Based Offline Planning](https://mlsanity.com/paper/2008.05556/) (2020)
- [When to use parametric models in reinforcement learning?](https://mlsanity.com/paper/1906.05243/) (2019)
- [Value Prediction Network](https://mlsanity.com/paper/1707.03497/) (2017)

## Built on by (most cited first, 30 of 48)

- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Foundation Models for Decision Making: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2303.04129/) (2023)
- [CycleResearcher: Improving Automated Research via Automated Review](https://mlsanity.com/paper/2411.00816/) (2024)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Straightening Out the Straight-Through Estimator: Overcoming Optimization Challenges in Vector Quantized Networks](https://mlsanity.com/paper/2305.08842/) (2023)
- [SafeDiffuser: Safe Planning with Diffusion Probabilistic Models](https://mlsanity.com/paper/2306.00148/) (2023)
- [Efficient Planning in a Compact Latent Action Space](https://mlsanity.com/paper/2208.10291/) (2022)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [PlayFusion: Skill Acquisition via Diffusion from Language-Annotated Play](https://mlsanity.com/paper/2312.04549/) (2023)
- [On Transforming Reinforcement Learning by Transformer: The Development Trajectory](https://mlsanity.com/paper/2212.14164/) (2022)
- [Efficient World Models with Context-Aware Tokenization](https://mlsanity.com/paper/2406.19320/) (2024)
- [Procedural Generalization by Planning with Self-Supervised World Models](https://mlsanity.com/paper/2111.01587/) (2021)
- [The Free Energy Principle for Perception and Action: A Deep Learning Perspective](https://mlsanity.com/paper/2207.06415/) (2022)
- [Addressing Optimism Bias in Sequence Modeling for Reinforcement Learning](https://mlsanity.com/paper/2207.10295/) (2022)
- [LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios](https://mlsanity.com/paper/2310.08348/) (2023)
- [Hierarchical Imitation Learning with Vector Quantized Models](https://mlsanity.com/paper/2301.12962/) (2023)
- [BetaZero: Belief-State Planning for Long-Horizon POMDPs using Learned Approximations](https://mlsanity.com/paper/2306.00249/) (2023)
- [Efficient Planning with Latent Diffusion](https://mlsanity.com/paper/2310.00311/) (2023)
- [Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning](https://mlsanity.com/paper/2406.03234/) (2024)
- [Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models](https://mlsanity.com/paper/2405.19878/) (2024)
- [Skill Decision Transformer](https://mlsanity.com/paper/2301.13573/) (2023)
- [Hybrid Search for Efficient Planning with Completeness Guarantees](https://mlsanity.com/paper/2310.12819/) (2023)
- [Efficient Monte Carlo Tree Search via On-the-Fly State-Conditioned Action Abstraction](https://mlsanity.com/paper/2406.00614/) (2024)
- [The Impact of Quantization and Pruning on Deep Reinforcement Learning Models](https://mlsanity.com/paper/2407.04803/) (2024)
- [Mitigating Adversarial Perturbations for Deep Reinforcement Learning via Vector Quantization](https://mlsanity.com/paper/2410.03376/) (2024)
- [Planning Transformer: Long-Horizon Offline Reinforcement Learning with Planning Tokens](https://mlsanity.com/paper/2409.09513/) (2024)
- [Few-Shot Task Learning through Inverse Generative Modeling](https://mlsanity.com/paper/2411.04987/) (2024)
- [Scalable Decision-Making in Stochastic Environments through Learned Temporal Abstraction](https://mlsanity.com/paper/2502.21186/) (2025)
- [GenCO: Generating Diverse Designs with Combinatorial Constraints](https://mlsanity.com/paper/2310.02442/) (2023)

## Abstract (arXiv)

Recent developments in the field of model-based RL have proven successful in a range of environments, especially ones where planning is essential. However, such successes have been limited to deterministic fully-observed environments. We present a new approach that handles stochastic and partially-observable environments. Our key insight is to use discrete autoencoders to capture the multiple possible effects of an action in a stochastic environment. We use a stochastic variant of Monte Carlo tree search to plan over both the agent's actions and the discrete latent variables representing the environment's response. Our approach significantly outperforms an offline version of MuZero on a stochastic interpretation of chess where the opponent is considered part of the environment. We also show that our approach scales to DeepMind Lab, a first-person 3D environment with large visual observations and partial observability.