Scaling Short-Term Memory of Visuomotor Policies for Long-Horizon Tasks

Resumen Técnico: Escalamiento de la Memoria de Corto Plazo en Políticas Visuomotoras para Tareas de Largo Horizonte

1. Planteamiento del Problema

Las tareas robóticas a menudo requieren memoria de corto plazo (que abarca segundos o minutos) para manejar la observabilidad parcial, como recuperar objetos que ya no son visibles o cronometrar acciones como apagar un electrodoméstico. Sin embargo, la mayoría de las políticas visuomotoras actuales entrenadas mediante aprendizaje por imitación son miope, ya que dependen únicamente de entradas sensoriales inmediatas sin aprovechar las experiencias pasadas.

Extender las políticas estándar basadas en Transformers para incluir ventanas de contexto más largas con el fin de abordar esto presenta dos desafíos críticos:

1. Correlaciones Espurias: El hecho de atender ingenuamente a historiales largos introduce información irrelevante (por ejemplo, ubicaciones de distractores pasados), lo que lleva a los modelos a aprender correlaciones falsas que degradan el rendimiento cuando el contexto cambia.
2. Escalabilidad Computacional: La huella computacional y de memoria de los mecanismos de atención estándar escala de forma lineal o cuadrática con la longitud del contexto. Procesar entradas visuales de alta dimensión a través de secuencias temporales largas (por ejemplo, minutos) resulta prohibitivamente costoso.

2. Metodología: PRISM

Los autores proponen PRISM (Arquitectura de Política con Memoria de Corto Plazo), una arquitectura basada en Transformers diseñada para utilizar eficazmente la memoria de corto plazo a través de dos innovaciones clave:

• Atención con Compuerta (Filtrado de Información): Para mitigar las correlaciones espurias, PRISM introduce un mecanismo de compuerta post-atención. Este componente filtra selectivamente la información recuperada modulando la salida de cada bloque de atención mediante una compuerta aprendible condicionada por las características de la entrada actual. Esto permite que la política adapte la supresión de detalles irrelevantes del historial, reduciendo el ruido en la predicción de la acción.
• Arquitectura Jerárquica (Eficiencia): Para abordar la escalabilidad, PRISM emplea un diseño jerárquico.
  • Codificación Local: Un codificador de memoria local procesa las entradas sensoriales y las comprime en tokens compactos utilizando un remuestreador (resampler) de estilo Perceiver.
  • Atención Global: Un codificador de memoria global computa la atención sobre estos tokens comprimidos a través del tiempo.
  • Fusión: La información global resultante se transmite de vuelta y se fusiona con los tokens locales originales mediante conexiones residuales.
  • Beneficio: Este enfoque reduce la complejidad computacional de $O((n \cdot k_{tot})^2)$ a aproximadamente $O(n \cdot k_{tot} \cdot m + (n \cdot k_{tot}/m)^2)$, donde $m$ es el factor de compresión. También reduce significativamente la huella de memoria en el tiempo de prueba al almacenar en caché únicamente los tokens de resumen comprimidos.

PRISM se entrena utilizando clonación de comportamiento estándar (aprendizaje por imitación) sin pérdidas auxiliares ni módulos de memoria externos, mapeando secuencias de entradas sensoriales y acciones pasadas hacia acciones futuras en un entorno de bucle cerrado.

3. Contribuciones Clave

• Arquitectura PRISM: Una novedosa arquitectura de política que combina la atención con compuerta y la sumarización jerárquica para escalar la memoria de corto plazo hasta dos minutos (a 2–3 Hz) manteniendo la eficiencia computacional y la robustez al ruido.
• REMEMBENCH: Un nuevo benchmark que comprende ocho diversas tareas de manipulación doméstica que abarcan cuatro categorías de memoria de corto plazo derivadas de la ciencia cognitiva:
  • Espacial: Recordar la ubicación de objetos (por ejemplo, recuperar fruta oculta).
  • Prospectiva: Retener intenciones durante un retraso (por ejemplo, cocinar carne durante una duración específica).
  • Asociativa de Objetos: Recordar asociaciones entre objeto y ubicación (por ejemplo, devolver un artículo lavado a su platillo original específico).
  • Conjunto de Objetos: Mantener y actualizar conjuntos de múltiples objetos (por ejemplo, contar palitos de pan movidos a un microondas).
  • A diferencia de los benchmarks previos que se centran en aspectos únicos o pruebas de memoria indirectas, REMEMBENCH requiere explícitamente memoria para tener éxito, ya que las entradas sensoriales inmediatas son insuficientes.

