Resumen Técnico: Modelos Causales Centrados en Objetos para la Planificación con Búsqueda en Árbol de Monte Carlo

Declaración del Problema

El aprendizaje por refuerzo basado en modelos (MBRL) ha logrado un éxito significativo al aprender la dinámica del entorno para generar experiencias imaginadas, reduciendo así la necesidad de interacciones en el mundo real. Sin embargo, aprender modelos de mundo precisos sigue siendo un desafío en entornos de alta dimensionalidad y no estacionarios compuestos por múltiples objetos interactuantes.

Los enfoques actuales de RL visual dependen típicamente de codificadores de Redes Neuronales Convolucionales (CNN) que producen una única representación holística de una imagen de entrada. Estas representaciones monolíticas a menudo fallan al capturar la estructura y las interacciones a nivel de objeto, las cuales son críticas para la toma de decisiones. En escenas complejas, la presencia de fondos dinámicos, entidades irrelevantes o objetos pequeños pero relevantes para la tarea puede degradar significativamente el rendimiento del agente. Mientras que los humanos perciben el mundo como una composición de entidades discretas, lo que permite un razonamiento eficiente, los agentes de RL estándar luchan por aislar estas entidades sin supervisión explícita o máscaras de segmentación de instancias.

Metodología: COMET

Los autores presentan COMET (Modelo Causal Centrado en Objetos para la Búsqueda en Árbol Eficiente), un algoritmo de aprendizaje por refuerzo basado en modelos que realiza la Búsqueda en Árbol de Monte Carlo (MCTS) dentro de un espacio latente estructurado por ranuras (slots). COMET integra tres componentes primarios:

1. Codificador Centrado en Objetos Congelado

COMET utiliza un codificador centrado en objetos no supervisado y congelado (por ejemplo, SLATE, DINOSAUR o Slot Contrast) para mapear las observaciones visuales en un conjunto de representaciones latentes centradas en objetos, o "ranuras" ($\bar{s}_t = \{s^1_t, \dots, s^n_t\}$).

• Inicialización de Ranuras: Para abordar la invarianza a la permutación y la naturaleza estocástica de la atención de ranura en modelos de imágenes estáticas, los autores inicializan las representaciones de las ranuras en el tiempo $t+1$ utilizando las ranuras del tiempo $t$. Esto asegura la consistencia temporal de las asignaciones de objetos dentro de un episodio.
• Congelación: El codificador se preentrena con datos recolectados mediante una política aleatoria y permanece congelado durante la fase de RL para proporcionar características estables a nivel de objeto.

2. Modelo de Mundo Basado en Transformer con Fusión Acción-Ranura

El núcleo de COMET es un modelo de mundo basado en un transformer que opera sobre estas ranuras. A diferencia de los enfoques estándar de estilo MuZero que procesan un vector de estado monolítico, COMET procesa una secuencia de ranuras de objetos.

• Vinculación Acción-Objeto: Un mecanismo novedoso fusiona las acciones con las ranas de objetos. En lugar de comprimir todos los latentes de los objetos en un solo vector para predecir el siguiente estado (lo que crea un cuello de botella), la incrustación de la acción se concatena con cada ranura de forma independiente. Estas incrustaciones de acción condicionadas por la ranura pasan a través de un proyector MLP compartido y luego a un backbone de transformer para predecir las siguientes ranuras de estado ($\hat{s}_{t+1}$) y las recompensas.
• Arquitectura: El modelo sigue un flujo de entrenamiento de estilo UniZero utilizando un backbone de transformer basado en nanoGPT. Predice las futuras ranuras y recompensas basándose en el conjunto actual de ranuras y la acción tomada.

3. Atención Causal de Objetos para Política y Valor

Las cabezas de política y valor utilizan una arquitectura de transformer modificada equipada con atención causal de objetos.

