PRTS: A Primitive Reasoning and Tasking System via Contrastive Representations

1. Planteamiento del Problema

Los modelos actuales Visión-Lenguaje-Acción (VLA) han avanzado el control robótico aprovechando los Modelos Visión-Lenguaje (VLM) como columnas vertebrales perceptuales, típicamente utilizando una arquitectura de doble sistema (Sistema 2 para el razonamiento, Sistema 1 para el control). Sin embargo, los enfoques existentes adolecen de una limitación fundamental:

• Estático vs. Temporal: Enmarcan predominantemente el pre-entrenamiento como clonación de comportamiento supervisada (BC), centrándose en correlaciones semánticas visuales-lingüísticas estáticas. Carecen de una comprensión intrínseca de la alcanzabilidad del objetivo—la probabilidad cuantitativa de que un par estado-acción actual conduzca a un objetivo especificado en lenguaje a lo largo del tiempo.
• Cuellos de botella en la estimación de valor: Los intentos de incorporar el aprendizaje de valor a menudo dependen de redes de valor separadas, requieren anotaciones manuales de recompensa o utilizan el aprendizaje de Diferencia Temporal (TD) que sufre de errores de bootstrap y alta varianza en tareas de largo horizonte.
• Brecha de generalización: Sin una conciencia temporal explícita, estos modelos luchan con la planificación de largo horizonte, instrucciones nuevas en cero disparos (zero-shot) y la recuperación de intervenciones humanas porque no pueden distinguir entre estados "fáciles pero erróneos" y "difíciles pero alcanzables".

2. Metodología: PRTS

Los autores proponen PRTS, un modelo base VLA que reformula el pre-entrenamiento mediante Aprendizaje por Refuerzo Condicionado por Objetivos (GCRL) utilizando Aprendizaje por Refuerzo Contrastivo (CRL).

A. Concepto Central: RL Contrastivo para VLA

En lugar de predecir acciones directamente o utilizar una red de valor separada, PRTS aprende un espacio de incrustación unificado donde el producto interno de una incrustación estado-acción ($\phi(s, a)$) y una incrustación de objetivo ($\psi(l)$) aproxima la ocupación de objetivo con descuento logarítmico.

• Base Teórica: Se basa en la equivalencia entre la función Q para recompensas de alcance de objetivo y la probabilidad de ocupación condicional (Lema 1).
• Objetivo: El modelo aprende a discriminar entre pares estado-acción que pueden alcanzar un objetivo de lenguaje específico (muestras positivas) y aquellos que no pueden (muestras negativas).

B. Innovaciones Técnicas Clave

1. RL Contrastivo Condicionado por Lenguaje:

   • A diferencia del CRL estándar que muestrea estados futuros geométricamente, PRTS trata la instrucción de lenguaje como el objetivo fijo para toda una trayectoria.
   • Problema de Múltiples Positivos: Dado que un único objetivo de lenguaje se aplica a todos los pasos de tiempo en una trayectoria, un lote contiene muchos pares positivos para un objetivo.
   • Pesado Temporal: Para recuperar la distribución de muestreo geométrico del CRL estándar, los autores introducen un esquema de pesado temporal. Los pares estado-acción más cercanos al final de una trayectoria (más cerca del objetivo) reciben un peso mayor ($\gamma^{T-t}$). Esto permite que el modelo aprenda la estructura temporal del alcance de objetivos sin etiquetas de recompensa explícitas.
   • Objetivos Bidireccionales:
     • Estado-Acción $\to$ Lenguaje: Alinea $(s, a)$ con su instrucción específica.
     • Lenguaje $\to$ Estado-Acción: Obliga al modelo a clasificar los estados-acciones dentro de una trayectoria por su proximidad al objetivo, codificando el progreso temporal.

2. Arquitectura de Paso Único (Sobrecarga Despreciable):

   • PRTS añade dos bloques de tokens auxiliares a la secuencia de entrada: <CRL_action> y <CRL_goal>.
   • Máscara Causal Consciente del Rol: Una máscara de atención personalizada (implementada mediante un kernel CuTe-FlashAttention fusionado) asegura el aislamiento de la información:
     • Los tokens <CRL_action> atienden solo a visión/propiocepción (no al lenguaje).
     • Los tokens <CRL_goal> atienden solo a la instrucción.
     • Los tokens de acción autoregresivos estándar mantienen la atención causal estándar.
   • Resultado: El modelo produce logits de acción discretos, incrustaciones estado-acción e incrustaciones de objetivo en un único paso hacia adelante, incurriendo en una sobrecarga computacional despreciable en comparación con la clonación de comportamiento estándar.

