What Must Generalist Agents Remember?

Título del artículo: ¿Qué deben recordar los agentes generalistas?
Autores: Khurram Yamin, Namrata Deka, Maitreyi Swaroop, Albert Ting, Jeff Schneider, Bryan Wilder (Carnegie Mellon University, Georgia Institute of Technology)

Planteamiento del problema

El artículo aborda una pregunta fundamental en el desarrollo de agentes de Aprendizaje por Refuerzo (RL) generalistas: ¿qué información interna debe almacenar un agente en su memoria para actuar de forma casi óptima a través de múltiples entornos (dominios) y diversos objetivos?

A diferencia de trabajos previos que asumen que al agente se le proporciona un índice o etiqueta explícita para la tarea o el dominio actual, este estudio analiza un escenario donde el agente no recibe dicho índice. El agente observa únicamente su propia trayectoria (historial de estados-acciones) y debe inferir la identidad del dominio latente necesaria para una selección de acciones competente. Los autores enmarcan esto como un proceso de decisión de Markov parcialmente observable (POMDP) de dominio latente, donde la variable oculta es el índice de dominio fijo $\sigma$ que altera la dinámica de las transiciones, mientras que el mecanismo de recompensa permanece compartido entre los dominios.

El desafío central es que dos dominios distintos pueden presentar la misma observación (estado) pero requerir acciones óptimas incompatibles. El artículo pregunta: ¿bajo qué condiciones es la memoria estrictamente necesaria para resolver esta ambigüedad, y permite una memoria suficientemente competente reconstruir la dinámica de transición local del entorno?

Metodología y marco teórico

Los autores formalizan la interacción en tiempo discreto con una familia de objetivos $\Psi$ y un conjunto de dominios $\Sigma$. Introducen dos resultados teóricos primarios:

1. Separación de representaciones (Teorema 1)
El primer resultado establece una condición necesaria para la memoria. Los autores definen un escenario de "cuello de botella compartido" (Supuesto 1) donde dos dominios distintos requieren conjuntos disjuntos de acciones casi óptimas en un mismo estado $s^\dagger$, separados por una brecha de valor significativa.

• Afirmación: Si un agente es uniformemente casi óptimo en todos los objetivos y dominios, su estado de memoria interna $M_t$ debe distinguir entre estos dominios en el estado de cuello de botella.
• Mecanismo: El teorema demuestra que las distribaciones de probabilidad condicional del estado de la memoria en el cuello de botella, dado el dominio, deben tener una distancia de Variación Total (TV) distinta de cero. Específicamente, si la política del agente es $\epsilon$-casi óptima, la distancia TV entre las distribuciones de memoria de los dos dominios está acotada inferiormente por una función de la brecha de valor, la probabilidad de alcanzar el cuello de botella y el factor de descuento.
• Implicación: La memoria no es meramente un estadístico suficiente para el control óptimo (como en los estados de creencia de un POMDP estándar), sino que es una condición necesaria para que las representaciones aprendidas logren una competencia uniforme cuando la identidad del dominio oculto es relevante para la acción.

2. Existencia de un decodificador (Teorema 2)
El segundo resultado investiga si la memoria requerida para la separación es también suficiente para reconstruir la dinámica del entorno.

• Configuración: Los autores introducen "objetivos de prueba" (una familia auxiliar $\Psi_F$) que actúan como pruebas predictivas de un solo paso. Definen la "memoria suficiente para el valor" como una memoria que puede predecir con precisión los valores de acción óptimos para estos objetivos de prueba.
• Afirmación: Si la memoria es suficiente en valor para una clase de objetivos de prueba que pueden distinguir los núcleos de transición locales (mediante una métrica de probabilidad integral), entonces existe un decodificador que puede reconstruir aproximadamente el núcleo de transición intervencional local $P^\sigma(\cdot | s, a)$.
• Implicación: Un estado de memoria que soporta la planificación casi óptima para un conjunto rico de objetivos contiene inherentemente la información suficiente para decodificar la dinámica local del entorno, permitiendo la planificación aproximada incluso sin que el entorno proporcione un modelo de mundo explícito.

