SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models

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, with Warp kernels accelerating parallel continuum-model execution, and is end-to-end differentiable with respect to states, inputs, and parameters. Sequential CPU rollouts are up to 27.0 times faster than SoRoSim, while GPU-parallel GVS rollouts increase throughput by up to 679.7 times. 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 98% 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 7 times faster than a CPU PyElastica discrete-rod baseline through massively parallel rollouts.

Robotics (cs.RO); Artificial Intelligence (cs.AI)

arXiv:2608.06650 [cs.RO]

(or arXiv:2608.06650v2 [cs.RO] for this version)

https://doi.org/10.48550/arXiv.2608.06650 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.