🤖 AI

SoRoMoX: Fast, Differentiable, and Parallelizable Soft Robot Models

SoRoMoX is a novel, fully differentiable, and GPU-parallelized JAX framework for soft robot modeling that significantly accelerates simulation and enables advanced control workflows—such as system identification, reinforcement learning, and safety-constrained control—by overcoming the computational limitations of existing Cosserat-rod implementations.

Original authors: Maximilian Stölzle, Solange Gribonval, Daniel Feliu-Talegon, Vito Daniele Perfetta, Michele Martini, Chuhan Zhang, Kiwan Wong, Mohammed Tarnini, Anup Teejo Mathew, Federico Renda, Daniela Rus, Cosimo

Published 2026-08-10

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Maximilian Stölzle, Solange Gribonval, Daniel Feliu-Talegon, Vito Daniele Perfetta, Michele Martini, Chuhan Zhang, Kiwan Wong, Mohammed Tarnini, Anup Teejo Mathew, Federico Renda, Daniela Rus, Cosimo Della Santina

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a world where robots aren't made of stiff metal gears and rigid joints, but are instead crafted from squishy, stretchy materials like silicone, rubber, or even living tissue. These are "soft robots," and they are incredible because they can squeeze through tiny cracks, handle fragile objects like eggs without breaking them, and move with a fluid grace that rigid machines can only dream of. However, there's a catch: because they are so squishy, they are incredibly hard to control. Unlike a rigid robot arm that moves in predictable, straight lines, a soft robot bends, twists, and warps in a million different ways at once. It's like trying to write a precise instruction manual for a jellyfish; if you pull one tentacle, the whole body ripples in a complex dance that is mathematically messy to predict. For a long time, scientists have had the right math to describe this squishiness, but the computer programs to actually run that math were slow, clunky, and couldn't easily talk to the advanced learning tools that modern robots use.

This is where a new tool called SoRoMoX steps in to save the day. Think of SoRoMoX as a super-powered, high-speed translator that turns the complex, wiggly math of soft robots into a language that modern computers can understand and execute at lightning speed. The researchers behind this project built a software framework that doesn't just simulate a soft robot; it makes the simulation "differentiable," which is a fancy way of saying the computer can instantly figure out how to tweak the robot's shape or movement to get a better result, just like a video game character learning to jump higher by trying and failing thousands of times in a split second. By running these simulations on powerful graphics cards (GPUs) and doing them in parallel—like having a thousand virtual robots practicing at the same time instead of one—the team showed that they could make soft robot control up to 234.6 times faster than previous methods. This speed unlocks new possibilities, allowing robots to learn complex tasks, avoid bumping into things too hard, and even figure out their own physical properties (like how stiff their body is) just by watching themselves move.

The Squishy Robot Revolution

Soft robots are the darlings of the future because they are safe and adaptable. Imagine a robot arm made of rubber that can gently wrap around a ripe strawberry to pick it, or a snake-like robot that can slither through the rubble of a collapsed building to find survivors. The problem is that controlling something that bends and stretches like a noodle is a nightmare for traditional computer programs. In the world of rigid robots (like the arms on a car assembly line), the math is straightforward: if you turn a motor a certain amount, the arm moves a specific distance. But with soft robots, the whole body deforms, creating a chaotic web of forces.

For years, scientists have developed theories to describe this behavior, often using a concept called "Cosserat-rod theory." Think of this as treating the soft robot not as a solid block, but as a flexible stick that can bend, twist, and stretch. While the math exists to describe this, the software tools to actually use it were stuck in the past. They were slow, couldn't run on modern super-fast computer chips, and didn't play nice with the "learning" algorithms that allow robots to teach themselves. It was like having a Ferrari engine (the math) but trying to drive it with a bicycle chain (the old software).

Enter SoRoMoX: The Speed Demon

The paper introduces SoRoMoX (Soft Robot Models in JAX), a new software framework designed to fix these bottlenecks. The authors built it using JAX, a powerful programming tool that lets computers perform complex math operations incredibly fast, especially when using Graphics Processing Units (GPUs)—the same chips that power video games and AI.

Here is what makes SoRoMoX special:

It's Differentiable: This is the magic ingredient. In simple terms, the software can calculate not just where the robot will go, but how to change its settings to get it exactly where you want it to go. It's like having a GPS that doesn't just tell you the route, but instantly recalculates the best path if you hit a traffic jam, all while you are driving.
It's Parallel: Instead of simulating one robot at a time, SoRoMoX can simulate hundreds or thousands of them simultaneously on a single GPU. This is like having a classroom of students all practicing a dance routine at once, rather than one student practicing alone.
It's Control-Ready: It provides all the specific numbers a robot controller needs—like how heavy the robot is, how gravity pulls on it, and how its joints move—so engineers can build better controllers without reinventing the wheel.

