Learning Abstract World Models with a Group-Structured Latent Space

Technical Summary: Learning Abstract World Models with a Group-Structured Latent Space

Problem Statement

Learning meaningful abstract models of Markov Decision Processes (MDPs) is essential for improving generalization from limited data in reinforcement learning (RL). While existing world model learning methods successfully embed high-dimensional data into lower-dimensional abstract spaces to improve sample efficiency, state-of-the-art approaches often require large training datasets to achieve superhuman performance. A key limitation is the lack of integration of known geometric structures (symmetries) inherent in the environment's dynamics. Without such priors, agents may struggle to generalize to unseen scenarios or require excessive data to learn simple symmetries, such as rotational or translational invariances.

Methodology

The authors propose a framework that imposes geometric priors on the low-dimensional representation manifold of a learned transition model. The core methodology involves:

1. Group-Structured Latent Space: Instead of treating the abstract state space $Z$ as a standard Euclidean vector space, the authors model it as a quotient space (e.g., $\mathbb{R}/k\mathbb{Z}$) or a product of such spaces. This choice reflects the underlying symmetry groups of the environment (e.g., cyclic groups for rotations).
2. Group Actions as Transition Dynamics: The transition model $\tau: Z \times A \to Z$ is constrained to operate via a group action. Specifically, the next latent state is predicted as the current state combined with a learned transition vector via a group operation $\oplus$:
   $$\hat{z}_{t+1} = z_t \oplus \Delta(z_t, a; \theta_{trans})$$
   For quotient spaces like $\mathbb{R}/k\mathbb{Z}$, this operation is implemented via modular arithmetic, naturally encoding periodic boundary conditions.
3. Disentangled Representations: To handle environments containing both symmetric (structured) and non-symmetric (unstructured) features, the framework promotes disentanglement. A sparsity regularization is applied to the transition vector $\Delta$, forcing features that should remain invariant under a specific action to be zero. This is achieved by masking specific coordinates of the transition vector based on a mapping $\sigma(a)$ that defines which latent dimensions are affected by action $a$.
4. Learning Objective: The model is trained using a self-supervised contrastive learning objective (InfoNCE loss) to align predicted transitions with ground truth, combined with a reward prediction loss and a volume regularization loss to keep the abstract space compact and interpretable. Unlike previous approaches that enforce equivariance through weight-tied network architectures (which can be computationally expensive), this method enforces structure directly within the latent space geometry.

Key Contributions

• Geometric Priors in Latent Space: The paper demonstrates how to incorporate known symmetric structures (e.g., rotational and translational symmetries) by selecting appropriate latent spaces (quotient spaces) and corresponding group actions, rather than modifying the neural network architecture.
• Handling Mixed Symmetries: The framework successfully learns representations that combine structured (symmetric) and unstructured features, allowing the model to scale to complex environments like first-person 3D views where not all dynamics are symmetric.
• Improved Generalization: The approach leads to simpler, more disentangled representations that generalize better to unseen transitions and require fewer training samples compared to fully unstructured baselines.
• Downstream RL Performance: The learned abstract models serve as effective priors for downstream RL tasks, significantly improving sample efficiency in low-data regimes.

Experimental Results

The authors evaluated their approach on several environments:

• Toy MDPs (Passage & Torus): In a simple cyclic grid world and a torus MDP (product of two cyclic groups), the group-structured model accurately predicted unseen transitions where unstructured baselines failed. The model recovered the correct topology (e.g., mapping the latent space to a torus in $\mathbb{R}^3$) using only 2 latent dimensions.
• MiniGrid (Top-Down View): The method successfully disentangled orientation (modeled as a cyclic group) from position (modeled as Euclidean space), resulting in a product space representation $\mathbb{R}/2\pi\mathbb{Z} \times \mathbb{R}^2$.
• VizDoom (First-Person 3D View): In a high-dimensional environment with RGB inputs, the model mapped complex visual data into a 3-dimensional latent space. It correctly captured the cyclic nature of rotation and the translational nature of movement, outperforming baselines in generalization metrics (Hits@k, Mean Reciprocal Rank) and downstream RL tasks.
• Quantitative Metrics: Across all environments, the "AWM + Geometric Priors" model significantly outperformed unstructured baselines (AWM without priors) and other methods (PRAE, Rotation Matrix) in generalization to unseen transitions, particularly in low-data regimes. In downstream RL tasks (Torus, MiniGrid, VizDoom), agents using the geometric prior achieved higher cumulative rewards with fewer training steps.

Significance and Claims

The paper claims that by leveraging geometric priors through the choice of latent space and group actions, it is possible to learn more sample-efficient and generalizable world models without the computational overhead of equivariant neural networks. The authors emphasize that their method allows for the flexible combination of symmetric and non-symmetric features, making it applicable to complex, real-world scenarios like 3D first-person games. The work suggests that respecting the underlying symmetry of the environment in the abstract representation leads to "simpler and more disentangled representations," which in turn facilitates better prediction and planning. The authors provide full code to ensure reproducibility.