# Equivariant MuZero

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 23 of its 36 references that are in the map.
- First posted on arXiv 9 Feb 2023.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

## Builds on (most cited first, 23 of 23)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [General E(2)-Equivariant Steerable CNNs](https://mlsanity.com/paper/1911.08251/) (2019)
- [Steerable CNNs](https://mlsanity.com/paper/1612.08498/) (2016)
- [How to Train Your Robot with Deep Reinforcement Learning; Lessons We've Learned](https://mlsanity.com/paper/2102.02915/) (2021)
- [Learning to Plan Chemical Syntheses](https://mlsanity.com/paper/1708.04202/) (2017)
- [Value Iteration Networks](https://mlsanity.com/paper/1602.02867/) (2016)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [Gauge Equivariant Mesh CNNs: Anisotropic convolutions on geometric graphs](https://mlsanity.com/paper/2003.05425/) (2020)
- [SO(2)-Equivariant Reinforcement Learning](https://mlsanity.com/paper/2203.04439/) (2022)
- [An investigation of model-free planning](https://mlsanity.com/paper/1901.03559/) (2019)
- [On the role of planning in model-based deep reinforcement learning](https://mlsanity.com/paper/2011.04021/) (2020)
- [Learning Symmetric Embeddings for Equivariant World Models](https://mlsanity.com/paper/2204.11371/) (2022)
- [MuZero with Self-competition for Rate Control in VP9 Video Compression](https://mlsanity.com/paper/2202.06626/) (2022)
- [Continuous MDP Homomorphisms and Homomorphic Policy Gradient](https://mlsanity.com/paper/2209.07364/) (2022)
- [Multi-Agent MDP Homomorphic Networks](https://mlsanity.com/paper/2110.04495/) (2021)
- [Equivariant Networks for Zero-Shot Coordination](https://mlsanity.com/paper/2210.12124/) (2022)

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

- [EDGI: Equivariant Diffusion for Planning with Embodied Agents](https://mlsanity.com/paper/2303.12410/) (2023)
- [The Larger They Are, the Harder They Fail: Language Models do not Recognize Identifier Swaps in Python](https://mlsanity.com/paper/2305.15507/) (2023)
- [Symmetric Replay Training: Enhancing Sample Efficiency in Deep Reinforcement Learning for Combinatorial Optimization](https://mlsanity.com/paper/2306.01276/) (2023)

## Abstract (arXiv)

Deep reinforcement learning repeatedly succeeds in closed, well-defined domains such as games (Chess, Go, StarCraft). The next frontier is real-world scenarios, where setups are numerous and varied. For this, agents need to learn the underlying rules governing the environment, so as to robustly generalise to conditions that differ from those they were trained on. Model-based reinforcement learning algorithms, such as the highly successful MuZero, aim to accomplish this by learning a world model. However, leveraging a world model has not consistently shown greater generalisation capabilities compared to model-free alternatives. In this work, we propose improving the data efficiency and generalisation capabilities of MuZero by explicitly incorporating the symmetries of the environment in its world-model architecture. We prove that, so long as the neural networks used by MuZero are equivariant to a particular symmetry group acting on the environment, the entirety of MuZero's action-selection algorithm will also be equivariant to that group. We evaluate Equivariant MuZero on procedurally-generated MiniPacman and on Chaser from the ProcGen suite: training on a set of mazes, and then testing on unseen rotated versions, demonstrating the benefits of equivariance. Further, we verify that our performance improvements hold even when only some of the components of Equivariant MuZero obey strict equivariance, which highlights the robustness of our construction.