Resumen Técnico: Latent Memory Palace (LMP)

Planteamiento del Problema

La toma de decisiones humana exhibe un espectro que va desde acciones reflexivas hasta el razonamiento deliberado de múltiples pasos, donde el tiempo y el esfuerzo invertidos escalan con la complejidad de la tarea. Si bien los Modelos de Lenguaje de Gran Escala (LLM) han demostrado con éxito capacidades de "razonamiento" adaptativo similares mediante la generación iterativa de tokens (por ejemplo, Chain-of-Thought), transferir esta capacidad al control robótico continuo sigue siendo un desafío. Razonar directamente en el espacio del lenguaje suele carecer de la granularidad necesaria para la comprensión espacial y el control de movimiento preciso. Además, los enfoques estándar de inferencia variacional para el control suelen depender de variables latentes de dimensión fija, las cuales no pueden adaptar su profundidad computacional a los requisitos específicos de cada entrada. El problema central es cómo permitir que las políticas robóticas realicen una computación iterativa y adaptativa en un espacio latente que se aprenda de extremo a extremo para la predicción de acciones, equilibrando así la toma de decisiones eficiente con una mejor generalización.

Metodología

Los autores proponen el Latent Memory Palace (LMP), un marco que formula el razonamiento para el control como inferencia variacional autorregresiva.

1. Formulación Central

LMP trata una acción $a$ no como una predicción directa de una observación $o$, sino como el resultado de una computación latente intermedia $z$. La política se define mediante un límite inferior variacional (ELBO):
$$\log p(a | o) \geq \mathbb{E}_{z \sim q_\theta(z|o,a)} [\log p_\phi(a | o, z)] - D_{KL}(q_\theta(z | o, a) \parallel p_\theta(z | o))$$
A diferencia de los enfoques estándar, la variable latente $z$ no es un único vector, sino una secuencia autorregresiva de longitud variable $z = (z_1, z_2, \dots, z_{T(z)})$ extraída de un vocabulario que incluye un token de Fin de Secuencia (EOS). La secuencia termina cuando se genera el EOS, lo que permite al modelo asignar computación de forma adaptativa según la complejidad de la decisión.

2. Razonamiento Adaptativo mediante Compresión

Para suscitar un comportamiento adaptativo donde el modelo solo utilice más pasos según sea necesario, LMP introduce un mecanismo de compresión. El decodificador de acciones $p_\phi$ está parametrizado como una Gaussiana isotrópica donde la varianza $\sigma(T(z))$ decae exponencialmente con la longitud de la secuencia latente $T(z)$.
$$\sigma(T(z)) = \gamma^{T(z)} \sigma_0, \quad \gamma \in (0, 1)$$
Esto impone una penalización por longitud: cada paso latente adicional agudiza la distribución Gaussiana. En consecuencia, el modelo solo puede "permitirse" generar otro paso si la precisión de la acción predicha justifica la reducción de la varianza. Esto obliga a la política a dedicar más cómputo a decisiones complejas y menos a decisiones simples o irreduciblemente inciertas (por ejemplo, el tiempo exacto del cierre de una pinza).

3. Optimización Tratable mediante RL en el Espacio Latente

Optimizar el ELBO a través de un muestreo autorregresivo no diferenciable es intratable mediante la retropropagación estándar. Los autores derivan un objetivo de Aprendizaje por Refuerzo (RL) sobre las trayectorias latentes:

• Episodio Latente: Se muestrea una trayectoria de la distribución posterior $q_\theta$ hasta el EOS.
• Recompensa: Compuesta por una recompensa de reconstrucción ($\log p_\phi(a|o,z)$) y una penalización KL contra la distribución a priori.
• Reducción de Varianza: El término KL se descompone en recompensas por paso utilizando la estructura autorregresiva, permitiendo la asignación de crédito paso a paso.
• Optimización: Los autores emplean un objetivo de sustituto recortado tipo PPO para estabilizar el entrenamiento y prevenir el colapso latente, utilizando rollouts on-policy recolectados en un búfer.

4. Tokenización de Acciones (LMP-tok)

Al eliminar el condicionamiento de la observación del marco variacional, la misma arquitectura instancia LMP-tok, un tokenizador de acciones de longitud variable. Este comprime fragmentos de acciones continuas en tokens discretos latentes, que luego pueden usarse para entrenar políticas autorregresivas descendentes. A diferencia de los tokenizadores que requieren la reconstrucción exacta de transformaciones discretas (por ejemplo, FAST), LMP-tok permite que secuencias de longitud arbitraria sean decodificadas en acciones válidas.

