Batched Differentiable Rigid Body Dynamics in PyTorch for GPU-Accelerated Robot Learning

Resumen Técnico: Dinámica de Cuerpos Rígidos por Lotes en PyTorch (bard)

Planteamiento del Problema
A medida que el control robótico transiciona hacia el aprendizaje por refuerzo (RL) a gran escala con computación de dinámica en el bucle, el campo se enfrenta a un cuello de botella crítico de rendimiento. Los pipelines de entrenamiento actuales, como los de Isaac Lab y MuJoCo Playground, dependen de PyTorch para la ejecución por lotes en GPU. Sin embargo, la biblioteca de dinámica estándar, Pinocchio, está limitada por la CPU y es secuencial. Esta arquitectura impone una sobrecarga de transferencia de datos CPU–GPU lineal para cada evaluación y deja los núcleos de la GPU ociosos durante las traversales de árboles secuenciales. Si bien los simuladores de GPU existentes como Brax y MuJoCo MJX ofrecen un alto rendimiento, están ligados a ecosistemas ajenos a PyTorch (JAX) o funcionan como simuladores monolíticos que no exponen operadores diferenciables independientes (por ejemplo, Jacobianos o matrices de masa) requeridos para funciones de pérdida personalizadas o recompensas basadas en la física. Por el contrario, las bibliotecas diferenciables compatibles con PyTorch existentes, como ADAM, sacrifican el rendimiento de la GPU en favor de la generalidad del backend, fallando al no explotar la fusión de kernels o evitar el flujo de control dependiente de los datos que rompe el trazado de grafos del compilador de PyTorch.

Metodología
Los autores presentan bard (Batched Articulated Rigid-body Dynamics), una biblioteca de puro PyTorch que implementa los algoritmos de la dinámica de cuerpos rígidos de Featherstone (FK, RNEA, CRBA, ABA y Jacobianos). La biblioteca está diseñada para la evaluación por lotes en GPU y la compatibilidad total con torch.compile. Las ganancias de rendimiento se logran mediante cuatro estrategias de diseño específicas:

1. Evaluación Diferida por Niveles (Tiered Lazy Evaluation): En lugar de calcular todas las cantidades cinemáticas en cada llamada, bard divide el estado cinemático en tres niveles protegidos por banderas de validez: (1) transformaciones locales padre-hijo, (2) adjuntos espaciales (requeridos solo para algoritmos de dinámica) y (3) poses en el marco del mundo (requeridos solo para consultas de FK/Jacobiano). Esto asegura que un bucle de RL que solo llama a RNEA evite el costo computacional de la propagación al marco del mundo.
2. Transformaciones de Articulación Libres de Multiplicación de Matrices: Para eliminar la sobrecarga de lanzamiento de miles de pequeñas multiplicaciones de matrices de $4 \times 4$ en la GPU, la biblioteca precalcula las constantes de rotación de Rodrigues en el momento de la construcción del modelo. En tiempo de ejecución, las transformaciones padre-hijo se recuperan utilizando únicamente multiplicaciones escalar-matriz con $\sin(q_i)$ y $\cos(q_i)$, lo que permite al compilador de PyTorch fusionar operaciones a través de todos los joints y elementos del lote en un solo kernel.
3. Propagación de Árbol en Paralelo por Niveles: Aunque los algoritmos de Featherstone son inherentemente secuenciales, bard reorganiza las traversales de árboles en operaciones por lotes paralelas por niveles. Al agrupar los nodos en la misma profundidad del árbol, la biblioteca reduce los pasos secuenciales de $n_b$ (número de cuerpos) a $d$ (profundidad del árbol) pasos por lotes. Aunque los pasos hacia atrás (acumulación de fuerzas) siguen siendo secuenciales debido a las dependencias de datos, todos los $B$ entornos se procesan simultáneamente en cada nodo.
4. Construcción Directa de Columnas del Jacobiano: Para articulaciones de un solo grado de libertad (single-DOF), la biblioteca evita la construcción de matrices adjuntas completas de $6 \times 6$. En su lugar, calcula directamente las columnas del Jacobiano utilizando offsets precalculados y poses disponibles en el marco del mundo, reemplazando las construcciones de adjuntos con productos vectoriales eficientes y multiplicaciones matriz-vector.

