SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models

2026-08-06

AI Summary

Existing modeling approaches for soft robots lack differentiability, GPU parallelization, and control-oriented workflows, hindering their applicability to advanced learning and control tasks. This work proposes a unified differentiable numerical framework grounded in Cosserat rod theory, supporting both piecewise-constant and variable-strain rod models while efficiently computing inertial matrices, generalized forces, and their Jacobian derivatives. It presents the first end-to-end differentiable simulation of rod/strain models on GPUs, enabling full gradient propagation with respect to states, inputs, and parameters. Implemented in Python using JAX, the framework leverages JIT compilation and automatic differentiation, achieving an 18.1× speedup over single-threaded CPU execution and a 234.6× increase in GPU throughput. These advances significantly enhance performance in system identification, residual learning, force-controlled tracking, gain optimization, safety-constrained control, and accelerate reinforcement learning training by up to 7×.

Abstract

Reduced-order models based on Cosserat-rod theory are now well established, and modeling theory is no longer the primary bottleneck in soft-robot control. Their implementations, however, do not support the differentiable, GPU-parallel, and control-oriented workflows that underpin advanced rigid-robotics applications. Here, we fill this gap with SoRoMoX (Soft Robot Models in JAX), a fully numerical, JIT-compilable Python/JAX framework. SoRoMoX implements articulated, Piecewise Constant Strain, and Variable Strain models through a unified, control-ready interface that provides inertia matrices, gravitational and elastic forces, Jacobians, and their derivatives. To our knowledge, it is the first rod/strain-based soft-robot modeling framework that runs directly on GPUs and is end-to-end differentiable with respect to states, inputs, and parameters. Sequential CPU rollouts are up to 18.1x faster than state-of-the-art alternatives, while GPU-parallel rollouts increase throughput by up to 234.6x. This performance enables workflows that were previously impractical or impossible: static-equilibrium system identification with 66% lower marker RMSE; residual-force learning with a further 64% reduction; computed-torque tracking with RMSE reduced by a factor of approximately 500 relative to model-free PD; control-gain optimization with up to 62% lower loss than untuned gains; safety-constrained control using high-order control barrier functions to keep the peak contact force within a prescribed 5 N bound, compared with 33.5 N without the safety constraint; and reinforcement-learning policy training up to 7x faster than a CPU PyElastica discrete-rod baseline through massively parallel rollouts.

Problem

• soft robot modeling
• differentiable simulation
• GPU parallelization
• control-oriented workflow
• Cosserat rod theory

Innovation

• differentiable simulation
• GPU parallelization
• soft robot modeling
• Cosserat rod theory
• JAX framework