Algebras of actions in an agent's representations of the world

This paper proposes a framework to extract and classify the algebras of world transformations from an agent's perspective, generalizing key symmetry-based representation learning concepts like equivariance and disentanglement to apply to any transformation algebra rather than just groups.

The Gist

Imagine you are teaching a robot to navigate a maze. The robot has to learn how to move, where walls are, and how to find a treasure. To do this efficiently, the robot needs a "mental map" of the world.

This paper is about how to build the best possible mental map for that robot.

The Old Way: The "Perfect Circle" Rule

For a long time, researchers tried to teach robots using a specific mathematical rule called Symmetry-Based Disentangled Representation Learning (SBDRL).

Think of this like teaching a robot about a perfectly round ball.

• If you roll the ball forward, backward, left, or right, it behaves exactly the same. It's symmetrical.
• If you roll it forward and then backward, you end up exactly where you started. This is called a "reversible" action.
• The old math only worked for worlds that acted like this perfect ball. It assumed every action could be undone and that the rules were the same everywhere.

The Problem: Real life isn't a perfect ball.

• What if the robot eats a cookie? You can't "un-eat" it. That's an irreversible action.
• What if there's a wall? You can't walk through it. The rules change depending on where you are.
• The old math broke down in these situations. It was too rigid, like trying to use a round peg in a square hole.

The New Way: The "Universal Toolkit"

The authors of this paper say: "Let's stop forcing the world to be a perfect ball. Let's build a toolkit that can handle any shape of world."

They propose a new mathematical framework that treats the robot's actions not just as symmetries, but as a language of transformations.

Analogy 1: The LEGO vs. The Clay

• The Old Way (SBDRL) was like trying to build with only LEGO bricks that snap together perfectly in a grid. It works great for simple, repetitive structures, but you can't make a smooth curve or a messy pile of sand.
• The New Way is like having modeling clay. You can mold it into a perfect sphere (the old way), but you can also mold it into a jagged rock, a flowing river, or a broken bridge. It handles the messy, real-world stuff where actions can't always be undone.

Analogy 2: The Map and the Territory

Imagine the robot is an explorer.

• The Old Map only showed roads that looped back on themselves. If the explorer tried to go down a dead end or cross a bridge that collapsed, the map said, "Error! This doesn't exist."
• The New Map is a living document. It records: "If you go here, you hit a wall (dead end)." "If you eat this apple, it disappears forever." It captures the true algebra (the rules of interaction) of the world, whether those rules are neat loops or messy one-way streets.

The Secret Sauce: Category Theory

To make this work, the authors used a branch of math called Category Theory.

Think of Category Theory as the "Grammar of Relationships."

• Instead of looking at the objects themselves (the walls, the cookies, the robot), it looks at how they relate to each other.
• It's like looking at a dance. You don't just study the dancers; you study the steps they take relative to each other.
• This allows the robot to understand that "eating a cookie" and "hitting a wall" are both valid parts of the world's structure, even if they don't fit the old "perfect symmetry" rules.

What Does This Mean for AI?

1. Smarter Learning: Robots can learn faster because they aren't confused by the fact that the world isn't perfect. They can understand that some things change forever (like eating food) and some things are blocked (like walls).
2. Better Generalization: If a robot learns the "grammar" of a specific type of messy world, it can apply that same grammar to a different messy world. It's like learning the rules of English grammar so you can read any book, not just one specific story.
3. Unlearning the "Perfect World" Bias: It frees AI developers from having to force their problems into neat, symmetrical boxes. They can model the world exactly as it is: complex, irreversible, and full of dead ends.

The Bottom Line

This paper is a blueprint for building more human-like mental models for AI.

Just as humans understand that you can't un-break a glass or un-eat a meal, this new framework allows AI to understand that the world is full of one-way streets and dead ends. By using a more flexible mathematical language (Category Theory), the authors have given AI a way to learn efficient, robust representations of the real world, not just the idealized, perfect worlds we used to study.

It's the difference between teaching a robot to drive on a perfect, empty racetrack versus teaching it to drive in a chaotic, rainy city with traffic jams and potholes. The new framework is the training manual for the city.

Technical Summary

1. Problem Statement

The core problem addressed is the limitation of current Symmetry-Based Disentangled Representation Learning (SBDRL) frameworks in Reinforcement Learning (RL).

• Current State: SBDRL (proposed by Higgins et al.) relies on algebraic groups to model world symmetries. It assumes that agent actions form groups, meaning actions must be reversible (have inverses) and the world must be "action-homogeneous" (the same action sequence yields the same result regardless of the starting state).
• The Gap: Many real-world RL scenarios involve irreversible actions (e.g., eating a consumable, breaking a block) or state-dependent constraints (e.g., hitting a wall). These scenarios violate the group axioms (specifically the existence of inverses and closure under all states). Consequently, SBDRL cannot formally represent or learn representations for these common, non-group-based transformations, limiting the robustness and generalization of AI agents.

2. Methodology

The authors propose a general mathematical framework based on Category Theory and Algebraic Structures to describe world transformations without restricting them to groups.

