Resumen Técnico: UniJEPA

Enunciado del Problema

El desarrollo de políticas robóticas generalistas capaces de manejar tareas diversas en entornos abiertos sigue siendo un desafío central en la IA encarnada. Los enfoques existentes suelen depender de Modelos Visión-Lenguaje (VLM) para la comprensión semántica o de modelos generativos para la dinámica visual. Sin embargo, estos métodos enfrentan limitaciones significativas:

• Los enfoques basados en VLM a menudo sufren una degradación de las capacidades fundamentales al ser ajustados finamente en conjuntos de datos robóticos escasos y heterogéneos. Frecuentemente pasan por alto las discrepancias fundamentales entre las tareas estáticas de visión-lenguaje y las tareas de acción robótica dinámica.
• Los enfoques basados en generación (por ejemplo, predicción de video) facilitan el aprendizaje de representaciones dinámicas, pero a menudo fallan en preservar la alineación crucial visión-lenguaje inherente a los VLM preentrenados.
• Los Modelos Unificados actuales a menudo intentan unificar la comprensión y la generación dentro de un marco de predicción de tokens discretos, lo que puede comprometer la robustez de la alineación visión-lenguaje y la capacidad de modelar dinámicas físicas continuas de manera efectiva.

El artículo postula que el aprendizaje de políticas robóticas requiere un paradigma que integre la comprensión de tareas discretas (comprensión) con el aprendizaje de representación de estados futuros continuos (planificación y dinámica).

Metodología: UniJEPA

Los autores proponen UniJEPA, un marco Visión-Lenguaje-Acción (VLA) que unifica el aprendizaje de representaciones discretas y continuas mediante una estrategia de entrenamiento en dos etapas. La arquitectura utiliza un diseño de Mezcla de Transformadores (MoT) con expertos especializados por modalidad.

1. Modelado Unificado de Incrustaciones Visión-Lenguaje (Etapa 1: Pre-entrenamiento)

La primera etapa establece un paradigma de pre-entrenamiento cruzado-encarnado para aprender representaciones conjuntas de texto e imagen.

• Arquitectura: El modelo se inicializa con un VLM preentrenado (específicamente Paligemma) e introduce un Experto Generador y un Experto de Acción (en etapas posteriores) dentro del marco MoT.
• Aprendizaje de Representación Discreta: El modelo se entrena en conjuntos de datos a gran escala de visión-lenguaje y descripciones de tareas encarnadas (anotadas en formatos estilo VQA) para aprender representaciones lingüísticas detalladas para comprender escenas e instrucciones.
• Modelado del Mundo Continuo: A diferencia de trabajos anteriores que predicen píxeles de imagen crudos, UniJEPA predice características visuales continuas de alta dimensión. Un codificador visual congelado (por ejemplo, SigLIP) codifica observaciones futuras, y el modelo aprende a predecir estas características continuas ($\hat{o}_{t+h}$) junto con tokens de texto discretos ($\hat{l}$).
• Objetivo de Entrenamiento: El modelo emplea una función de pérdida conjunta que combina la pérdida de Entropía Cruzada para la rama de lenguaje y la pérdida de Error Cuadrático Medio (MSE) para la rama de predicción visual continua.

2. Modelado Unificado de Acción (Etapa 2: Ajuste Fino)

En la segunda etapa, el modelo se ajusta finamente en datos específicos de la encarnación para mapear representaciones predictivas a tokens de acción.

• Experto de Acción: Se entrena un experto distinto desde cero para manejar el espacio de acciones. Aprovecha el ajuste de flujo (flow matching) para capturar la distribución continua e intrínsecamente multimodal de las acciones.
• Optimización Conjunta: El modelo se entrena para predecir simultáneamente estados visuales futuros continuos y ejecutar secuencias de acciones. La función objetivo combina la pérdida de predicción visual continua con una pérdida de ajuste de flujo para la generación de acciones.
• Estrategia de Datos: El pre-entrenamiento utiliza una mezcla de videos robóticos (con descripciones de subtareas), videos de demostración humana y datos genéricos de preguntas y respuestas visión-lenguaje. El ajuste fino se realiza exclusivamente en datos VLA de simulación y robots del mundo real.

Contribuciones Clave

1. Arquitectura VLA Novel: La introducción de un marco VLA que integra la predicción de tokens discretos (para la comprensión) con la predicción visual continua (para la dinámica) dentro de una arquitectura MoT unificada.
2. Marco de Entrenamiento en Dos Etapas: Una estrategia que alinea las representaciones de acción mientras preserva las representaciones intermedias aprendidas durante el pre-entrenamiento a gran escala, transfiriendo efectivamente el conocimiento desde datos a escala de internet a tareas robóticas.
3. Predicción de Características Continuas: La demostración de que predecir características visuales continuas (en lugar de píxeles crudos o tokens discretos) sirve como una señal más efectiva para aprender información dinámica crucial para la generación de acciones.

Resultados Experimentales

UniJEPA fue evaluado en dos puntos de referencia de simulación y dos plataformas robóticas del mundo real (un brazo Franka Emika Panda de 7 grados de libertad y una mano hábil de 12 grados de libertad).

• Puntos de Referencia de Simulación:
  • Calvin Benchmark: UniJEPA logró una longitud promedio de secuencia de subtareas de 4.11, superando al mejor anterior (UP-VLA en 4.08) y superando significativamente a líneas base como GR-1 (3.06).
  • SimplerEnv: En las plataformas WindowX y Google Robot, UniJEPA logró tasas de éxito de vanguardia de 71.0% y 78.4%, respectivamente. Demostró un alto rendimiento consistente en todas las subtareas, evitando los perfiles de rendimiento "puntiagudos" de otros métodos.
• Experimentos en el Mundo Real:
  • Brazo Franka Panda: UniJEPA logró las tasas de éxito más altas en las cuatro categorías de tareas (Agarrar y Colocar, Presionar Botón, Enrutamiento de Cable, Cajón), incluyendo en tareas no vistas que involucraban objetos nuevos.
  • Mano Hábil de 12 Grados de Libertad: El modelo logró la tasa de éxito promedio más alta en nueve categorías de habilidades, mostrando una mejora del 12% en tareas no vistas en comparación con las líneas base.
• Estudios de Ablación:
  • Eliminar el pre-entrenamiento resultó en una caída de aproximadamente el 16% en la tasa de éxito en tareas de agarrar y colocar del mundo real.
  • El uso de características visuales continuas para la predicción produjo un impulso de rendimiento de aproximadamente el 20% en comparación con modelos sin este componente.
  • Entre los métodos de codificación continua, las características SigLIP proporcionaron la mejor alineación con la base del VLM, superando a los enfoques basados en DINO y destilación.

Significado y Afirmaciones

El artículo afirma que UniJEPA aborda con éxito el compromiso entre la comprensión semántica y la predicción dinámica en el aprendizaje de políticas robóticas. Al aprovechar el pre-entrenamiento a gran escala en datos de demostración tanto robóticos como humanos, el modelo adquiere capacidades robustas de generalización.

Los autores enfatizan que su enfoque permite a los robots manejar objetos completamente nuevos y escenarios fuera de distribución (OOD) de manera más efectiva que los métodos existentes. La integración del aprendizaje de representación futura continua permite a la política anticipar mejor las dinámicas físicas, mientras que la alineación visión-lenguaje preservada asegura una interpretación precisa de instrucciones complejas. Los resultados sugieren que un enfoque unificado que combina comprensión, planificación y aprendizaje de representaciones continuas es un camino viable hacia políticas robóticas generalistas.