DynaWM: Dynamics-Aware Distillation with World Model and Momentum Targets for Smooth Locomotion over Continuous Stairs

Resumen Técnico: DynaWM: Destilación Consciente de la Dinámica con Modelo de Mundo y Objetivos de Momento para una Locomoción Suave sobre Escaleras Continuas

Planteamiento del Problema
Si bien los robots bípedos con ruedas han demostrado éxito en la navegación de pendientes y obstáculos de un solo escalón mediante el control predictivo basado en modelos y el aprendizaje por refuerzo profundo, navegar por escaleras largas y continuas sigue siendo un desafío significativo. Las soluciones actuales que dependen de marcos de destilación profesor-alumno enfrentan tres limitaciones principales:

1. Representaciones de la Dinámica Debilitadas: Los codificadores del profesor optimizados únicamente mediante gradientes de política tienden a codificar solo características del terreno directamente relevantes para las recompensas inmediatas, ignorando la información geométrica más amplia esencial para el modelado de la dinámica del suelo.
2. Colapso Dimensional: En el entrenamiento concurrente profesor-alumno, las actualizaciones rápidas de las representaciones del profesor pueden causar que el codificador del alumno colapse en representaciones triviales (ajustándose solo a las predicciones medias) en lugar de capturar características esenciales del terreno.
3. Opacidad: El espacio latente en los marcos existentes actúa como una "caja negra", lo que dificulta la verificación de la verdadera capacidad de modelado de la altura del terreno o garantiza la interpretabilidad durante el entrenamiento.

Estos problemas resultan en una falta de suavidad en el movimiento y de adaptabilidad al terreno cuando los robots se enfrentan a escaleras continuas, lo que dificulta un despliegue sim-to-real fiable.

Metodología: DynaWM
Los autores proponen DynaWM (Destilación consciente de la dinámica con modelo de mundo y objetivos de momento), un marco de aprendizaje de representaciones diseñado para mejorar la codificación del terreno y estabilizar la transferencia de conocimiento. El marco consta de tres módulos interconectados:

1. Modelo de Mundo como Regularizador:
   Para imponer la conciencia de la dinámica hacia adelante, se introduce un modelo de mundo como un regularizador para el codificador del profesor. El codificador del profesor ($E_{\theta_t}$) genera variables latentes ($z_t^t$) que, combinadas con observaciones privilegiadas, se alimentan a un modelo de mundo ($W_\omega$) para predecir el siguiente estado privilegiado ($\hat{s}_{t+1}$).

   • Mecanismo: Se minimiza una pérdida de predicción ($L_{pred}$) para entrenar el modelo de mundo y propagar los gradientes de vuelta al codificador del profesor.
   • Efecto: Esto obliga al codificador a retener información esencial para la dinámica hacia adelante (por ejemplo, la geometría del terreno y la fricción) en lugar de descartarla como ruido irrelevante para las recompensas inmediatas. El modelo de mundo se implementa como una red residual profunda con bloques Dense-LayerNorm-Swish para capturar dependencias temporales de largo alcance.

2. Codificador de Objetivo de Momento para una Destilación Estable:
   Para abordar la no estacionariedad del codificador del profesor durante el entrenamiento concurrente, el marco emplea un codificador de objetivo de momento inspirado en BYOL (Bootstrap Your Own Latent).

   • Mecanismo: En lugar de que el alumno realice una regresión directa hacia el profesor en línea que cambia rápidamente, este aprende de un codificador de objetivo ($E_{\hat{\theta}_s}$) que mantiene un promedio móvil exponencial (EMA) de sus propios pesos del alumno. Se añade una cabeza de predictor al alumno en línea para evitar el colapso.
   • Efecto: Esto proporciona objetivos de destilación consistentes y estables, evitando el colapso dimensional y permitiendo que el alumno aprenda representaciones robustas conscientes de la dinámica basándose únicamente en la historia propioceptiva.

3. Diseño de Recompensa y Entrenamiento:
   La función de recompensa se aumenta con restricciones específicas (rango de articulación, distancia de la distancia entre ejes, límites de posición de la base y equilibrio de las piernas) para mitigar problemas como la dependencia excesiva de la cadera y la oscilación del cuerpo observados en el ascenso de escaleras. El entrenamiento utiliza un paradigma de Profesor-Alumno Concurrente (CTS) con PPO, incorporando la randomización de dominio a través de tipos de terreno (plano, pendientes, escalones, irregulares) y parámetros físicos (masa, fricción, ganancias PD) para cerrar la brecha sim-to-real.

Contribuciones Clave

• Regularización por Modelo de Mundo: La introducción de un modelo de mundo para imponer la conciencia de la dinámica hacia adelante, asegurando que las representaciones aprendidas capturen la geometría y la dinámica del terreno en lugar de solo maximizar las recompensas.
• Destilación Estable mediante Objetivos de Momento: El uso de un codificador de objetivo de momento para proporcionar objetivos de destilación estables, previniendo eficazmente el colapso dimensional causado por las actualizaciones no estacionarias del profesor.
• Evaluación Transparente: La aplicación de Análisis de Componentes Principales (PCA) y métricas cuantitativas (PC-Correlation, Linear Probing $R^2$, CCA) para verificar que los codificadores capturan jerárquicamente la altura del terreno, ofreciendo una visión transparente de la calidad de la representación.

Resultados Experimentales
El método fue evaluado en simulación (IsaacGym) y en un robot bípedo con ruedas JiaRan en el mundo real.

• Calidad de la Representación: Las visualizaciones de PCA y las métricas cuantitativas mostraron que el método propuesto logró una PC-Correlation de 0.96 para la altura del terreno en el codificador del profesor y de 0.87 en el codificador del alumno, superando significativamente a las ablaciones sin el modelo de mundo o los objetivos de momento (que sufrieron representaciones entrelazadas o colapso).
• Desempeño en Simulación: En escaleras continuas diversas (variando anchura y altura hasta 20 cm), DynaWM logró tasas de éxito de hasta el 100% para escalones de 26 cm de ancho y del 89.3% para escalones anchos de 20 cm de alto, superando a las líneas de base como CTS y RMA.
• Suavidad del Movimiento: El método demostró una suavidad de movimiento superior (menor jerk y aceleración de la base) y un consumo de energía comparable a CTS.
• Despliegue en el Mundo Real: En experimentos en el mundo real sobre terrenos no vistos, incluyendo escaleras irregulares con protuberancias y escalones continuos que superan las alturas de entrenamiento, el robot logró una tasa de éxito del 100% (5/5 ensayos) con un movimiento significativamente más suave que la línea de base CTS.

Significancia y Reivindicaciones
El artículo afirma que DynaWM cierra con éxito la brecha entre el aprendizaje privilegiado y el control desplegable para terrenos complejos. Al asegurar que el codificador capture la geometría completa del terreno mediante la regularización consciente de la dinámica y estabilizar el proceso de destilación, el método permite que los robots bípedos con ruedas atraviesen escaleras largas y continuas con alta adaptabilidad y suavidad. Los autores sostienen que este enfoque supera las limitaciones de los marcos profesor-alumno actuales, que a menudo fallan en mantener la suavidad del movimiento y la adaptabilidad al terreno en escenarios de escaleras continuas. El trabajo demuestra que el modelado explícito de la dinámica hacia adelante y la generación de objetivos estables son críticos para una transferencia sim-to-real robusta en la locomoción de piernas con ruedas.