PocketDP3: Efficient Pocket-Scale 3D Visuomotor Policy

Resumen Técnico: PocketDP3: Política Visomotora 3D Eficiente de Escala de Bolsillo

Declaración del Problema
Los avances recientes en la manipulación robótica han aprovechado las políticas de difusión basadas en visión 3D para aprender habilidades complejas a partir de datos. Sin embargo, los autores identifican un desajuste arquitectónico crítico en los modelos actuales de vanguardia (SOTA), como DP3. Estos modelos suelen emparejar un codificador de nubes de puntos ligero y eficiente con un decodificador de base excesivamente grande, a menudo una U-Net condicional. Aunque las U-Nets son adecuadas para la síntesis de imágenes de alta resolución donde se requiere una reconstrucción de píxeles finamente detallada, los autores argumentan que este diseño es ineficiente para el control visomotor 3D. En este contexto, el codificador ya produce una representación de la escena compacta y rica en información, y el objetivo de predicción es una trayectoria de acción de baja frecuencia. Por consiguiente, el masivo decodificador introduce una redundancia de parámetros sustancial, aumenta la huella de memoria y obstaculiza el despliegue en tiempo real en robots con recursos limitados sin proporcionar ganancias proporcionales en el rendimiento del control. Además, el muestreo de difusión estándar requiere muchos pasos de eliminación de ruido iterativos, lo que crea cuellos de botella de latencia, necesitando a menudo técnicas complejas de destilación de consistencia para acelerar la inferencia.

Metodología
El artículo propone PocketDP3, una política de difusión 3D de escala de bolsillo diseñada para resolver el desequilibrio arquitectónico y acelerar la inferencia.

• Arquitectura (Diffusion Mixer - DiM): En lugar de una pesada U-Net condicional, PocketDP3 emplea un decodificador Diffusion Mixer (DiM) ligero construido sobre bloques MLP-Mixer. La arquitectura conserva el codificador de nubes de puntos eficiente de DP3 (utilizando Farthest Point Sampling y MLPs para generar una representación 3D compacta), pero reemplaza el decodificador con $K$ bloques DiM apilados.

  • Mecanismo de Fusión: Los bloques DiM realizan una fusión de información eficiente alternando entre la mezcla temporal (agregando información a través de la dimensión temporal de la trayectoria) y la mezcla de canales (agregando a través de las dimensiones de las características).
  • Condicionamiento: El modelo incorpora representaciones de la escena y estados del robot mediante una proyección lineal e inyecta esta información de condicionamiento en los bloques DiM utilizando capas FiLM (Feature-wise Linear Modulation) combinadas con conexiones residuales con compuerta (gated residual connections).
  • Eficiencia de Parámetros: Los autores proponen dos configuraciones, "Tiny" (0.53M de parámetros) y "Base" (1.73M de parámetros), que representan menos del 1% de los parámetros utilizados por métodos previos como DP3 (~250M de parámetros).

• Estrategia de Inferencia: El método utiliza muestreo de estilo DDIM. Notablemente, los autores demuestran que PocketDP3 soporta inferencia de dos pasos (2 evaluaciones de función) sin requerir destilación de consistencia adicional u objetivos de entrenamiento especializados (por ejemplo, Difusión de un solo paso o Modelos de Consistencia). Esto reduce significativamente la latencia de inferencia manteniendo el rendimiento.

• Entrenamiento: El modelo se entrena utilizando un objetivo de predicción $x_0$, minimizando el error cuadrático medio entre la trayectoria limpia predicha y la trayectoria experta real, condicionada por la observación y el paso de tiempo (timestep).

Contribuciones Clave

