Gated Memory Policy

1. El Problema

Las tareas de manipulación robótica presentan requisitos de memoria muy variados:

• Tareas Markovianas: Requieren poca o ninguna memoria histórica (ej. agarrar y colocar objetos simples).
• Tareas No Markovianas: Requieren memoria intra-episodio (dentro de una sola ejecución) o memoria inter-episodio (agregar información a través de múltiples intentos para inferir propiedades físicas no observables, como la fricción).

Desafíos actuales:

• Extensión ingenua del historial: Aumentar simplemente la longitud de la historia de observaciones en políticas visuomotoras (como los Difusion Policies) suele provocar una caída significativa en el rendimiento en tareas Markovianas debido al sobreajuste (overfitting) y al desplazamiento de distribución (distribution shift).
• Costo computacional: Procesar historiales largos en arquitecturas basadas en Transformers escala cuadráticamente ($O(H^2)$), lo que es ineficiente.
• Ruido y Robustez: Las políticas son sensibles a observaciones ruidosas o imprecisas en el historial, lo que puede desestabilizar la inferencia.

2. Metodología: Gated Memory Policy (GMP)

Los autores proponen GMP, una política visuomotora que aprende dinámicamente cuándo recordar memoria y qué recordar. Se basa en un Diffusion Transformer (DiT) y consta de tres componentes clave:

A. Mecanismo de Puerta de Memoria (Memory Gate)

Para aprender cuándo recordar:

• Se utiliza un MLP que toma la observación actual ($I_t$) y la propiocepción ($P_t$) para generar un valor binario $\mu_t \in \{0, 1\}$.
• Si $\mu_t = 1$, se activa la atención cruzada al historial; si $\mu_t = 0$, se ignora el historial.
• Calibración de la puerta: Entrenar la puerta de forma end-to-end es difícil (tiende a usar memoria siempre). En su lugar, proponen un proceso de calibración auto-supervisada:
  1. Se entrenan dos políticas separadas en el conjunto de entrenamiento: una sin memoria ($\pi$) y otra con memoria siempre activa ($\pi_{mem}$).
  2. Se evalúan en un conjunto de validación.
  3. Si el error de predicción de la política sin memoria es significativamente mayor que el de la política con memoria en un paso de tiempo $t$, se etiqueta ese paso como "requiere memoria" ($\mu_t=1$).
  4. Se entrena la puerta con estas etiquetas y luego se congelan sus pesos para entrenar la política final en todo el dataset.

B. Módulo de Atención Cruzada (Cross-Attention)

Para aprender qué recordar eficientemente:

• En lugar de concatenar todo el historial en una secuencia larga (costoso), GMP utiliza un módulo de atención cruzada.
• Los tokens de historial (acciones y características visuales) se almacenan en una caché y se consultan solo cuando la puerta está activa.
• Esto reduce la complejidad computacional de cuadrática a lineal respecto a la longitud del historial.

C. Augmentación con Ruido de Difusión

Para mejorar la robustez:

• Durante el entrenamiento (y la inferencia), se inyecta ruido de difusión programado en las acciones históricas.
• Esto evita que la política dependa de historias "perfectas" o limpias, mitigando la sensibilidad a errores acumulados y mejorando la generalización en entornos reales.

3. Contribuciones Clave

1. Gated Memory Policy (GMP): Un marco práctico que selecciona la memoria bajo demanda, manteniendo la reactividad en tareas Markovianas y la capacidad de adaptación en tareas No Markovianas.
2. MemMimic: Un nuevo conjunto de benchmarks (simulación y mundo real) diseñado específicamente para evaluar el uso de memoria intra-episodio e inter-episodio, llenando un vacío en las evaluaciones existentes.
3. Estudio Empírico Exhaustivo: Análisis en simulación y realidad que demuestra cómo las decisiones de diseño (puerta, ruido, atención cruzada) impactan el rendimiento, la robustez y la eficiencia computacional.

4. Resultados

El método se evaluó en MemMimic (tareas No Markovianas) y RoboMimic (tareas Markovianas):

• Rendimiento en Tareas No Markovianas: GMP logró una mejora promedio del 30.1% en la tasa de éxito en comparación con las líneas base de "historial largo" (Long-hist baselines) en el benchmark MemMimic.
  • Ejemplos destacados: Iterative Pushing (inferencia de fricción), Iterative Flinging (inferencia de masa) y Iterative Casting (ajuste de velocidad de lanzamiento).
• Rendimiento en Tareas Markovianas: A diferencia de las políticas de historial largo que sufren degradación, GMP mantuvo un rendimiento competitivo en RoboMimic (Square, Transport, Tool Hang) gracias a que la puerta de memoria se apaga ($\mu_t=0$) cuando no es necesario.
• Eficiencia Computacional:
  • Al usar atención cruzada y saltar la atención cuando la puerta está cerrada, el tiempo de inferencia se mantiene constante y bajo cuando no se necesita memoria, evitando el costo cuadrático de los Transformers estándar.
• Robustez: La inyección de ruido en las acciones históricas demostró ser crucial para evitar el sobreajuste y manejar entornos ruidosos del mundo real.

5. Significado e Impacto

El trabajo de GMP es significativo porque resuelve la dicotomía tradicional entre reactividad (baja latencia, tareas simples) y memoria a largo plazo (adaptación, tareas complejas).

• Generalización: Permite a un solo modelo manejar tareas que van desde la manipulación simple hasta la adaptación en contexto (in-context adaptation) sin necesidad de cambiar la arquitectura o hiperparámetros.
• Eficiencia: Demuestra que no es necesario procesar todo el historial en cada paso; la selección inteligente de qué recordar es más eficiente y robusta.
• Aplicabilidad Real: La validación en tareas de "mundo real" (Flip and Place Back, Iterative Casting) con robots físicos (UR5, ARX X5) y cámaras móviles (iPhone) confirma que la metodología es viable fuera de la simulación, superando las limitaciones de las políticas puramente basadas en Transformers o RNNs tradicionales.

En resumen, GMP introduce un mecanismo de "memoria bajo demanda" que hace que las políticas de robots sean más inteligentes, rápidas y robustas, permitiendo aprender de la experiencia pasada solo cuando es estrictamente necesario para la tarea.