# SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models

arXiv: https://arxiv.org/abs/2608.06650 · HTML: https://mlsanity.com/paper/2608.06650/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 30 of its 81 references that are in the map.
- First posted on arXiv 6 Aug 2026.

## Builds on (most cited first, 30 of 30)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning](https://mlsanity.com/paper/2108.10470/) (2021)
- [Open3D: A Modern Library for 3D Data Processing](https://mlsanity.com/paper/1801.09847/) (2018)
- [Learning to Walk in Minutes Using Massively Parallel Deep Reinforcement Learning](https://mlsanity.com/paper/2109.11978/) (2021)
- [Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation](https://mlsanity.com/paper/2106.13281/) (2021)
- [Control Barrier Function Based Quadratic Programs for Safety Critical Systems](https://mlsanity.com/paper/1609.06408/) (2016)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [DiffuseBot: Breeding Soft Robots With Physics-Augmented Generative Diffusion Models](https://mlsanity.com/paper/2311.17053/) (2023)
- [Viser: Imperative, Web-based 3D Visualization in Python](https://mlsanity.com/paper/2507.22885/) (2025)
- [Model Based Control of Soft Robots: A Survey of the State of the Art and Open Challenges](https://mlsanity.com/paper/2110.01358/) (2021)
- [SoftZoo: A Soft Robot Co-design Benchmark For Locomotion In Diverse Environments](https://mlsanity.com/paper/2303.09555/) (2023)
- [Soft Robots Modeling: a Structured Overview](https://mlsanity.com/paper/2112.03645/) (2021)
- [Optimistix: modular optimisation in JAX and Equinox](https://mlsanity.com/paper/2402.09983/) (2024)
- [Design and Control of a Bipedal Robotic Character](https://mlsanity.com/paper/2501.05204/) (2025)
- [Elastica: A compliant mechanics environment for soft robotic control](https://mlsanity.com/paper/2009.08422/) (2020)
- [Physics-informed Neural Networks to Model and Control Robots: a Theoretical and Experimental Investigation](https://mlsanity.com/paper/2305.05375/) (2023)
- [Sim-to-Real of Soft Robots with Learned Residual Physics](https://mlsanity.com/paper/2402.01086/) (2024)
- [Discrete Cosserat Approach for Multi-Section Soft Robots Dynamics](https://mlsanity.com/paper/1702.03660/) (2017)
- [A MATLAB Toolbox for Hybrid Rigid Soft Robots Based on the Geometric Variable Strain Approach](https://mlsanity.com/paper/2107.05494/) (2021)
- [Rod models in continuum and soft robot control: a review](https://mlsanity.com/paper/2407.05886/) (2024)
- [An Experimental Study of Model-based Control for Planar Handed Shearing Auxetics Robots](https://mlsanity.com/paper/2310.13889/) (2023)
- [DisMech: A Discrete Differential Geometry-based Physical Simulator for Soft Robots and Structures](https://mlsanity.com/paper/2311.18126/) (2023)
- [Guiding Soft Robots with Motor-Imagery Brain Signals and Impedance Control](https://mlsanity.com/paper/2401.13441/) (2024)
- [Learning Low-Dimensional Strain Models of Soft Robots by Looking at the Evolution of Their Shape with Application to Model-Based Control](https://mlsanity.com/paper/2411.00138/) (2024)
- [Input Decoupling of Lagrangian Systems via Coordinate Transformation: General Characterization and its Application to Soft Robotics](https://mlsanity.com/paper/2306.07258/) (2023)
- [Mastering Contact-rich Tasks by Combining Soft and Rigid Robotics with Imitation Learning](https://mlsanity.com/paper/2410.07787/) (2024)
- [Safe Autonomous Environmental Contact for Soft Robots using Control Barrier Functions](https://mlsanity.com/paper/2504.14755/) (2025)
- [Contact-Aware Safety in Soft Robots Using High-Order Control Barrier and Lyapunov Functions](https://mlsanity.com/paper/2505.03841/) (2025)
- [UMArm: Untethered, Modular, Portable, Soft Pneumatic Arm](https://mlsanity.com/paper/2505.11476/) (2025)
- [A Closed-Form CLF-CBF Controller for Whole-Body Continuum Soft Robot Collision Avoidance](https://mlsanity.com/paper/2603.19424/) (2026)

## Abstract (arXiv)

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.