1. Perspectiva Arquitectónica: El artículo identifica que emparejar codificadores 3D compactos con decodificadores U-Net masivos es una ineficiencia generalizada en las políticas de difusión 3D, lo que conduce al desperdicio de parámetros.
2. Arquitectura PocketDP3: La introducción de una política de escala de bolsillo que reemplaza la U-Net con un Mezclador de Difusión (DiM), permitiendo una fusión temporal y de canales efectiva con un presupuesto de parámetros mínimo.
3. Muestreo Rápido Libre de Destilación: La demostración de que es posible lograr un muestreo de dos pasos de alto rendimiento sin destilación de consistencia, simplificando el flujo de entrenamiento y permitiendo el despliegue de baja latencia.
4. Validación Empírica: Una evaluación exhaustiva que muestra que PocketDP3 iguala o supera el rendimiento SOTA reduciendo drásticamente el tamaño del modelo y el tiempo de inferencia.

Resultados Experimentales
El método fue evaluado en tres entornos de simulación (RoboTwin2.0, Adroit y MetaWorld) y experimentos en el mundo real.

• Rendimiento en Simulación:

  • RoboTwin2.0: PocketDP3-Base logró una tasa de éxito promedio del 71.6%, superando al SOTA anterior (DP3) en un 20.8%. PocketDP3-Tiny también superó a DP3 por un 15.2%.
  • Adroit y MetaWorld: PocketDP3-Base logró una tasa de éxito promedio del 77.4%, superando a DP3 (73.0%) y FlowPolicy (72.6%).
  • Estudios de Ablación: Las comparaciones con un "Vanilla-MLP" (que falló en converger) y un "Vanilla-U-Net" reducido (que tuvo un rendimiento inferior) confirmaron que el mecanismo de fusión específico del bloque DiM es crítico para el rendimiento bajo presupuestos estrictos de parámetros.
  • Pasos de Inferencia: El estudio de ablación sobre el Número de Evaluaciones de Función (NFE) mostró que el rendimiento se estanca después de 2 pasos; aumentar a 5 o 10 pasos no produjo ganancias significativas.

• Eficiencia:

  • Tamaño del Modelo: PocketDP3-Base utiliza 1.73M de parámetros, en comparación con los 255.1M de DP3 (una reducción de >99%).
  • Latencia: Con 2 pasos de inferencia, PocketDP3-Base logra una latencia de 4.80ms en una RTX 5880 Ada, significativamente más rápido que DP3 (51.4ms con 10 pasos) y comparable a FlowPolicy (7.04ms con 1 paso) pero con dos órdenes de magnitud menos de parámetros.

• Despliegue en el Mundo Real:

  • Probado en un robot AgileX Piper con una cámara RealSense D455.
  • PocketDP3-Base demostró una fuerte transferencia de simulación a realidad (sim-to-real), superando a la línea base DP3 en un 15.3% en promedio en tareas como colocación de objetos y apilamiento de bloques.

Significancia y Reivindicaciones
El artículo afirma que PocketDP3 ofrece un camino práctico hacia el control visuomotor 3D eficiente. Al desafiar la suposición de que los decodificadores grandes son necesarios para las políticas basadas en difusión cuando la representación de entrada ya está comprimida, los autores demuestran que "más pequeño es mejor" en este contexto específico. El método logra un rendimiento de vanguardia con <1% de los parámetros de los métodos anteriores y permite la inferencia en tiempo real sin un complejo entrenamiento de destilación. Los autores posicionan este trabajo como un paso hacia el despliegue de políticas robóticas sofisticadas en hardware con recursos limitados.

Limitaciones y Trabajo Futuro
Los autores reconocen que la arquitectura DiM puede enfrentar un "cuello de botella espacial" para tareas que requieren un razonamiento geométrico local de grano fino (por ejemplo, inserción submilimétrica) debido a la falta de sesgos espaciales alineados con los píxeles inherentes a las U-Nets. También señalan que el mecanismo subyacente que permite la inferencia de dos pasos de alta calidad sigue sin estar claro y requiere mayor investigación. Se sugiere que el trabajo futuro se centre en el diseño de codificadores de condicionamiento más eficientes y ricos en información, así como en el estudio sistemático de las compensaciones entre la capacidad de condicionamiento, la complejidad de la tarea y la escala del decodificador.