Contribuciones Clave

1. Latent Memory Palace (LMP): Una formulación novedosa del razonamiento como inferencia variacional con una distribución latente autorregresiva, que permite la computación iterativa y adaptativa para el control robótico.
2. LMP-$\pi$ (Política): Una política de control que genera trazas de "razonamiento" de longitud variable durante el tiempo de prueba, asignando computación de forma adaptativa según la dificultad y la incertidumbre de la tarea.
3. LMP-tok (Tokenizador): Un tokenizador de acciones de longitud variable que mejora significamente el rendimiento de las políticas autorregresivas descendentes al comprimir las acciones en espacios latentes discretos.
4. Marco de Optimización: Un método tratable basado en RL para optimizar el límite inferior variacional para el muestreo autorregresivo de longitud variable, incluyendo técnicas para estabilizar el entrenamiento contra el colapso latente.

Resultados Experimentales

Los autores evaluaron LMP tanto en entornos reales (DROID) como simulados (LIBERO, D3IL, RoboMimic).

• Rendimiento de la Política (LMP-$\pi$):

  • Zero-Shot y Fine-tuning: En la plataforma DROID, LMP-$\pi$ superó a Diffusion Policy (DP) en configuraciones zero-shot y alcanzó tasas de éxito más altas (por ejemplo, 95% en limpieza de objetos múltiples) en tareas de alta precisión con ajuste fino.
  • Generalización: En LIBERO-90, LMP-$\pi$ mostró un "suelo" de rendimiento significativamente más alto en las 10 tareas más difíciles en comparación con DP, lo que indica una reducción de la interferencia entre tareas y una mejor capacidad de recuperación.
  • Iterativo vs. No Iterativo: LMP-$\pi$ superó consistentemente a las líneas base no iterativas (VAE Policy, VQ-BeT), demostando que las representaciones latentes de longitud fija carecen de la granularidad requerida para la manipulación compleja.

• Rendimiento del Tokenizador (LMP-tok):

  • Al utilizarse para entrenar políticas autorregresivas descendentes, LMP-tok superó significativamente a las líneas base (VQ-VAE, FAST, OAT) en las tareas de RoboMimic y LIBERO.
  • Notablemente, LMP-tok fue el único método que logró un rendimiento no trivial en la tarea de alta precisión "tool-hang" de RoboMimic, donde otros tokenizadores fallaron debido a restricciones de precisión o reconstrucción.

• Análisis de Computación Adaptativa:

  • Interpretabilidad: Las visualizaciones muestran que el número de pasos de razonamiento latente se correlaciona con la complejidad de la tarea. El modelo utiliza menos pasos durante los movimientos de la pinza (donde el tiempo es irreduciblemente variable) y más pasos durante tareas de alineación precisa.
  • Compresión: Los estudios de ablación confirmaron que una compresión más fuerte (decaimiento de varianza más rápido) conduce a una mejor generalización, siempre que se equilibre para evitar un error de reconstrucción excesivo.
  • Búsqueda de Modos (Mode-Seeking): Truncar los pasos latentes en el tiempo de prueba no degradó significamente el rendimiento; en su lugar, el modelo exhibió un comportamiento de "búsqueda de modo", comprometiéndose con modos de acción válidos específicos en lugar de promediarlos.

Significado y Reivindicaciones

El artículo afirma establecer la inferencia variacional autorregresiva como un marco práctico y fundamentado para la computación iterativa y adaptativa en el aprendizaje robótico. Al trazar una analogía con la mnemotecnia del "palacio de la memoria", el trabajo demuestra que organizar la información en un espacio latente de longitud variable permite a las políticas:

1. Imitar la Variabilidad en la Decisión Humana: Escalar el esfuerzo computacional con la complejidad de la decisión.
2. Mejorar la Generalización: Lograr un rendimiento robusto a través de diversas tareas y entornos al aprender a asignar el cómputo de manera eficiente.
3. Mejorar la Tokenización: Proporcionar un método superior para discretizar acciones continuas para el modelado autorregresivo.

Los autores posicionan a LMP como un puente entre las capacidades de razonamiento de los modelos fundacionales y los requisitos de precisión del control robótico, ofreciendo una nueva perspectiva donde el "razonamiento" es un proceso latente implícito optimizado mediante inferencia variacional y aprendizaje por refuerzo.