mlsanityData to 14 Oct 2026

Learning Abstract World Models with a Group-Structured Latent Space

arXiv 2506.01529 ↗First posted 2 Jun 2025cs.LGcited by 1

On the map: World models → Latent predictive world models

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.

Key facts. Cited by 1 paper 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.

Builds on

34 of 34 references.

Closest ancestors · 3

Same direction, newest first

  1. 2022Learning Symmetric Embeddings for Equivariant World Models24 Apr
  2. 2019Contrastive Learning of Structured World Models27 Nov
  3. 2018Learning Latent Dynamics for Planning from Pixels12 Nov

Other references · 31

  1. 2018Representation Learning with Contrastive Predictive Codingcited by 7,777
  2. 2015Deep Reinforcement Learning with Double Q-learningcited by 1,745
  3. 2020Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hyperspherecited by 1,405
  4. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  5. 2023Mastering Diverse Domains through World Modelscited by 1,253
  6. 2021E(n) Equivariant Graph Neural Networkscited by 965
  7. 2018Recurrent World Models Facilitate Policy Evolutioncited by 785
  8. 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
  9. 2023Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Taskscited by 327
  10. 2016Dynamic Filter Networkscited by 319
  11. 2018Towards a Definition of Disentangled Representationscited by 296
  12. 2021Mastering Atari Games with Limited Datacited by 237
19 more
  1. 2019DeepMDP: Learning Continuous Latent Space Models for Representation Learningcited by 207
  2. 2016ViZDoom: A Doom-based AI Research Platform for Visual Reinforcement Learningcited by 179
  3. 2020MDP Homomorphic Networks: Group Symmetries in Reinforcement Learningcited by 142
  4. 2022Equivariance with Learned Canonicalization Functionscited by 109
  5. 2021Residual Pathway Priors for Soft Equivariance Constraintscited by 75
  6. 2022Causal Dynamics Learning for Task-Independent State Abstractioncited by 65
  7. 2021Equivariant Q Learning in Spatial Action Spacescited by 62
  8. 2020Plannable Approximations to MDP Homomorphisms: Equivariance under Actionscited by 61
  9. 2020Symmetry-Aware Actor-Critic for 3D Molecular Designcited by 55
  10. 2022Sample Efficient Grasp Learning Using Equivariant Modelscited by 53
  11. 2022On-Robot Learning With Equivariant Modelscited by 42
  12. 2018Combined Reinforcement Learning via Abstract Representationscited by 37
  13. 2024Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Productscited by 36
  14. 2022Continuous MDP Homomorphisms and Homomorphic Policy Gradientcited by 27
  15. 2021Multi-Agent MDP Homomorphic Networkscited by 25
  16. 2020Group Equivariant Deep Reinforcement Learningcited by 24
  17. 2024Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inferencecited by 17
  18. 2020Learning Group Structure and Disentangled Representations of Dynamical Environmentscited by 13
  19. 2023rm E(3)-Equivariant Actor-Critic Methods for Cooperative Multi-Agent Reinforcement Learningcited by 6

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2026: 1.

1 paper.

  1. 2026Flow Equivariant World Models: Memory for Partially Observed Dynamic Environmentscited by 8

Similar, not cited

Close in wording, no citation either way

  1. 2026Geometric Priors for Generalizable World Models via Vector Symbolic Architecture
  2. 2026Platonic Representation Hypothesis on World Models
  3. 2021Representation Learning for Online and Offline RL in Low-rank MDPs
  4. 2024Isometric Representation Learning for Disentangled Latent Space of Diffusion Models
  5. 2025A Geometric Theory of Cognition for Machine Intelligence
  6. 2026Learning POMDP World Models from Observations with Language-Model Priors
  7. 2026Subspace-Decomposed JEPAs: Disentangling Progression and Content in Latent World Models
  8. 2025Cloning Deterministic Worlds: The Critical Role of Latent Geometry in Long-Horizon World Models