# Learning Abstract World Models with a Group-Structured Latent Space

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

## Key facts

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

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

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

- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere](https://mlsanity.com/paper/2005.10242/) (2020)
- [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)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Dynamic Filter Networks](https://mlsanity.com/paper/1605.09673/) (2016)
- [Towards a Definition of Disentangled Representations](https://mlsanity.com/paper/1812.02230/) (2018)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [DeepMDP: Learning Continuous Latent Space Models for Representation Learning](https://mlsanity.com/paper/1906.02736/) (2019)
- [ViZDoom: A Doom-based AI Research Platform for Visual Reinforcement Learning](https://mlsanity.com/paper/1605.02097/) (2016)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [Equivariance with Learned Canonicalization Functions](https://mlsanity.com/paper/2211.06489/) (2022)
- [Residual Pathway Priors for Soft Equivariance Constraints](https://mlsanity.com/paper/2112.01388/) (2021)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Equivariant Q Learning in Spatial Action Spaces](https://mlsanity.com/paper/2110.15443/) (2021)
- [Plannable Approximations to MDP Homomorphisms: Equivariance under Actions](https://mlsanity.com/paper/2002.11963/) (2020)
- [Symmetry-Aware Actor-Critic for 3D Molecular Design](https://mlsanity.com/paper/2011.12747/) (2020)
- [Learning Symmetric Embeddings for Equivariant World Models](https://mlsanity.com/paper/2204.11371/) (2022)
- [Sample Efficient Grasp Learning Using Equivariant Models](https://mlsanity.com/paper/2202.09468/) (2022)
- [On-Robot Learning With Equivariant Models](https://mlsanity.com/paper/2203.04923/) (2022)
- [Combined Reinforcement Learning via Abstract Representations](https://mlsanity.com/paper/1809.04506/) (2018)
- [Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products](https://mlsanity.com/paper/2401.10216/) (2024)
- [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)

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

- [Flow Equivariant World Models: Memory for Partially Observed Dynamic Environments](https://mlsanity.com/paper/2601.01075/) (2026)

## Abstract (arXiv)

Learning meaningful abstract models of Markov Decision Processes (MDPs) is crucial for improving generalization from limited data. In this work, we show how geometric priors can be imposed on the low-dimensional representation manifold of a learned transition model. We incorporate known symmetric structures via appropriate choices of the latent space and the associated group actions, which encode prior knowledge about invariances in the environment. In addition, our framework allows the embedding of additional unstructured information alongside these symmetries. We show experimentally that this leads to better predictions of the latent transition model than fully unstructured approaches, as well as better learning on downstream RL tasks, in environments with rotational and translational features, including in first-person views of 3D environments. Additionally, our experiments show that this leads to simpler and more disentangled representations. The full code is available on GitHub to ensure reproducibility.