Exploring Conditions for Diffusion models in Robotic Control

1. Planteamiento del Problema

El trabajo aborda el desafío de adaptar modelos de difusión preentrenados (específicamente modelos de texto-a-imagen como Stable Diffusion) para su uso en control robótico mediante aprendizaje por imitación.

• Limitación de los enfoques actuales: Las representaciones visuales preentrenadas (como CLIP o VC-1) suelen ser "agnósticas a la tarea" (task-agnostic). Se mantienen congeladas durante el aprendizaje de la política, lo que significa que no se adaptan a los requisitos específicos de cada tarea de control, limitando su rendimiento.
• Fallo de las condiciones textuales: En tareas de visión general (como segmentación semántica), usar descripciones de texto (prompts) como condiciones para modelos de difusión ha demostrado ser muy efectivo. Sin embargo, los autores descubren que aplicar esta misma estrategia en entornos de control robótico es ineficaz o incluso perjudicial.
• Causa raíz: Existe una brecha de dominio entre los datos de entrenamiento del modelo de difusión (imágenes de la web con objetos comunes) y los entornos de control robótico (agentes especializados en tareas específicas). Además, el control requiere una granularidad visual fina y dinámica (por cuadro), algo que las descripciones de texto generadas automáticamente no pueden capturar con precisión, llevando a una "grounding" (asociación texto-imagen) errónea.

2. Metodología: ORCA

Para superar estas limitaciones, los autores proponen ORCA (un marco para aprovechar modelos de difusión con condiciones adaptativas a la tarea). La metodología se basa en dos pilares principales, evitando el ajuste fino (fine-tuning) del modelo de difusión completo para mantener la generalización:

A. Prompts de Tarea Aprendibles (Learnable Task Prompts)

En lugar de usar descripciones de texto fijas o generadas por LLMs (que fallan en la asociación con regiones específicas), ORCA introduce prompts de tarea aprendibles.

• Estos son parámetros entrenables que se optimizan directamente durante el aprendizaje de la política de control.
• Actúan como "palabras implícitas" que aprenden a enfocarse en las regiones relevantes de la imagen para una tarea específica (por ejemplo, el brazo robótico y el botón en una tarea de presionar un botón), sin necesidad de descripciones textuales explícitas.

B. Prompts Visuales (Visual Prompts)

Dado que el control robótico opera sobre flujos de video dinámicos, se requiere información visual detallada por cuadro.

• ORCA utiliza un codificador de visión (como DINOv2) para extraer representaciones visuales densas (no globales) de cada cuadro.
• Estas representaciones se proyectan a través de una pequeña capa convolucional y se utilizan como prompts visuales.
• Esto permite al modelo de difusión capturar detalles finos y movimientos específicos del agente en cada frame, guiando la acción de manera dinámica.

C. Aprendizaje de la Política

• Objetivo: Se utiliza el objetivo estándar de clonación de comportamiento (Behavior Cloning).
• Proceso: La red de política toma las representaciones visuales extraídas del modelo de difusión (condicionado por los prompts de tarea y visuales) y predice la acción.
• Optimización: Los prompts de tarea y visuales se aprenden de extremo a extremo (end-to-end) junto con la red de política, sin modificar los pesos del modelo de difusión subyacente (Stable Diffusion v1.5).

3. Contribuciones Clave

1. Identificación del fallo de los prompts de texto: Demostración empírica y análisis (mediante mapas de atención cruzada) de que los prompts de texto, aunque exitosos en visión estática, fallan en control robótico debido a la brecha de dominio y la incapacidad de capturar la dinámica frame-a-frame.
2. Propuesta de ORCA: Un marco simple pero efectivo que combina prompts de tarea aprendibles y prompts visuales para condicionar modelos de difusión en tareas de control.
3. Eficiencia y Adaptabilidad: A diferencia del ajuste fino completo (que causa sobreajuste en datos limitados de imitación), ORCA aprende solo los prompts, manteniendo el modelo base congelado y permitiendo cambiar de tarea simplemente reentrenando los prompts.
4. Nuevo Estado del Arte (SOTA): Logra el mejor rendimiento en múltiples benchmarks de control robótico sin necesidad de ajustar el modelo de difusión.

4. Resultados Experimentales

Los autores evaluaron ORCA en tres conjuntos de datos estándar: DeepMind Control (DMC), MetaWorld y Adroit.

• Rendimiento General: ORCA superó consistentemente a todos los baselines, incluyendo modelos agnósticos a la tarea (CLIP, VC-1) y otros métodos basados en difusión (SCR).
  • En DMC, ORCA obtuvo una puntuación media normalizada de 74.3, superando a SCR (68.3) y VC-1 (53.1).
  • En MetaWorld, alcanzó un 95.2% de éxito medio, superando a SCR (90.4%) y VC-1 (84.2%).
  • En Adroit (manipulación dextrómana), obtuvo un 65.3% de éxito medio, superando a SCR (58.0%) y VC-1 (47.3%).
• Análisis de Componentes:
  • La combinación de task prompts y visual prompts fue superior al uso individual de cualquiera de los dos.
  • Se demostró que el ajuste fino completo de modelos como SCR o VC-1 en datos de imitación limitados degrada drásticamente el rendimiento (caída de >80% en algunos casos), mientras que ORCA mantiene la generalización.
• Visualización: Los mapas de atención mostraron que ORCA logra una "grounding" precisa, enfocándose en las partes relevantes del robot y los objetos objetivo, a diferencia de los prompts de texto que a menudo generan ruido o atención en fondos irrelevantes.

5. Significado e Impacto

Este trabajo es significativo porque:

• Cambia el paradigma de condicionamiento: Sugiere que para el control robótico, las condiciones deben ser visuales y dinámicas, no solo textuales.
• Puente entre Generación y Control: Establece una conexión robusta entre los potentes modelos generativos de difusión (entrenados en internet) y la robótica, sin requerir costosos datos de interacción robótica para preentrenar el modelo base.
• Eficiencia Computacional: Ofrece una vía para obtener representaciones visuales adaptadas a la tarea con un costo computacional bajo (solo aprenden los prompts), evitando los problemas de sobreajuste del ajuste fino completo.
• Futuro: Abre la puerta a utilizar modelos de difusión como extractores de características superiores para agentes de visión-acción (VLA) especializados, complementando o reemplazando a los encoders visuales tradicionales.

En resumen, ORCA demuestra que la clave para utilizar modelos de difusión en robótica no es describir la tarea con texto, sino aprender a "preguntar" al modelo visualmente y contextualmente para cada frame, logrando así una adaptación superior a la tarea.