Validación experimental

Los autores validan sus afirmaciones teóricas utilizando un entorno personalizado llamado ForkWorld:

• Entorno: Un mundo de rejilla (gridworld) con una intersección en T donde un pasillo conduce a una bifurcación (cuello de botella). El dominio es latente: en el dominio "normal", 'arriba' mueve al norte y 'abajo' mueve al sur; en el dominio "invertido", estas asignaciones se invierten. El objetivo (llegar arriba o abajo) es observable, pero el dominio no lo es.
• Cuello de botella: En la bifurcación, la acción óptima depende enteramente del dominio oculto. Un agente sin memoria no puede distinguir los dominios y se ve forzado a adivinar, limitando su rendimiento a un 50% de éxito.
• Agentes: El estudio compara agentes sin memoria (DQN), agentes con ventanas de historial explícitas (Stacked DQN) y agentes recurrentes (DRQN) frente a líneas de base de oráculo que conocen el dominio.

Resultados clave:

1. Necesidad de memoria: Los agentes sin memoria no lograron superar la línea de base del 50% (alcanzando aproximadamente un 19%), confirmando que sin la memoria para separar los dominios, la política no puede converger a una solución estable. Los agentes equipados con memoria (Stacked DQN y DRQN) alcanzaron tasas de éxito del 75–88%, cerrando la brecha con el oráculo.
2. Separación de representaciones: Una sonda lineal entrenada para predecir el dominio a partir del estado interno del agente en el cuello de botella mostró que los agentes con memoria lograron una clasificación de dominio casi perfecta (precisión >97%) en visitas subsecuentes a la bifurcación, mientras que los agentes sin memoria permanecieron en el nivel del azar (50%). Esto confirma empíricamente el Teorema 1.
3. Dinámicas decodificables: Los autores entrenaron un decodificador para predecir el siguiente estado a partir de la representación de la memoria del agente. Los agentes con memoria redujeron significamente el error de predicción de transición en comparación con una línea de base, mientras que los agentes sin memoria no mostraron mejora. Esto respalda el Teorema 2, demostrando que la memoria utilizada para la planificación también codifica la dinámica de transición local.

Contribuciones clave y significancia

El artículo realiza las siguientes contribuciones:

1. Necesidad formal de la memoria: Proporciona una prueba de control-teórica de que, para agentes generalistas que operan en dominios latentes, la memoria no es opcional, sino una condición necesaria para la casi optimalidad uniforme cuando los dominios comparten observaciones pero requieren acciones disjuntas.
2. La memoria como sustrato para modelos de mundo: Demuestra que la misma representación de memoria requerida para la selección de acciones, si es suficientemente rica para soportar la predicción de valor para objetivos de prueba, soporta inherentemente la decodificación de la dinámica de transición local.
3. Cambio de análisis: A diferencia de los resultados previos de "buen regulador" o de RL causal que asumen que el agente recibe un índice de tarea, este trabajo analiza el mapeo historial-a-memoria, mostrando que el agente debe aprender implícitamente a separar los dominios a partir de los datos de la trayectoria por sí solo.

Significancia:
Los autores argumentan que estos resultados caracterizan la memoria como el sustrato que soporta la disambiguación de dominios, la reconstrucción de la transición del modelo y la planificación para agentes generalistas. El trabajo sugiere que la capacidad de inferir el contexto latente y reconstruir la dinámica local no es una característica arquitectónica que deba añadirse, sino un requisito fundamental para cualquier agente que funcione bien a través de una familia de entornos cambiantes. Los hallazgos implican que los agentes de RL generalistas exitosos deben codificar internamente una representación que separe los dominios latentes, y que esta codificación facilita naturalmente la reconstrucción de la estructura causal del entorno.