Resumen Técnico: RoboSSM – Aprendizaje de Imitación In-Context Escalable mediante Modelos de Espacio de Estados

1. Planteamiento del Problema

El Aprendizaje de Imitación In-Context (ICIL, por sus siglas en inglés) permite que los robots se adapten a tareas novedosas utilizando pocas demostraciones sin actualizaciones de parámetros durante el despliegue. Si bien los métodos recientes de ICIL han adoptado con éxito arquitecturas basadas en Transformers para formular el aprendizaje de imitación como un problema de modelado de secuencias, estos enfrentan limitaciones significativas de escalabilidad. Los Transformers exhiben una complejidad temporal cuadrática ($O(L^2)$) relativa a la longitud de la secuencia y tienen dificultades para extrapolar más allá de las longitudes de prompt vistas durante el entrenamiento. En consecuencia, los métodos de ICIL existentes se degradan drásticamente al manejar prompts de prueba más largos (por ejemplo, más demostraciones o tareas de largo horizonte) que los utilizados durante el entrenamiento, lo que limita su aplicabilidad en escenarios del mundo real complejos que requieren un contexto extenso.

2. Metodología: RoboSSM

Los autores proponen RoboSSM, un marco de trabajo de ICIL escalable que reemplaza los Transformers con Modelos de Espacio de Estados (SSMs), utilizando específicamente Longhorn como la arquitectura base.

Arquitectura y Codificación

• Procesamiento de Entrada: RoboSSM codifica observaciones multimodales (datos visuales mediante CNNs y datos propioceptivos mediante MLPs) en embeddings. Notablemente, la representación de entrada excluye tokens de acción e instrucciones de lenguaje de la tarea para forzar al modelo a inferir la intención de la tarea únicamente a partir de las demostraciones in-context.
• Arquitectura Base Longhorn: La secuencia de embeddings de observación es procesada por Longhorn, que mantiene una matriz de estado de memoria $s_t$ actualizada mediante una recurrencia lineal. A diferencia del mecanismo de autoatención de los Transformers, Longhorn utiliza una regla de actualización recurrente derivada de la programación convexa en línea:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  Esta formulación equilibra la retención de la información de estado previa con la incorporación de nuevas entradas.
• Escalamiento $\beta$: Para fomentar específicamente la atención a los prompts de demostración en el entorno de ICIL, los autores introducen un mecanismo de escalado en tiempo de prueba donde el vector de ponderación $\beta_t$ se escala mediante un factor $\gamma \in (0, 1]$. Esto modula el equilibrio entre la retención del estado y la integración de nueva información.

Entrenamiento e Inferencia

• Entrenamiento: El modelo se entrena en un conjunto de datos multitarea donde la entrada consiste en un prompt (un conjunto de $N_{train}$ trayectorias de demostración) y una trayectoria de consulta. El modelo predice acciones para la trayectoria de consulta condicionada por el prompt.
• Inferencia: En el tiempo de prueba, la política se adapta a tareas no vistas utilizando un prompt de $N_{test}$ demostraciones. El modelo predice iterativamente las acciones basándose en la observación inicial y el contexto proporcionado por el prompt, sin realizar actualizaciones de parámetros.

3. Contribuciones Clave

1. Marco de Trabajo de ICIL Escalable: La introducción de RoboSSM, el primer enfoque que aprovecha los SSMs para el aprendizaje de imitación in-context, permitiendo una inferencia de tiempo lineal ($O(L)$) y fuertes capacidades de extrapolación.
2. Extrapolación de Contexto Largo: Demostrar que RoboSSM puede procesar eficazmente prompts de tiempo de prueba hasta 16 veces más largos que los vistos durante el entrenamiento, una capacidad en la que fallan los métodos basados en Transformers.
3. Inferencia Eficiente: Lograr tiempos de ejecución de inferencia significativamente menores en comparación con los baselines de Transformers a medida que aumenta la longitud del prompt, debido a la complejidad lineal de los SSMs frente a la complejidad cuadrática de los Transformers.
4. Generalización a Tareas de Largo Horizonte: Mostrar que el modelo puede generalizar desde distribuciones de entrenamiento de corto horizonte hacia tareas de largo horizonte y de múltiples etapas, así como escenarios con dilatación temporal.

4. Resultados Experimentales

Los experimentos se realizaron en el benchmark LIBERO, cubriendo diversas tareas de manipulación (suites de Objeto, Meta, Espacial, Largo y 90).

• Generalización de la Longitud del Prompt:
  • Cuando se entrenó con solo 2 demostraciones ($N_{train}=2$), RoboSSM mantuvo o mejoró las tasas de éxito a medida que aumentaba el número de demostraciones de prueba ($N_{test}$) hasta 32.
  • Por el contrario, el baseline basado en Transformer (ICRT) experimentó una caída drástica en el rendimiento y colapsó al 0% de éxito cuando $N_{test} > N_{train}$.
• Largo Horizonte y Dilatación Temporal:
  • RoboSSM resolvió con éxito tareas de largo horizonte (por ejemplo, "Poner tanto el Objeto A como el Objeto B en el Objeto C") y tareas con demostraciones con dilatación temporal (hasta 16 veces el estiramiento temporal), mientras que ICRT falló completamente.
• Rendimiento In-Distribution:
  • Incluso cuando $N_{train} = N_{test}$, RoboSSM superó consistentemente a ICRT en la mayoría de los escenarios, particularmente en las escenas de LIBERO-90 Study y Living Room.
• Comparación con Aprendizaje Multitarea (MTL):
  • RoboSSM e ICRT superaron significativamente a los baselines estándar de MTL (que requieren instrucciones de lenguaje y no utilizan prompting in-context), confirmando la eficacia del paradigma ICIL.
• Análisis del Espacio Latente:
  • Las visualizaciones de las representaciones latentes mostraron que RoboSSM genera trayectorias que permanecen cerca de los clústeres de las demostraciones del prompt, indicando una extrapolación estable. Por el contrario, las trayectorias de ICRT se desviaron significativamente fuera del manifold del prompt a medida que el contexto aumentaba.
• Tiempo de Ejecución:
  • RoboSSM demostró un escalado casi lineal con la longitud del prompt, mientras que el tiempo de ejecución de ICRT aumentó cuadráticamente, creando una brecha de rendimiento creciente a favor de RoboSSM para contextos largos.

5. Significado y Reivindicaciones

El artículo afirma que RoboSSM representa un cambio de paradigma para el aprendizaje de imitación in-context al establecer los Modelos de Espacio de Estados como un backbone eficiente y escalable. El significado principal radica en superar los cuellos de botella computacionales y de extrapolación de los Transformers, permitiendo que los robots:

• Manejen ventanas de contexto sustancialmente más largas en el tiempo de prueba sin necesidad de reentrenamiento.
• Se adapten a tareas no vistas, de largo horizonte y temporalmente variables utilizando prompts de pocos disparos (few-shot).
• Soporten la adaptación continua para el aprendizaje de por vida simplemente ingiriendo nuevos prompts de demostración sin actualizaciones de parámetros específicos de la tarea.

Los autores concluyen que, si bien RoboSSM demuestra el potencial de los SSMs para el ICIL de largo contexto, el trabajo futuro puede requerir corpora de entrenamiento más amplios y diversos para mejorar aún más la generalización. También señalan que, aunque la validación en robots reales no es la contribución central de este trabajo, sigue siendo una vía potencial para la demostración futura de la generalización.