3. Pipeline de Entrenamiento en Dos Etapas:

   • Pre-entrenamiento: Optimiza conjuntamente la columna vertebral VLM en clonación de comportamiento (predicción del siguiente token), tareas auxiliares (Robot-QA) y las pérdidas contrastivas bidireccionales. Aún no se adjunta ningún experto en acciones.
   • Post-entrenamiento: Se adjunta un Experto en Acción de Ajuste de Flujo (basado en un Transformador de Difusión) y se ajusta finamente en datos de tareas posteriores. La columna vertebral pre-entrenada proporciona los priores de alcanzabilidad de objetivos.

3. Contribuciones Clave

1. Reformulación del Pre-entrenamiento VLA: Primer enfoque que escala el RL Contrastivo sin recompensas y condicionado por lenguaje directamente a la fase de pre-entrenamiento VLA, eliminando la necesidad de anotaciones manuales de recompensa o redes de valor separadas.
2. Conciencia Intrínseca de Alcanzabilidad de Objetivos: Las representaciones aprendidas codifican cuantitativamente la probabilidad de alcanzar un objetivo, cerrando la brecha entre el razonamiento semántico y el progreso temporal de la tarea.
3. Arquitectura Eficiente: La "Máscara Causal Consciente del Rol" y el kernel CuTe personalizado permiten una supervisión densa de alcanzabilidad de objetivos con una sobrecarga computacional cercana a cero, haciéndolo escalable a modelos grandes (4B de parámetros) y conjuntos de datos masivos (167B de tokens).
4. Marco Unificado: Unifica el razonamiento semántico, la predicción de acciones discretas y la estimación de valor dentro de una única arquitectura VLM.

4. Resultados Experimentales

PRTS fue pre-entrenado con 167.8 mil millones de tokens de diversos conjuntos de datos (AgiBotWorld, RoboMind, Open X-Embodiment, etc.) y evaluado en simulaciones y benchmarks del mundo real.

• Benchmarks de Simulación (LIBERO, LIBERO-Plus, LIBERO-Pro, SimplerEnv):

  • Rendimiento SOTA: PRTS logra resultados de última generación, igualando o superando a bases fuertes (por ejemplo, $\pi^0.5$, ABot-M0) mientras utiliza significativamente menos cómputo de post-entrenamiento (por ejemplo, 1/8 del presupuesto de $\pi^0.5$ en LIBERO).
  • Tareas de Largo Horizonte: Mejoras significativas en la suite "Long" de LIBERO (96.6% vs. 94.8% para bases con presupuesto igualado), demostrando la eficacia de la conciencia temporal.
  • Robustez: Supera a las bases en generalización de cero disparos bajo cambios de distribución (cámara, estado del robot, iluminación, disposición de objetos) y perturbaciones de Tarea/Posición en LIBERO-Pro.
  • Instrucción de Cero Disparos: Mejora masiva en instrucciones novedosas (31.5% vs. 0.8% para $\pi^0.5$ en el eje "Tarea" de LIBERO-Pro).

• Evaluación en el Mundo Real (RealMan de Doble Brazo & Flexiv de Brazo Único):

  • Manipulación Compleja: Logró una tasa de éxito promedio del 95.9% en 11 tareas de doble brazo y 90.0% en 3 tareas de brazo único de alta precisión.
  • Recuperación de Intervención Humana: PRTS demostró una robustez superior cuando los humanos interrumpían tareas (por ejemplo, retirando objetos durante la ejecución). A diferencia de las bases que fallaban o reproducían trayectorias obsoletas, PRTS reevaluó con éxito el estado y re-planificó para alcanzar el objetivo.
  • Generalización: Mantuvo altas tasas de éxito bajo cambios en iluminación, identidad de objetos y disposición espacial.

• Visualización de Valor:

  • Se demostró que el producto interno $\phi(s, a)^\top \psi(l)$ se correlaciona fuertemente con el progreso de la tarea, aumentando a medida que el robot se acerca al objetivo y permaneciendo bajo para instrucciones no coincidentes, confirmando que el modelo aprende una función de valor válida condicionada al objetivo.

5. Significado

PRTS representa un cambio de paradigma en los modelos base robóticos:

• De la Imitación al Razonamiento: Va más allá de la simple clonación de comportamiento (imitación) para aprender conciencia de alcanzabilidad de objetivos, permitiendo que el robot entienda qué tan cerca está de un objetivo y qué tan factible es un camino.
• Escalabilidad: Demuestra que los objetivos contrastivos, previamente limitados a dominios de toma de decisiones más pequeños, pueden escalar a columnas vertebrales VLA grandes (4B+ parámetros) y conjuntos de datos masivos sin costos prohibitivos.
• Inteligencia de Propósito General: Al inyectar priores temporales y conscientes de objetivos durante el pre-entrenamiento, PRTS logra una generalización y robustez notables de cero disparos, sugiriendo un camino hacia una inteligencia robótica verdaderamente de propósito general y dirigida a objetivos, capaz de manejar entornos reales complejos, de largo horizonte y dinámicos.