A. Formal Framework (Section 2)

• World Model: Defined as a directed multigraph $\mathcal{W} = (W, \hat{D}, s, t)$ where $W$ are world states and $\hat{D}$ are minimum transitions.
• Agent Actions: Actions are defined as sequences of minimum transitions. The authors introduce a labeling map that associates transitions with actions.
• Equivalence Relation ($\sim$): Two actions $a, a'$ are equivalent if they produce the same resulting state for all initial world states ($a * w = a' * w$). This creates a quotient set of actions $A/\sim$.
• Algebraic Structure: The set of equivalence classes $A/\sim$ forms an algebraic structure under composition. The paper investigates whether this structure is a Group, Monoid, or Small Category depending on the world's properties.

B. Algorithmic Exploration (Section 3 & 4)

The authors developed algorithms to generate State Cayley Tables and Action Cayley Tables to empirically analyze the algebraic structure of specific worlds:

1. Reproducing SBDRL: They showed that for worlds satisfying specific conditions (unrestricted actions and global inverses), $A/\sim$ forms a Group, thereby recovering SBDRL as a special case.
2. Beyond Groups: They applied the framework to worlds with:
   • Irreversible actions: (e.g., consuming an item).
   • Restricted actions: (e.g., hitting a wall, where the action is either masked or treated as an identity/no-op).
   • Action-inhomogeneous worlds: Where the effect of an action depends on the specific state.
   • Results: These scenarios resulted in algebraic structures that are Monoids (if actions are total but not invertible) or Small Categories (if actions are partial/undefined in some states).

C. Categorical Generalization (Section 5)

Using Category Theory, the authors generalized two core concepts of SBDRL:

1. Equivariance Condition:
   • SBDRL: Requires a natural transformation between group actions.
   • Generalization: They defined equivariance for Monoids (single-object categories) and Small Categories (multi-object categories). This allows the equivariance condition to hold even when actions are irreversible or state-dependent.
2. Disentanglement:
   • They proved that if the action algebra decomposes into sub-algebras, the representation can be disentangled.
   • Crucially, they showed that disentangled sub-algebras can have independent equivariance conditions. This means different parts of the agent's representation can be learned using different learning algorithms tailored to their specific algebraic structure (e.g., one part for reversible group actions, another for irreversible monoid actions).

3. Key Contributions

1. Unified Mathematical Framework: A formal system to describe world transformations as algebras (Groups, Monoids, Categories) rather than restricting them to Groups.
2. Derivation and Limitation Analysis: A rigorous proof that SBDRL is a subset of this framework, specifically limited to worlds satisfying "World Condition 1" (unrestricted actions) and "World Condition 2" (global inverses).
3. Algorithmic Generation: The creation of algorithms to automatically generate Cayley tables for arbitrary agent-world interactions, identifying the underlying algebraic structure (Group vs. Monoid vs. Category).
4. Categorical Generalization of SBDRL:
   • Extending the equivariance condition from groups to monoids and small categories.
   • Extending the disentanglement definition to allow for independent learning of sub-algebras with different structural properties.
5. Empirical Validation: Demonstrated through examples (walls, movable blocks, consumables) that common RL scenarios form Monoids or Categories, not Groups, and are thus outside the scope of standard SBDRL but within the proposed framework.

4. Results

• SBDRL Recovery: The framework successfully reproduces SBDRL representations when the world satisfies group axioms.
• Complexity of Non-Group Worlds:
  • Adding a wall to a grid world (treating blocked moves as identity) changed the algebra from a 4-element Group to a 26-element Monoid.
  • Masking blocked moves (treating them as undefined) resulted in a 59-element Small Category.
  • Irreversible actions (consuming items) resulted in Monoids with 64 elements (identity treatment) or 20 elements (masked treatment).
• Theoretical Proof: The paper proves that the Yoneda Lemma perspective (focusing on relationships/morphisms rather than internal object structure) makes Category Theory the natural formalism for these generalizations.
• Independent Learning: The generalized disentangling definition implies that an agent can learn different subspaces of its representation independently, even if those subspaces correspond to different algebraic structures (e.g., one subspace for reversible navigation, another for irreversible resource management).

5. Significance and Impact

• Broader Applicability: This work removes the "group constraint" from symmetry-based learning, allowing AI developers to apply rigorous mathematical symmetry principles to a much wider range of RL problems, including those with irreversible dynamics and partial observability.
• Foundation for Efficient Learning: By identifying the correct algebraic structure (Monoid/Category) for a specific environment, learning algorithms can be tailored to exploit these structures, potentially improving data efficiency and generalization.
• Explainable AI (XAI): The framework offers a way to predict and analyze the algebraic structures an agent should learn, providing a theoretical basis for interpreting agent representations.
• Future Directions: The authors suggest this framework could unify various equivariance conditions in Deep Learning (CNNs, Graph Neural Networks) and improve World Models in RL by incorporating non-group symmetries.

In summary, the paper provides a foundational shift from Group-based Symmetry to Algebraic/Category-based Symmetry, offering a robust mathematical toolkit for the next generation of representation learning in Artificial Intelligence.