Resumen Técnico: Optimización de Trayectorias Amortizada para MPC Residual mediante Diferenciación de Contacto Implícita

Planteamiento del Problema

La simulación diferenciable es esencial para la optimización y el aprendizaje de trayectorias ricas en contactos, ya que expone las sensibilidades locales de los resultados de la tarea respecto a los controles. Sin embargo, la escalabilidad de este flujo de trabajo se ve obstaculizada por el coste computacional de diferenciar los resolvedores de contacto iterativos. Los enfoques existentes enfrentan un compromiso fundamental entre convergencia y paralelismo:

1. Diferencias Finitas (FD): Costoso y sensible al tamaño del paso.
2. Diferenciación Automática (AD) Desenrollada: Diferencia todo el rastro del resolvedor. Aunque es preciso, almacena una traza de computación creciente para cada iteración del resolvedor. El endurecimiento de la solución de contacto (aumentando las iteraciones) consume significativamente más memoria, lo que limita el número de trayectorias paralelas que se pueden procesar.
3. Derivaciones de Sensibilidad KKT: Requieren derivaciones intrincadas y ensambladas a mano, específicas de cada resolvedor.

El desafío central es que la AD de modo inverso estándar diferencia el rastro ejecutado del resolvedor en lugar de la solución convergida. Esto crea un cuello de botella de memoria donde el aumento del esfuerzo del resolvedor (para mejorar la precisión) reduce directamente la capacidad de lote (batch), impidiendo una optimización eficiente de todo el horizonte.

Metodología

Los autores proponen un marco que desacopla el resolvedor hacia adelante de su derivada hacia atrás utilizando el Teorema de la Función Implícita (IFT).

1. Diferenciación Implícita vía IFT

En lugar de desenrollar las iteraciones del resolvedor, el método diferencia el residuo de estacionariedad de la solución suavizada y regularizada del problema de contacto.

• Formulación: El problema de contacto en MuJoCo (MJX) se formula como una optimización convexa regularizada. La solución $a^\star$ se define como la raíz del residuo de estacionariedad $F(a, \theta) = 0$, donde $\theta$ representa los parámetros del sistema (estado, controles, fricción, etc.).
• Derivación: Al aplicar el IFT al residuo $F$, se deriva la derivada de la solución con respecto a los parámetros como:
  $$\frac{da^\star}{d\theta} = -F_a^{-1} F_\theta$$
  donde $F_a$ es el Jacobiano del residuo con respecto a la aceleración.
• Implementación: El método utiliza jax.lax.custom_root para implementar este paso hacia atrás. Requiere solo una resolución lineal transpuesta y un producto vector-Jacobiano. Crucialmente, no almacena los estados intermedios de las $K$ iteraciones del resolvedor utilizadas en el paso hacia adelante.
• Alcance: La derivada es local a regiones de contacto suave fijas y puede cambiar a través de transiciones de conjuntos de contacto o modos de fricción, pero dentro de un régimen fijo, proporciona gradientes exactos de la solución convergida.

2. Destilación de Optimizador para MPC Residual

Para aprovechar la escalabilidad de la derivada implícita, los autores introducen un flujo de destilación de optimizador para el Control Predictivo basado en Modelo (MPC):

• Profesor (Teacher): Un iLQR (Regulador Lineal Cuadrático Iterativo) de horizonte completo y por lotes utiliza la dinámica de MJX diferenciable implícitamente para generar trayectorias de alta calidad ricas en contactos.
• Estudiante (Student): Una política de red neuronal es entrenada para imitar las acciones del profesor iLQR de horizonte completo.
• Despliegue (Controlador Híbrido): En tiempo de ejecución, el sistema utiliza un enfoque de MPC Residual:
  • La política destilada proporciona la acción nominal de largo horizonte ($\bar{u}_t$).
  • Un iLQR residual de corto horizonte calcula una corrección local ($r_t$) y las ganancias de retroalimentación basadas en el modelo de contacto diferenciable.
  • El control final es $u_t = \text{sg}[\pi_\theta(o_t)] + r_t$, donde $\text{sg}$ denota una operación de parada de gradiente (stop-gradient), evitando que la política sea actualizada por el optimizador de corto horizonte durante el despliegue.