4. Resultados Experimentales

El artículo evalúa PRISM frente a arquitecturas recurrentes (LSTM, Mamba), transformers recurrentes de nivel de segmento (Transformer-XL) y otras variantes de Transformers (Atención Lineal, Scene Memory Transformer) en REMEMBENCH y benchmarks estándar.

• Rendimiento en REMEMBENCH: PRISM logra una tasa de éxito promedio del 41%, superando a las bases más fuertes (SMT con 36%, PTP con 28%, Mamba con 25%) por un margen absoluto del 5%–12%.
  • Escalamiento: El rendimiento escala con el tamaño de la memoria; aumentar la ventana de 1 a 512 pasos de tiempo mejora las tasas de éxito del 16% al 42% sin saturación.
  • Ablación: Eliminar el mecanismo de compuerta provoca una caída de 16 puntos en el rendimiento, confirmando su necesidad para filtrar el historial irrelevante.
• Benchmarks Estándar (ROBOCASA & LIBERO):
  • En ROBOCASA, PRISM con memoria (n=256) logra un 43% de éxito, una mejora absoluta del 11% sobre su variante sin memoria y un 11% de mejora sobre la base ajustada GR00T-N1-2B.
  • En LIBERO, PRISM logra un 89% de éxito promedio, una mejora del 15% sobre la base ajustada OpenVLA y una mejora del 12% sobre su variante sin memoria.
  • Los autores señalan que añadir memoria ayuda a resolver la ambigüedad de la acción en estados visualmente idénticos (por ejemplo, distinguir entre "recién agarrado" y "a punto de colocar"), reduciendo la multimodalidad incluso en tareas no diseñadas explícitamente para probar la memoria.
• Evaluación en el Mundo Real: En una adaptación del mundo real de la tarea "Lavar y Devolver al Contenedor", PRISM logró un 30% de éxito, comparado con el 15% de una base de Transformer sin memoria y el 0% de una política sin memoria.
• Eficiencia: Con una longitud de contexto de 512 pasos, PRISM reduce el uso de memoria GPU pico en un 80% y el tiempo de ejecución en un 40% en comparación con la atención global completa y los enfoques híbridos de ventana deslizante.

5. Significado y Reivindicaciones

El artículo afirma que PRISM y REMEMBENCH establecen una base para desarrollar y evaluar políticas visuomotoras aumentadas con memoria de corto plazo capaces de manejar tareas de largo horizonte.

• Generalización: Al cubrir cuatro categorías distintas de memoria, REMEMBENCH fomenta el desarrollo de mecanismos de memoria generales en lugar de soluciones aisladas y específicas para cada tarea.
• Escalabilidad: PRISM demuestra que la memoria de corto plazo puede escalarse a minutos de interacción sin los costos prohibitivos asociados con las extensiones ingenuas de la ventana de contexto.
• Robustez: El mecanismo de atención con compuerta demuestra ser esencial para evitar que el modelo aprenda correlaciones espurias en datos de largo horizonte ruidosos.

Los autores concluyen que, si bien los éxitos actuales en REMEMBENCH (alrededor del 30–40%) indican que las tareas de largo horizonte siguen siendo desafiantes, las ganancias consistentes del aumento de la memoria a través de diversas arquitecturas y benchmarks validan el enfoque. Identifican direcciones futuras tales como la integración de la memoria persistente de largo plazo y el aprovechamiento del preentrenamiento a escala de internet para los módulos de memoria, pero enfatizan que su trabajo actual se centra estrictamente en la mecánica de la memoria de corto plazo dentro del marco del aprendizaje por imitación.