JAX Soft Robot Modeling

Important

This project is archived and is no longer under active development. SoRoMoX (Soft Robot Models in jaX) is the successor to jsrm. Most features provided by jsrm are also available in SoRoMoX through a slightly modified interface.

SoRoMoX replaces the symbolic derivation pipeline with scalable, fully numerical JAX implementations and adds articulated, Piecewise Constant Strain (PCS), and Geometric Variable Strain (GVS) models behind a common control-oriented interface. It supports planar and spatial robots, additional actuation models, model-based controllers, multiple rendering backends, and workflows for parameter identification, optimization, residual learning, safety-constrained control, and reinforcement learning. Its JAX-native implementation also provides JIT compilation, automatic differentiation, vectorization, batching, and execution on CPUs, GPUs, and TPUs. In particular, SoRoMoX is significantly faster when simulating more than one planar PCS system, because its fully numerical dynamics implementation is no longer restricted by the limitations of symbolic dynamics derivation.

This repository contains symbolic derivations of the kinematics and dynamics of various soft robots using Sympy. The symbolic expressions are then implemented in JAX and can be used for fast, parallelizable, and differentiable simulations. So far, we have focused on planar settings and implemented the following soft robots:

• N-link pendulum
• Planar Piecewise Constant Strain (PCS) continuum soft robot
• Planar Handed Shearing Auxetics (HSA) robot

We are happy to receive contributions for other soft robots and/or other settings (e.g., 3D).

Citation

This simulator is part of the publication An Experimental Study of Model-based Control for Planar Handed Shearing Auxetics Robots presented at the 18th International Symposium on Experimental Robotics. You can find the publication online in the Springer Proceedings on Advanced Robotics (SPAR): https://doi.org/10.1007/978-3-031-63596-0_14

Please use the following citation if you use our software in your (scientific) work:

@inproceedings{stolzle2023experimental,
  title={An experimental study of model-based control for planar handed shearing auxetics robots},
  author={St{\"o}lzle, Maximilian and Rus, Daniela and Della Santina, Cosimo},
  booktitle={International Symposium on Experimental Robotics},
  pages={153--167},
  year={2023},
  organization={Springer}
}

Installation

The plugin can be installed from PyPI:

pip install jsrm

or locally from the source code:

pip install -e .

If you want to run the examples, you will also need to install the following dependencies:

pip install -e ".[examples]"

Usage

Always first source all necessary environment variables when opening a new terminal:

source 01-configure-env-vars.sh

Then, we can symbolically derive the pendulum kinematics and dynamics:

python examples/derive_pendulum.py

Finally, we can simulate the pendulum

python examples/simulate_pendulum.py

See also

You might also be interested in the following repositories:

• The jax-spcs-kinematics repository contains an implementation of the Selective Piecewise Constant Strain (SPCS) kinematics in JAX. We have shown in our paper that this kinematic model is suitable for representing the shape of HSA rods.
• The HSA-PyElastica repository contains a plugin for PyElastica for the simulation of HSA robots.
• The hsa-planar-control repository contains JAX and ROS2 implementations of model-based control algorithms for planar HSA robots.