Contribuciones Clave

1. Paso Atrás Implícito Asistido por AD: Los autores introducen un paso hacia atrás para el residuo de estacionariedad del solve de contacto regularizado de MJX. Este enfoque deja el resolvedor hacia adelante sin cambios, evita sistemas KKT ensamblados a mano y elimina el almacenamiento de memoria dependiente de las iteraciones del resolvedor requerido por la AD desenrollada.
2. Validación y Escalabilidad: El método se valida contra diferencias finas centrales y AD desenrollada. Los resultados muestran que el IFT iguala la precisión de la AD desenrollada con un alto número de iteraciones, mientras supera sustancialmente a esta última en memoria temporal compilada.
   • Escalabilidad de Memoria: La memoria del IFT permanece casi constante a medida que aumentan las iteraciones del solver (cambiando <4% de 1 a 10 iteraciones), mientras que la AD desenrollada crece 10.6×.
   • Escalabilidad de Contacto/Dimensión: El IFT utiliza 20× menos memoria con 256 contactos activos y 6× menos con 16 contactos y 96 Grados de Libertad (DoF) en comparación con la AD desenrollada.
3. MPC Residual Guiado por Política: El artículo demuestra un flujo de destilación donde las acciones optimizadas de horizonte completo se destilan en una política. Esta política guía un iLQR residual de corto horizonte, permitiendo la ejecución exitosa en tareas complejas donde el iLQR de corto horizonte estándar falla.

Resultados Experimentales

El marco fue evaluado en tres tareas ricas en contactos: Finger (rotar un spinner), Franka (empujar una caja) y Unitree (aceleración de un cuadrúpedo).

• Optimización de Trayectorias: En la optimización por lotes de horizonte completo, el iLLQR utilizando la derivada implícita logró costos de ejecución medios más bajos y rangos de percentiles más estrechos en comparación con los métodos de primer orden (Adam) en todas las tareas.
• Rendimiento de la Destilación:
  • El controlador híbrido (Política + iLQR Residual de Corto Horizonte) superó significativamente al iLQR de corto horizonte estándar.
  • Con un horizonte de planificación de 6 pasos, el enfoque híbrido mejoró las tasas de éxito entre 28 y 98 puntos porcentuales sobre el iLQR estándar en las tres tareas.
  • Franka: El iLQR de corto horizonte estándar logró un éxito cercano a cero, mientras que el controlador híbrido tuvo éxito con un lookahead sustancialmente menor. La política proporcionó la estrategia de empuje de largo horizonte, mientras que el iLCR se encargó de las correcciones de contacto locales.
  • Finger & Unitree: La política permitió al sistema tener éxito con horizontes de anticipación (lookahead) más cortos de lo requerido por el iLQR puro. En ciertos regímenes (ej. Finger), un controlador de solo política superó a los de horizonte residual muy corto, ya que el optimizador residual podía perturbar una política que ya era correcta.

Significancia y Reivindicaciones

El artículo afirma que la diferenciación implícita basada en residuos hace que las derivadas de contacto de MJX con tolerancia convergente sean prácticas a escala de aceleradores. Al eliminar el cuello de botella de memoria asociado con las iteraciones del solver, el método permite:

1. Convergencia más Estricta: Los resolvedores pueden ejecutarse a mayor precisión sin sacrificar el paralelismo de lote.
2. Supervisión de Horizonte Completo: La capacidad de ejecutar iLQR de horizonte completo por lotes como un profesor en línea, lo cual era previamente inviable debido a las restricciones de memoria.
3. Amortización Efectiva: El controlador híbrido resultante logra tasas de éxito sustancialmente mayores que el iLQR estándar con horizontes de planificación cortos, amortizando efectivamente la computación costosa específica de la tarea a través de las condiciones iniciales y las metas.

Los autores liberan la implementación como código abierto, proporcionando un patrón modular para la simulación de contacto diferenciable que separa el resolvedor hacia adelante de su derivada sin requerir derivaciones analíticas específicas del solver.