La biblioteca utiliza un diseño de memoria de primero-el-lote (batch-first) y almacena las transformaciones espaciales transpuestas de forma contigua para evitar el desperdicio de la línea de caché. Soporta tanto el modo eager (saltándose el seguimiento de gradientes para el máximo rendimiento) como el modo funcional (diferenciación automática completa) y es totalmente compatible con torch.compile mediante el uso de flujo de control estático y tensores de índice constantes.

Contribuciones Clave

1. bard: Una biblioteca de dinámica de puro PyTorch que permite la evaluación por lotes en GPU de los algoritmos de Featherstone.
2. Técnicas de Nivel de Sistema: Métodos novedosos para ejecutar computaciones recursivas de árboles en GPUs, incluyendo el almacenamiento en caché de la cinemática, espacios de trabajo preasignados y propagación paralela por niveles.
3. Validación: Pruebas exhaustivas en identificación de sistemas basada en gradientes e integración de RL a gran escala con 4,096 entornos paralelos.

Resultados Experimentales
Los autores evaluaron bard en cinco modelos de robots (que van desde 7 hasta 23 DOF) frente a Pinocchio y ADAM en una GPU NVIDIA H200:

• Precisión Numérica: bard iguala a Pinocchio dentro de los errores de redondeo de punto flotante. Las discrepancias máximas fueron cercanas al epsilon de la máquina ($\approx 10^{-16}$) para FK y Jacobianos, y por debajo de $10^{-13}$ para RNEA y CRBA. ABA mostró discrepancias ligeramente mayores (hasta $10^{-10}$) debido a la acumulación de errores en su estructura de tres pasadas, pero se mantuvo dentro de los límites prácticos.
• Rendimiento (Throughput): Con un tamaño de lote de 4,096, bard logró un rendimiento hasta 64× mayor que Pinocchio para la Cinemática Directa (Forward Kinematics) y 63× para los Jacobianos. Para los algoritmos de dinámica, las aceleraciones oscilaron entre 3× y 7×.
• Identificación de Sistemas: En una tarea de identificación basada en gradientes en un manipulador de 7 DOF con ruido de torque del 5%, bard recuperó con éxito las masas de los eslabones con un error medio de 1.24%.
• Integración de RL: Integrado en un pipeline de entrenamiento AMP de Isaac Lab para un cuadrúpedo de 11 DOF con 4,096 entornos paralelos, bard logró un rendimiento 8.5× mayor que Pinocchio y 2.0× mayor que ADAM.

Significado y Reivindicaciones
El artículo afirma que bard llena un vacío crítico en el ecosistema de la robótica al proporcionar una biblioteca de dinámica modular, de alto rendimiento y nativa para PyTorch. Permite el entrenamiento de RL a gran escala con dinámica en el bucle en una sola GPU, una tarea que anteriormente se veía obstaculizada por cuellos de botella de la CPU o incompatibilidad de ecosistema. Los autores señalan que la ventaja de rendimiento es más pronunciada para las cargas de trabajo dominadas por la cinemática con tamaños de lote grandes. También advierten que torch.compile puede degradar el rendimiento para cargas de trabajo intensivas en dinámica en GPUs de gama baja (por ejemplo, NVIDIA L4) debido al vertido de registros (register spilling), recomendando el modo eager para tales casos. El trabajo demuestra que reorganizar algoritmos seriales para la ejecución paralela en GPU puede generar ganancias de rendimiento significativas sin sacrificar la precisión numérica o la diferenciabilidad.