Resumen Técnico: Z-1 – Aprendizaje por Refuerzo Eficiente para Modelos de Visión-Lenguaje-Acción

1. Planteamiento del Problema

Los modelos de Visión-Lenguaje-Acción (VLA) han emergido como un marco prometedor para la manipulación robótica, unificando instrucciones de lenguaje, observaciones visuales y control continuo. Sin embargo, las políticas actuales de vanguardia (SOTA) se entrenan predominantemente mediante Clonación de Comportamiento (BC) o Ajuste Fino Supervisado (SFT) sobre demostraciones humanas fijas. Esta dependencia del aprendizaje por imitación presenta dos limitaciones críticas:

1. Generalización Limitada: El rendimiento está limitado por la cobertura de los datos de demostración. En tareas de largo horizonte o compositivas, los pequeños errores de predicción pueden acumularse, empujando al robot hacia estados pobremente representados en los datos de entrenamiento.
2. Incapacidad de Autocorrección: Los modelos BC y SFT carecen de mecanismos para mejorar a partir de sus propios fallos. No pueden optimizar el éxito de la tarea mediante la interacción de ensayo y error, lo cual es esencial para refinar comportamientos de manipulación complejos.

Si bien el post-entrenamiento mediante Aprendizaje por Refuerzo (RL) ofrece una vía para superar estas limitaciones optimizando las recompensas a nivel de tarea, aplicar RL a los modelos VLA basados en flujo sigue siendo un desafío. Los obstáculos clave incluyen el alto costo de la generación de rollouts, la escasez de recompensas y la dificultad de corregir fallos que derivan del anclaje visual o el razonamiento espacial cuando el núcleo de visión-lenguaje permanece congelado.

2. Metodología: El Marco Z-1

Z-1 es un marco de post-entrenamiento de RL modular diseñado específicamente para modelos VLA basados en flujo. Construido sobre el modelo base preentrenado $\pi_0.5$, Z-1 emplea un flujo de trabajo de entrenamiento en dos etapas utilizando únicamente las demostraciones de RoboCasa liberadas públicamente.

Etapa 1: Ajuste Fino Supervisado (SFT) por Escena

El flujo de trabajo comienza adaptando $\pi_0.5$ al entorno RoboCasa. En lugar de entrenar un modelo separado para cada tarea, Z-1 realiza un SFT por escena (agrupando tareas por categoría de escena, ej. puertas, cajones, fregaderos). Este enfoque mejora la eficiencia de los datos al aprovechar las disposiciones de objetos compartidas y las primitivas de manipulación dentro de una escena. Tanto el módulo de visión-lenguaje como el experto de acción se actualizan durante esta etapa para producir una inicialización robusta.

Etapa 2: Optimización de Política Relativa de Grupo por Tarea (GRPO)

Tras el SFT, Z-1 aplica una estrategia de post-entrenamiento GRPO específica para la tarea. A diferencia del RL estándar, que trata los rollouts como trayectorias completas independientes, Z-1 introduce cuatro componentes modulares clave para mejorar la eficiencia y la estabilidad:

• GRPO de Prefijo Compartido (Shared-Prefix GRPO): Para abordar el problema donde las fases iniciales de "aproximación" varían significamente y diluyen las señales de recompensa, Z-1 obliga a los rollouts dentro de un grupo GRPO a compartir un prefijo común. El estado del simulador en el punto de ramificación se clona, y solo los sufijos (acciones post-ramificación) se muestrean de forma independiente. Esto asegura que las comparaciones relativas de grupo se centren en las fases de interacción críticas en lugar de en variaciones redundantes de aproximación.
• Ramificación de Prefijo Estructurada en Árbol (Tree-Structured Prefix Branching): Para tareas que requieren refinamiento durante la fase de aproximación, Z-1 generaliza el enfoque de prefijo compartido plano. Introduce puntos de ramificación intermedios dentro del prefijo, creando un árbol de clusters de rollouts. Esto reduce el exceso de enmascaramiento, permitiendo que los segmentos parcialmente compartidos contribuyan a la señal de aprendizaje mientras se reduce la computación redundante.
• Decaimiento de Recompensa Consciente del Éxito (Success-Aware Reward Decay): Para manejar recompensas dispersas donde todos los trayectos exitosos son tratados por igual, Z-1 introduce un factor de decaimiento basado en el tiempo de finalización ($\tilde{r}_i = r \cdot \gamma^{d_i}$). Las finalizaciones exitosas más rápidas reciben mayores recompensas. Crucialmente, para grupos donde todos los rollouts tienen éxito, Z-1 utiliza un ordenamiento de línea base mínima no negativa ($A_i = R_i - \min_j R_j$) en lugar de la normalización de la media. Esto evita que los rollouts exitosos pero más lentos reciban ventajas negativas, evitando la penalización excesiva de ejecuciones fiables pero lentas.
• Entrenamiento Conjunto Selectivo VLM–Experto de Acción (Sel-JT): Mientras que el RL-VLA estándar suele congelar el núcleo de visión-lenguaje, Z-1 selecciona dinámicamente si actualizar el Modelo de Visión-Lenguaje (VLM) junto con el Experto de Acción. Esta decisión se basa en diagnósticos de la etapa de entrenamiento (ej. tasas de éxito de SFT, modos de fallo). Si los fallos se atribuyen a un pobre anclaje visual o razonamiento espacial en lugar de a un control de bajo nivel, el VLM se descongela para permitir la co-adaptación percepción-acción.

