Learning Abstract World Models with a Group-Structured Latent Space

Resumen Técnico: Aprendizaje de Modelos del Mundo Abstractos con un Espacio Latente Estructurado por Grupos

Enunciado del Problema

Aprender modelos abstractos significativos de Procesos de Decisión de Markov (MDP) es esencial para mejorar la generalización a partir de datos limitados en el aprendizaje por refuerzo (RL). Si bien los métodos existentes de aprendizaje de modelos del mundo incorporan con éxito datos de alta dimensión en espacios abstractos de menor dimensión para mejorar la eficiencia de las muestras, los enfoques más avanzados a menudo requieren grandes conjuntos de datos de entrenamiento para lograr un rendimiento superior al humano. Una limitación clave es la falta de integración de las estructuras geométricas conocidas (simetrías) inherentes a la dinámica del entorno. Sin tales priores, los agentes pueden tener dificultades para generalizar a escenarios no vistos o requerir datos excesivos para aprender simetrías simples, como invariancias rotacionales o traslacionales.

Metodología

Los autores proponen un marco que impone priores geométricos en la variedad de representación de baja dimensión de un modelo de transición aprendido. La metodología central implica:

1. Espacio Latente Estructurado por Grupos: En lugar de tratar el espacio de estados abstractos $Z$ como un espacio vectorial euclidiano estándar, los autores lo modelan como un espacio cociente (por ejemplo, $\mathbb{R}/k\mathbb{Z}$) o un producto de tales espacios. Esta elección refleja los grupos de simetría subyacentes del entorno (por ejemplo, grupos cíclicos para rotaciones).
2. Acciones de Grupo como Dinámicas de Transición: El modelo de transición $\tau: Z \times A \to Z$ se restringe a operar mediante una acción de grupo. Específicamente, el siguiente estado latente se predice como el estado actual combinado con un vector de transición aprendido mediante una operación de grupo $\oplus$:
   $$\hat{z}_{t+1} = z_t \oplus \Delta(z_t, a; \theta_{trans})$$
   Para espacios cociente como $\mathbb{R}/k\mathbb{Z}$, esta operación se implementa mediante aritmética modular, codificando naturalmente condiciones de frontera periódicas.
3. Representaciones Desentrelazadas: Para manejar entornos que contienen tanto características simétricas (estructuradas) como no simétricas (no estructuradas), el marco promueve el desentrelazado. Se aplica una regularización de dispersión al vector de transición $\Delta$, forzando a que las características que deben permanecer invariantes bajo una acción específica sean cero. Esto se logra enmascarando coordenadas específicas del vector de transición basándose en una mapeo $\sigma(a)$ que define qué dimensiones latentes se ven afectadas por la acción $a$.
4. Objetivo de Aprendizaje: El modelo se entrena utilizando un objetivo de aprendizaje contrastivo auto-supervisado (pérdida InfoNCE) para alinear las transiciones predichas con la verdad fundamental, combinado con una pérdida de predicción de recompensa y una pérdida de regularización de volumen para mantener el espacio abstracto compacto e interpretable. A diferencia de enfoques anteriores que imponen equivariancia mediante arquitecturas de redes neuronales con pesos atados (lo cual puede ser computacionalmente costoso), este método impone la estructura directamente dentro de la geometría del espacio latente.

Contribuciones Clave

• Piores Geométricos en el Espacio Latente: El artículo demuestra cómo incorporar estructuras simétricas conocidas (por ejemplo, simetrías rotacionales y traslacionales) seleccionando espacios latentes apropiados (espacios cociente) y acciones de grupo correspondientes, en lugar de modificar la arquitectura de la red neuronal.
• Manejo de Simetrías Mixtas: El marco aprende con éxito representaciones que combinan características estructuradas (simétricas) y no estructuradas, permitiendo que el modelo escale a entornos complejos como vistas 3D en primera persona donde no todas las dinámicas son simétricas.
• Mejora en la Generalización: El enfoque conduce a representaciones más simples y mejor desentrelazadas que generalizan mejor a transiciones no vistas y requieren menos muestras de entrenamiento en comparación con líneas base completamente no estructuradas.
• Rendimiento en RL Aguas Abajo: Los modelos abstractos aprendidos sirven como priores efectivos para tareas de RL aguas abajo, mejorando significativamente la eficiencia de las muestras en regímenes de bajos datos.

Resultados Experimentales

Los autores evaluaron su enfoque en varios entornos:

• MDP de Juguetes (Pasaje y Toro): En un mundo de cuadrícula cíclico simple y un MDP de toro (producto de dos grupos cíclicos), el modelo estructurado por grupos predijo con precisión transiciones no vistas donde las líneas base no estructuradas fallaron. El modelo recuperó la topología correcta (por ejemplo, mapeando el espacio latente a un toro en $\mathbb{R}^3$) utilizando solo 2 dimensiones latentes.
• MiniGrid (Vista desde Arriba): El método desentrelazó con éxito la orientación (modelada como un grupo cíclico) de la posición (modelada como espacio euclidiano), resultando en una representación de espacio producto $\mathbb{R}/2\pi\mathbb{Z} \times \mathbb{R}^2$.
• VizDoom (Vista 3D en Primera Persona): En un entorno de alta dimensión con entradas RGB, el modelo mapeó datos visuales complejos en un espacio latente de 3 dimensiones. Capturó correctamente la naturaleza cíclica de la rotación y la naturaleza traslacional del movimiento, superando a las líneas base en métricas de generalización (Hits@k, Recíproco Medio de Rango) y tareas de RL aguas abajo.
• Métricas Cuantitativas: En todos los entornos, el modelo "AWM + Piores Geométricos" superó significativamente a las líneas base no estructuradas (AWM sin priores) y a otros métodos (PRAE, Matriz de Rotación) en la generalización a transiciones no vistas, particularmente en regímenes de bajos datos. En tareas de RL aguas abajo (Toro, MiniGrid, VizDoom), los agentes que utilizan el prior geométrico lograron recompensas acumuladas más altas con menos pasos de entrenamiento.

Significado y Afirmaciones

El artículo afirma que, al aprovechar priores geométricos a través de la elección del espacio latente y las acciones de grupo, es posible aprender modelos del mundo más eficientes en muestras y generalizables sin la sobrecarga computacional de las redes neuronales equivariantes. Los autores enfatizan que su método permite la combinación flexible de características simétricas y no simétricas, haciéndolo aplicable a escenarios complejos del mundo real como juegos en primera persona 3D. El trabajo sugiere que respetar la simetría subyacente del entorno en la representación abstracta conduce a "representaciones más simples y mejor desentrelazadas", lo que a su vez facilita una mejor predicción y planificación. Los autores proporcionan el código completo para garantizar la reproducibilidad.