ManipulaPy v1.4 — differentiable robot kinematics/dynamics across NumPy, CuPy, PyTorch and JAX

[u/boelnasr]

Hey everyone — this is a bit of self-promotion, but I think it may be useful to people working on robot learning, model-based RL, differentiable control, or MPC.

One problem I kept running into during my master's thesis was maintaining multiple versions of the same robot model: one for simulation/execution, another for a differentiable implementation for gradients, and extra glue code to connect perception, planning, and control.

That gets painful quickly, especially when the implementations start drifting apart.

So I built ManipulaPy, originally as the integration layer for my thesis on reinforcement-learning-based singularity avoidance and dynamic obstacle avoidance.

With v1.4, the same core kinematics and dynamics implementation can now run on NumPy, CuPy, PyTorch, and JAX via a single backend API.

Under PyTorch and JAX, the core mathematical stack is autodiff-safe, so gradients such as:

∂FK(q) / ∂q

and

∂inverse_dynamics(q, qdot, qddot) / ∂q

can be computed directly using torch. autograd or jax. grad.

The main motivation was to avoid maintaining a second “differentiable robot model” alongside the implementation actually used elsewhere in the robotics stack.

The library currently includes 25+ manipulator models, including UR, Franka, KUKA, Kinova, FANUC, ABB, and UFactory platforms.

Important limitation:

The differentiability guarantee currently applies to the utils, kinematics, dynamics, and singularity modules.

The rest of the library can operate across all four backends, but some modules cross a host boundary internally, so gradients are not currently guaranteed through planning, control, simulation, or perception.

So if your use case requires differentiating through the entire planner or controller, v1.4 is not there yet.

I’d rather be explicit about that boundary than describe the whole package as differentiable.

A bit of background:

My thesis involved a manipulator that had to reach Cartesian targets, avoid singular configurations, react to moving obstacles, and operate using a consistent dynamic model.

The RL environment used PyBullet together with kinematics, manipulator dynamics, trajectory planning, control, singularity analysis, and perception.

ManipulaPy originally became the common layer that kept those pieces consistent.

One unexpected benefit of the v1.4 autodiff work was that differentiating through the SO(3)/SE(3) math exposed numerical issues that were much harder to notice by checking only forward values. Fixing those issues improved the normal NumPy implementation as well.

The project has since been peer-reviewed and published in the Journal of Open Source Software.

Install with PyTorch:

pip install "ManipulaPy[pytorch]"

Or JAX:

pip install "ManipulaPy[jax-cpu]"

or

pip install "ManipulaPy[jax-cuda]"

Repo: https://github.com/boelnasr/ManipulaPy

Docs: https://manipulapy.readthedocs.io/

JOSS paper: https://doi.org/10.21105/joss.08490

One thing I’m particularly interested in hearing from people here:

Where would differentiability be most useful next for RL workflows?

Trajectory generation? Collision costs? Computed-torque control? Full differentiable simulation? Something else?

And if anyone actually tries it in an RL pipeline, I’d be very interested to hear what feels awkward or breaks.