What the Team Discovered

The researchers didn't just build the tool; they put it through a gauntlet of tests to see how it performed compared to the best existing software.

Speed and Efficiency
In their tests, SoRoMoX was a clear winner. When running simulations on a standard computer processor (CPU), it was up to 18.1 times faster than the current state-of-the-art alternative called SoRoSim. But the real magic happened when they switched to a GPU. By running simulations in parallel, they achieved a massive speedup of up to 234.6 times. This means tasks that used to take hours could be done in minutes, or even seconds.

Learning and Optimization
The team used this speed to tackle some tough problems:

Figuring Out the Robot's Body: They tested if the software could help a robot "learn" its own physical properties, like how stiff its material is. By comparing the robot's actual movement to the simulation, they could adjust the model's parameters. This reduced the error in predicting the robot's position by 66%.
Fixing the "Oops" Moments: Even with a perfect model, real robots have small errors. The team used SoRoMoX to teach a neural network to predict these leftover errors (residuals). This further reduced the error by another 64%, making the robot's movements incredibly precise.
Better Control: They tested different control strategies. A standard "model-free" controller (which doesn't know the robot's physics) struggled. But when they used a "computed-torque" controller (which uses the full physics model), the tracking error dropped by a factor of approximately 500. That's a massive leap in precision.
Safety First: One of the most exciting applications was safety. They programmed the robot to reach a target while avoiding hitting a wall too hard. Using a special safety algorithm, they kept the contact force below 5 Newtons (about the weight of a small apple). Without this safety constraint, the robot hit with a force of 33.5 Newtons, which could damage fragile objects.
Reinforcement Learning: Finally, they used the tool to train a robot using reinforcement learning (where the robot learns by trial and error). Because SoRoMoX could run thousands of trials at once, the training was up to 7 times faster than using older, slower methods.

The Bottom Line

SoRoMoX isn't just a faster calculator; it's a bridge that connects the complex, squishy world of soft robotics with the high-speed, data-driven world of modern AI and control. By making soft robot models fast, differentiable, and easy to use, the authors have opened the door for robots that can learn, adapt, and interact with the world in ways that were previously too slow or too difficult to achieve. While the results shown here are based on simulations, the dramatic improvements in speed and accuracy suggest that the next generation of soft robots could soon be learning to move with the grace and intelligence of living creatures.

Technical Summary: SoRoMoX – Fast, Differentiable, and Parallelizable Soft Robot Models

1. Problem Statement

While reduced-order modeling of soft robots based on Cosserat-rod theory (specifically Piecewise Constant Strain [PCS] and Variable Strain [VS] formulations) has matured, a significant bottleneck remains in model deployment. Existing implementations often suffer from three critical limitations that hinder their integration into advanced control and learning pipelines:

Lack of Differentiability: They do not support end-to-end automatic differentiation (AD) with respect to states, inputs, and parameters, preventing their use in gradient-based optimization, system identification, and learning.
Computational Inefficiency: They are typically CPU-bound, legacy-code based (e.g., MATLAB), or closed-source, failing to leverage modern hardware accelerators (GPUs/TPUs) for parallel execution.
Incomplete Interfaces: They often function as "black-box" simulators that provide only rollout trajectories, lacking direct access to essential control-oriented quantities such as inertia matrices, Jacobians, Coriolis terms, and actuation maps required for model-based control and safety filtering.

Consequently, workflows common in rigid robotics—such as large-scale system identification, trajectory optimization, reinforcement learning (RL), and safety-constrained control using Control Barrier Functions (CBFs)—remain impractical or impossible for soft robots.

2. Methodology: The SoRoMoX Framework

The authors introduce SoRoMoX (Soft Robot Models in JAX), a fully numerical, Just-In-Time (JIT) compilable Python/JAX framework designed to bridge this gap.

Core Architecture