• Puntuaciones de Causalidad: Para cada ranura latente $h^k_t$, la red estima una puntuación de causalidad $\alpha^k_t \in [0, 1]$, que representa la probabilidad de que el objeto correspondiente sea relevante para la decisión actual.
• Modulación de la Atención: Estas puntuaciones se utilizan para modular la atención de producto punto escalado. Una matriz de pesos aprendida $W_t$ eleva una estructura de grafo causal predefinida al nivel de los tokens, permitiendo que el mecanismo de atención se concentre en las entidades relevantes para la tarea (objetos causales) mientras suprime las irrelevantes (objetos no causales). Esto permite que el agente enfoque su planificación en el subconjunto específico de objetos que actualmente interactúan con el entorno.

Contribuciones Clave

1. Algoritmo COMET: La introducción de un algoritmo de MBRL basado en objetos y MCTS que combina representaciones de nivel de objeto congeladas con un modelo de mundo basado en transformer para la planificación en un espacio latente estructurado.
2. Mecanismo de Fusión Acción-Ranura: Un método novedoso para vincular acciones con ranuras centradas en objetos dentro de un backbone de transformer unificado. Esto permite el modelado de mundo y la predicción de política/valor centrados en objetos sin colapsar la información del objeto en un solo vector.
3. Atención Causal de Objetos: La integración de puntuaciones de causalidad aprendidas para modular las interacciones de los tokens, permitiendo que la política se enfoque dinámicamente en las entidades relevantes para la tarea durante la toma de decisiones.
4. Evaluación Empírica: Evaluación exhaustiva a través de ocho tareas visual y dinámicamente diversas, incluyendo el benchmark de RL visual centrado en objetos, ManiSkill, Robosuite y VizDoom.

Resultados Experimentales

COMET fue evaluado frente a varias líneas base:

• Baselines Centrados en Objetos: OCRL (PPO con codificador de objetos), OCDreamer (RSSM centrado en objetos) y SOLD (dinámica basada en ranuras no supervisada).
• Baselines Monolíticos: UniZero (planificación latente estándar de estilo MuZero).

Hallazgos de Rendimiento:

• Eficiencia de Muestreo: A través de la diversa suite de tareas, COMET logra un Puntaje Normalizado Medio más alto durante las etapas tempranas del entrenamiento en comparación tanto con los métodos de MCTS monolíticos fuertes como con otros baselines de RL centrados en objetos.
• Especificidad de la Tarea:
  • En tareas visualmente simples (ej. Object Goal, Object Comparison), COMET aprovecha las representaciones de objetos de alta calidad y la atención causal para identificar con precisión los objetos relevantes, lo que conduce a una convergencia más rápida.
  • En tareas de control continuo (ej. Block Lifting, Cube Pushing), COMET demuestra un rendimiento consistente, aunque no supera universalmente a todos los baselines en cada métrica específica (por ejemplo, no superó a los competidores en Block Lifting).
• Generalización: El método muestra robustez a través de entornos con variaciones en la complejidad visual, desde formas 2D sobre fondos monocromáticos hasta manipulación robótica 3D compleja y escenarios de disparos en primera persona.

Significado y Reivindicaciones

El artículo postula que modelar explícitamente la importancia de los objetos individuales para la toma de decisiones mejora el aprendizaje de la política en entornos orientados a objetos. Al introducir un sesgo inductivo de nivel de objeto explícito en la planificación latente de estilo MuZero, COMET permite al modelo razonar sobre entidades individuales y sus interacciones en lugar de tratar la escena como una masa holística.

Los autores afirman que COMET aborda eficazmente el desafío de representar observaciones en entornos complejos de múltiples objetos. Al desacoplar las representaciones de los objetos y utilizar la atención causal para filtrar la relevancia, el algoritmo logra una mayor eficiencia de muestreo y una toma de decisiones más enfocada. El trabajo sugiere que los sesgos inductivos centrados en objetos son un camino viable y efectivo para mejorar la planificación basada en modelos en RL visual, particularmente cuando no se dispone de segmentación de instancias o cuando se requiere un aprendizaje de representación no supervisado.