Learning Abstract World Models with a Group-Structured Latent Space

2025-06-02

AI Summary

To address the weak generalization of abstract world models under few-shot learning conditions, this paper proposes a geometrically grounded abstract MDP modeling framework. Our method explicitly embeds group actions—such as rotations and translations—into the latent space and leverages group representation theory to design symmetry-constrained variational autoencoders that jointly encode structured geometric priors and unstructured dynamics. This work is the first to incorporate explicit group actions into the abstract MDP formalism, thereby enhancing representational disentanglement and downstream reinforcement learning (RL) performance. Evaluated in 3D first-person environments, our approach achieves higher state-transition prediction accuracy than unstructured baselines, improves RL sample efficiency by 37%, and yields more compact, semantically interpretable latent representations.

Technology Category

• Machine Learning: Representation Learning
• Reasoning under Uncertainty: Uncertainty Representations
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

Abstract

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.

Problem

• Learning abstract models of MDPs for better generalization
• Imposing geometric priors on low-dimensional representation manifolds
• Incorporating symmetric structures and unstructured information in latent space

Innovation

• Group-structured latent space for MDPs
• Geometric priors on representation manifold
• Combining symmetries with unstructured information