Unified Interface: SoRoMoX implements a common API for articulated robots, PCS models, and Geometric Variable Strain (GVS) models. This interface exposes:
- Kinematics: Forward kinematics $g(q, s)$ , Jacobians $J(q, s)$ , and Jacobian time derivatives $\dot{J}(q, \dot{q}, s)$ .
- Dynamics: Inertia matrices $M(q)$ , Coriolis/Centrifugal terms $C(q, \dot{q})$ , gravity $G(q)$ , elasticity $K(q)$ , damping $D$ , and actuation maps $A(q)$ .
- Simulation: Open-loop and closed-loop rollouts integrated with the Diffrax ODE solver.
Computational Stack: Built on JAX, the framework supports:
- Automatic Differentiation (AD): Differentiation with respect to states, inputs, and physical parameters.
- JIT Compilation: Optimization of kernels for CPU, GPU, and TPU execution.
- Vectorization: Massively parallel batched execution (via jax.vmap) to simulate thousands of environments simultaneously.
Model Families:
- Articulated Soft Robots: Serial chains with compliant joints or soft actuators (e.g., McKibben muscles).
- Continuum Robots:
  - PCS: Constant strain within segments, utilizing closed-form kinematics for efficiency.
  - GVS: Variable strain approximated by basis functions, allowing for more complex deformation modes while maintaining a finite-dimensional ODE structure.
Control Integration: The framework includes packaged model-based controllers (PID, Computed Torque, Impedance Control) and supports coordinate transformations to operational and actuation spaces.

Implementation Details

Singularity Handling: The code explicitly handles removable singularities (e.g., $0/0$ in Lie-group exponentials) using conditional logic (lax.cond) to ensure gradients remain well-defined near straight configurations or zero arclength.
Structure-Aware Evaluation: The system distinguishes between specialized models (e.g., planar PCS) and general formulations, using lower-dimensional algebra and closed-form maps where possible to reduce computational overhead.
Extensibility: Users can inherit from base classes to define custom robots by overriding specific parameters (geometry, material properties, actuation routing) without rewriting the entire dynamics engine.

3. Key Contributions

First GPU-Parallel, Differentiable Rod-Based Framework: SoRoMoX is the first framework implementing PCS and GVS models that runs natively on GPUs/TPUs and is fully differentiable with respect to all model components.
Control-Ready API: It provides a unified interface for all quantities required by model-based control algorithms, eliminating the need for separate simulation and control codebases.
Performance Benchmarks:
- Sequential CPU: Up to 18.1× faster than state-of-the-art alternatives (SoRoSim) for sequential rollouts.
- Parallel GPU: Up to 234.6× higher throughput via batched GPU execution compared to CPU baselines.
Comprehensive Case Studies: The paper demonstrates the framework's utility across six distinct application domains:
- Static-equilibrium system identification.
- Residual-force learning.
- Model-based control (Computed Torque, Impedance).
- Control-gain optimization.
- Safety-constrained control using High-Order Control Barrier Functions (HOCBFs).
- Parallel Reinforcement Learning (RL).

4. Experimental Results

The paper validates SoRoMoX through quantitative benchmarks and application case studies:

System Identification: Using gradient-based optimization on static equilibrium data, the framework identified physical parameters (Young's modulus, Poisson's ratio, density) for a tendon-driven arm, reducing marker Root Mean-Squared Error (RMSE) by 66% (from 56.9 mm to 19.3 mm).
Residual Learning: By learning residual generalized forces to correct model mismatch, the RMSE was further reduced by 64% relative to the identified model, achieving a final error of ~6.9 mm.
Model-Based Control:
- Computed Torque (CT): CT tracking reduced RMSE by a factor of approximately 500 compared to model-free PD control.
- Gain Optimization: Automated gain tuning reduced control loss by up to 62% compared to untuned gains.
Safety-Constrained Control: In a task requiring contact with obstacles, a safety-aware controller (HOCLF+HOCBF) successfully kept peak contact forces within a 5 N bound, whereas a safety-unaware controller exceeded this limit by reaching 33.5 N.
Reinforcement Learning: Training a policy for a soft arm using parallel rollouts on a single GPU was 7× faster than a CPU-based PyElastica baseline (using 512 parallel environments).

5. Significance and Claims

The authors claim that SoRoMoX represents a paradigm shift in soft robotics software, transferring the "differentiable, accelerator-based" workflow standard in rigid robotics to the soft robotics domain.

Bridging the Gap: The paper argues that the primary bottleneck is no longer the derivation of models, but their deployment. SoRoMoX fills this gap by making reduced-order models directly usable in large-scale optimization and learning pipelines.
Enabling New Workflows: The performance gains enable workflows previously considered impractical, such as real-time safety filtering with high-order barrier functions and training RL policies in minutes rather than days.
Reproducibility and Standardization: By providing an open-source, well-documented, and control-ready interface, the authors aim to facilitate benchmarking, reproducibility, and the sharing of robot designs across the community. They suggest that the field would benefit from a standardized robot-description format (analogous to URDF) for continuum soft robots, a direction SoRoMoX supports through its modular design.

The paper concludes that modern software design choices—specifically automatic differentiation, batched execution, and hardware acceleration—can substantially broaden the practical application of soft robot models, moving them from theoretical constructs to tools for real-world control and learning.