Paper & Results¶

SoRoMoX is presented in “SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models”.

Paper overview¶

The paper's primary contribution is a fully numerical, JIT-compilable Python/JAX implementation of control-oriented models for articulated and continuum soft robots. SoRoMoX implements articulated soft-robot, piecewise-constant strain (PCS), and geometric variable strain (GVS) models through one interface for kinematics, dynamics, energies, Jacobians and derivatives, actuation maps, and forward dynamics.

Because the numerical core is JAX-native, these model computations can be JIT-compiled, automatically differentiated with respect to states, inputs, and physical parameters, vectorized, and executed on CPUs, GPUs, and TPUs. Model-based controllers and rendering backends are supporting contributions built around this model layer.

Model evaluation¶

Sequential CPU rollouts¶

The paper reports wall-clock time for three seconds of simulated motion on an Intel Core Ultra 7 165H CPU with 16 GB RAM. SoRoMoX is up to 27.0× faster than SoRoSim across the four cases:

The case labels describe the strain-based models in SoRoMoX and SoRoSim, which use two piecewise elements. PyElastica uses a discrete Cosserat-rod model with 10–100 elements. SoRoSim uses MATLAB's adaptive ode45 solver; SoRoMoX uses Tsit5 with a fixed 0.1 ms step. PyElastica uses the second-order symplectic PositionVerlet solver with the same fixed step size.

The paper also compares SoRoSim and SoRoMoX tip positions at identical saved timestamps, without temporal or phase alignment:

These RMSE values correspond to 0.14–1.20% of the 0.6 m beam length. The largest instantaneous tip-position difference is 14.14 mm for spatial PCS. See the CPU benchmark reproduction instructions for the simulation scripts, saved trajectories, and plotting command.

Batched GPU rollouts¶

The exact scaling depends on the model, number of links or segments, hardware, and discretization. In this FP64 run, automatic backend selection resolves to Warp for GVS. At batch size 1,024, GVS reaches 679.7× the single-environment throughput for one segment and 133.6× for 32 segments. Throughput increases monotonically through batch size 1,024 for the GVS cases with 1, 2, 4, 8, or 16 segments; the 32-segment case peaks at batch size 512 and is 1.2% lower at 1,024.

Application case studies¶

1. Static-equilibrium system identification¶

This study differentiates through a two-link GVS model to identify Young's modulus, Poisson's ratio, and mass density from static marker measurements of a conical tendon-driven arm. Optimizing the physical parameters reduces the aggregate marker-position RMSE from 56.9 mm to 19.3 mm, an improvement of about 66%.

2. Residual-force learning¶

Using the same differentiable static-equilibrium model, the study first estimates input-dependent residual generalized forces and then trains a neural network to predict them. The optimized and learned residuals reduce the overall marker RMSE to 6.6 mm and 6.9 mm, respectively.

3. Model-based control¶

This study exercises the quantities exposed by the common model interface. A one-segment spatial PCS robot compares model-free PD/PID control with potential, feedforward, and computed-torque compensation; a fully actuated two-segment PCS robot tracks a spherical figure-eight pose trajectory with operational-space impedance control. The operational-space example reports 2.84 mm position RMSE and 2.29° orientation RMSE.

Watch the MP4 · View the animated GIF

4. Control-gain optimization¶

Gradients are propagated through full closed-loop rollouts while six gain initializations are evaluated in parallel. For actuation-space and operational-space control, the optimized losses decrease by 47% and 99%, respectively, relative to their initial median values.

5. Safety-constrained control¶

The safety study uses differentiable PCS kinematics and dynamics to formulate high-order control Lyapunov and control barrier function constraints. The safety-unaware controller reaches about 33.5 N of pairwise normal force. The HOCBF-constrained controller stays at or below the prescribed 5 N limit while accepting a larger final goal distance.

The video shows the HOCBF-constrained controller approaching the green goal among the purple obstacles.

Watch the MP4 · View the animated GIF

6. Parallel reinforcement learning¶

The final study vectorizes tendon-driven PCS environments on a GPU for PPO training. The trained policy reaches a 6.19 mm final mean tracking error and a 91.4% step-wise success rate. The reported 4× and 7× training speedups at 256 and 512 environments relative to the CPU PyElastica baseline jointly reflect model formulation, JIT compilation, hardware acceleration, and rollout parallelization.

The video shows the trained policy tracking the green target in one environment.

Watch the MP4 · View the animated GIF

SoRoMoX can roll out many soft-robot environments in parallel, accelerating the experience collection that dominates reinforcement-learning training. The 20 second animation below visualizes a trained policy acting simultaneously in 64 independently simulated environments, with each robot rendered as a swept surface.

Watch the MP4 · View the animated GIF

Reproducing the results¶

The paper_results/ directory contains experiment entry points, canonical outputs, and workflow-specific notes. Install its dependencies from an editable checkout:

python -m pip install -e ".[paper_results]"

Start with the README nearest to the experiment you want to reproduce. Record the SoRoMoX version, hardware, solver configuration, and random seed with new results. For package setup, see Installation.

If you publish work based on these models or results, follow the Citation guide.