Policy Contrastive Decoding for Robotic Foundation Models

1. Planteamiento del Problema

Los modelos fundamentales robóticos (políticas de robots generalistas) han demostrado un gran potencial para la manipulación flexible y de propósito general. Sin embargo, los autores identifican un defecto crítico: estos modelos a menudo aprenden correlaciones espurias a partir de trayectorias de preentrenamiento. En lugar de centrarse en características de objetos relevantes para la tarea, las políticas dependen de señales de fondo irrelevantes (por ejemplo, iluminación, texturas u objetos de fondo específicos) para predecir acciones.

• Consecuencia: Cuando el entorno de despliegue cambia (por ejemplo, cambios de fondo, ajustes de iluminación o reubicación de objetos), el rendimiento de la política se degrada significativamente porque los "atajos" aprendidos ya no son válidos.
• Limitaciones actuales: Las soluciones existentes a menudo requieren reentrenamiento o ajuste fino, lo cual es computacionalmente costoso y requiere acceso a los pesos del modelo. Existe una falta de métodos libres de entrenamiento y plug-and-play para corregir estos fallos en tiempo de inferencia.

2. Metodología: Decodificación Contrastiva de Políticas (PCD)

Los autores proponen la Decodificación Contrastiva de Políticas (PCD), un plugin libre de entrenamiento diseñado para tiempo de inferencia que redirige la atención de una política robótica desde características espurias hacia pistas visuales relevantes para el objeto.

Mecanismo Central

PCD opera contrastando las distribuciones de probabilidad de acción generadas a partir de dos entradas visuales diferentes:

1. Observación Original ($o_i$): La entrada visual estándar.
2. Observación enmascarada del objeto ($\hat{o}_i$): La misma entrada visual donde los objetos relevantes para la tarea objetivo están enmascarados (eliminados/inpainteados).

El método calcula una nueva distribución de probabilidad de acción contrastiva ($\pi^*$) utilizando la siguiente fórmula:
$$\pi^*_\theta(a_i | \ell, o_i) = \frac{1}{C} \cdot \pi_\theta(a_i | \ell, o_i) \cdot \left( \frac{\pi_\theta(a_i | \ell, o_i)}{\pi_\theta(a_i | \ell, \hat{o}_i)} \right)^\alpha$$

• Lógica: Si la política depende de características espurias (fondo), la distribución de probabilidad será similar tanto para $o_i$ como para $\hat{o}_i$, resultando en una relación cercana a 1 (sin cambio). Si la política depende del objeto, la distribución para $\hat{o}_i$ diferirá significativamente, amplificando la probabilidad de la acción correcta en la salida final.
• Hiperparámetro ($\alpha$): Controla la fuerza del contraste. $\alpha=0$ recupera la política de referencia.

Componentes Técnicos Clave

Para hacer que PCD sea aplicable a diversas políticas robóticas, los autores introducen dos estrategias específicas:

1. Seguimiento a Enmascarado (Track2Mask):

   • Objetivo: Generar automáticamente observaciones enmascaradas del objeto para cada cuadro de una trayectoria sin intervención manual.
   • Proceso:
     • Anotación: Los objetos objetivo en el cuadro inicial se identifican mediante indicaciones humanas (Punto/Caja) o detectores de vocabulario abierto disponibles (por ejemplo, Grounding DINO).
     • Seguimiento: SAM2 (Segment Anything Model 2) rastrea y segmenta estos objetos en cuadros subsiguientes.
     • Inpainting: Los objetos segmentados se inpaintean (eliminan) para crear la observación enmascarada $\hat{o}_i$.

2. Modelado Probabilístico basado en KDE (KDE-PM):

   • Desafío: Las políticas autoregresivas (por ejemplo, OpenVLA) emiten probabilidades de token directamente, pero las políticas basadas en difusión (por ejemplo, Octo, $\pi_0$) generan acciones mediante un proceso de eliminación de ruido y no emiten nativamente una distribución de probabilidad $\pi_\theta(a|o)$.
   • Solución: Los autores utilizan Estimación de Densidad de Kernel (KDE) para aproximar la distribución de probabilidad de acción. Muestrean $N$ acciones candidatas de la política de difusión y ajustan una densidad de kernel gaussiano a estas muestras para estimar $\pi_\theta(a|o)$ y $\pi_\theta(a|\hat{o})$, permitiendo que la fórmula de decodificación contrastiva funcione con modelos de difusión.

3. Contribuciones Clave

1. Primera Solución Libre de Entrenamiento: PCD es el primer trabajo que propone un enfoque plug-and-play y libre de entrenamiento para mitigar las correlaciones espurias en políticas robóticas durante la inferencia.
2. Agnóstico al Modelo: Funciona como un plugin de caja negra compatible tanto con políticas autoregresivas (OpenVLA) como basadas en difusión (Octo, $\pi_0$) sin requerir acceso a pesos preentrenados ni ajuste fino.
3. Robustez ante Cambios de Distribución: Al contrastar explícitamente entradas relevantes para el objeto frente a entradas irrelevantes para el objeto, PCD obliga a la política a centrarse en características causales, mejorando significativamente la generalización a escenarios no vistos.

4. Resultados Experimentales

Los autores evaluaron PCD en 15 tareas diversas a través de entornos de simulación (SIMPLER) y entornos del mundo real utilizando tres políticas de última generación: OpenVLA, Octo y $\pi_0$.

• Rendimiento en Simulación:
  • OpenVLA: Mejoró la tasa de éxito en un 50.6% en promedio.
  • Octo: Mejoró la tasa de éxito en un 29.7% en promedio.
  • $\pi_0$: Mejoró la tasa de éxito en un 8.9% en promedio (a pesar de que $\pi_0$ ya era una línea base fuerte).
• Rendimiento en el Mundo Real:
  • En tareas de manipulación del mundo real utilizando la política $\pi_0$, PCD logró una mejora masiva del 108% en la tasa de éxito en comparación con la línea base.
  • El método demostró ser robusto frente a diversos distractores (cambios de iluminación, desorden de fondo, desplazamientos de posición de objetos).
• Estudios de Ablación:
  • El método es robusto ante diferentes estrategias de detección de objetos (indicaciones humanas vs. Grounding DINO vs. YOLO World).
  • Es robusto ante diferentes algoritmos de inpainting (LaMa obtuvo el mejor rendimiento).
  • El hiperparámetro $\alpha$ muestra un impacto claramente positivo cuando es $>0$.

5. Significado y Limitaciones

• Significado: PCD aborda una debilidad fundamental en los modelos fundamentales robóticos actuales (la dependencia de correlaciones espurias) sin el alto costo del reentrenamiento. Demuestra que intervenciones simples en tiempo de inferencia pueden mejorar drásticamente la generalización y la seguridad en el despliegue del mundo real.
• Limitaciones:
  • Sobrecarga Computacional: PCD requiere un segundo pase forward (para la imagen enmascarada) y seguimiento/inpainting de objetos, aumentando la latencia de inferencia en aproximadamente un 24-100% (dependiendo de la política y la paralelización). Sin embargo, los autores argumentan que esto es aceptable dados las ganancias masivas en las tasas de éxito.
  • Alcance: Actualmente se centra en la corrección en tiempo de inferencia; el trabajo futuro apunta a abordar las correlaciones espurias durante la fase de entrenamiento.

Conclusión: El artículo presenta un marco altamente efectivo y versátil (PCD) que mejora significativamente la fiabilidad y la generalización de los modelos fundamentales robóticos al obligarlos a atender las características visuales correctas, cerrando la brecha entre el rendimiento en simulación y la robustez en el mundo real.