3. Contribuciones Clave

El artículo presenta las siguientes contribuciones:

1. Marco Z-1: Un marco de post-entrenamiento GRPO modular para modelos VLA basados en flujo que opera eficazmente utilizando solo datos de demostración públicos, evitando la necesidad de conjuntos de datos privados.
2. Nueva Construcción de Rollout: La propuesta de Shared-Prefix GRPO y Tree-Structured Prefix Branching para reducir la simulación redundante y centrar la optimización en segmentos de interacción críticos.
3. Innovaciones en Recompensa y Entrenamiento: El desarrollo de Success-Aware Reward Decay para proporcionar una retroalimentación más fina sobre la velocidad de finalización sin modelado de recompensa denso, y Selective Joint Training para adaptar los módulos de percepción solo cuando es necesario.
4. Nuevo Rendimiento SOTA: El establecimiento de un nuevo promedio de éxito de vanguardia en el benchmark RoboCasa bajo configuraciones de demostración pública.

4. Resultados Experimentales

Z-1 fue evaluado en 24 tareas de manipulación estándar de RoboCasa a través de 7 categorías (puertas, cajones, pick-and-place, microondas, fregadero/grifo, estufa y preparación de café).

• Ganancia de Rendimiento: El RL de Z-1 mejoró la tasa de éxito promedio del 67.4% (inicialización SFT) al 80.6%, una ganancia de 13.2 puntos porcentuales.
• Comparación con SOTA: Z-1 RL superó los resultados publicados anteriormente, incluyendo GR00T, GR00T N1.5, Video Policy y X-WAM. Específicamente, superó a X-WAM (el baseline más fuerte en la comparación) por 1.4 puntos porcentuales en promedio, a pesar de que X-WAM utilizaba información de profundidad mientras que Z-1 dependía únicamente de demostraciones públicas.
• Mejoras Específicas por Tarea:
  • Cajones y Fregadero/Grifo: Se observaron mejoras significativas (ej. Cajón: 83.4% $\to$ 96.1%; Fregadero: 63.2% $\to$ 94.3%), demostrando la eficacia de la construcción de rollouts basada en prefijos para tareas que requieren una manipulación precisa tras la aproximación.
  • Pick-and-Place: Se lograron ganancias sustanciales (53.9% $\to$ 70.9%).
  • Tareas de Estufa: Aunque Z-1 mejoró respecto a su baseline de SFT, se mantuvo ligeramente por debajo de X-WAM, indicando que las tareas que requieren transporte de largo horizonte complejo y recuperación robusta aún presentan desafíos.
• Estudios de Ablación:
  • Las variantes basadas en prefijos (Shared-Prefix y Tree-Structured) superaron consistentemente al GRPO vanilla, siendo Tree-Structured la que ofreció un mayor rendimiento máximo cuando el entrenamiento se mantenía estable.
  • El Entrenamiento Conjunto Selectivo (Sel-JT) superó significativamente al entrenamiento de solo Experto de Acción en tareas difíciles como "TurnOnStove", validando la necesidad de adaptar las representaciones visuales cuando el anclaje es insuficiente.
  • El Success-Aware Reward Decay proporcionó mejoras de tasa de éxito más rápidas y estables en comparación con las penalizaciones de recompensa estándar, que tendían a ser excesivamente agresivas.

5. Significado y Afirmaciones

El artículo afirma que el post-entrenamiento sistemático de GRPO puede mejorar sustancialmente las políticas VLA basadas en flujo sin depender de demostraciones privadas adicionales.

Los autores enfatizan que Z-1 demuestra la viabilidad del RL online para modelos VLA al abordar ineficiencias específicas en la generación de rollouts y la calibración de recompensas. Los resultados sugieren que, con un diseño cuidadoso —específicamente mediante la construcción de prefijos compartidos, la calibración de recompensa consciente de la finalización y la selección adaptativa de módulos—, el RL puede refinar eficazmente las políticas más allá de los límites del aprendizaje por imitación. El trabajo establece un nuevo benchmark para el entrenamiento VLA basado únicamente en datos públicos, mostrando que las mejoras algorítmicas en el post-entrenamiento pueden generar ganancias de rendimiento comparables o superiores a las de modelos entrenados con datos más extensos o propietarios.