Optimize Jacobian algorithm

[diegoferigo]

This PR enhances the computation of the Jacobians of all links. Before this PR, we computed in parallel with jax.vmap the free-floating left-trivialized Jacobian of all links ${}^L J_{W,L/B}$. Then, we were adjusting the input and output representations to match the desired ones.

This PR, instead of calling a parallelized version of jaxsim.rbda.jacobian, computes only once a full doubly-left Jacobian ${}^B J_{W,\text{\_}/B}$ (note that there is a $\text{\_}$ instead of $L$), defined as the free-floating Jacobian with all rows filled. The free-floating doubly-left Jacobian of the i-th link ${}^B J_{W,L/B}$ is then computed by filtering the columns of the full Jacobian with the support parent array $\kappa(i)$ of the link, that sets to zero the columns corresponding to all links not part of the path $\pi_B(L)$.

All of this allows to compute the Jacobians of all links by only vmapping the full Jacobian using $\kappa(i)$, therefore the expensive algorithm to compute the full Jacobian is executed only once. The vmapped operation is just a column filtering, therefore its cost is almost zero.

This enhancement should speed up the computation of jaxsim.api.model.forward_dynamics_crb. While at this stage is always better to use jaxsim.api.model.forward_dynamics_aba, the CRB version might become useful in all experiments that extend the forward dynamics with e.g. actuators, friction models, or any other stateless second-order dynamics elements¹. In these cases, operating on ABA could become quite difficult, and prototyping using the CRB version might be pretty helpful.

cc @traversaro @DanielePucci

---

📚 Documentation preview 📚: https://jaxsim--121.org.readthedocs.build//121/

Footnotes

1. Under the assumption that they do not extend the integrated state vector. In that case, things become more complex since also the integrator has to properly updated. ↩

[traversaro]

Cool! Did you did any benchmark on the difference before and after this change?

[diegoferigo]

I was running them while you were typing your comment :)

Here below a benchmark on a 58-DoFs ErgoCub model on CPU.

Before

In [2]: %timeit -r1 -n1 _ = js.model.generalized_free_floating_jacobian(model, data0)
28 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [3]: %timeit -r 10 -n 1000 _ = js.model.generalized_free_floating_jacobian(model, data0)
1.37 ms ± 82.6 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

In [4]: %timeit -r1 -n1 _ = js.model.forward_dynamics_crb(model, data0)
1min 27s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [5]: %timeit -r10 -n1000 _ = js.model.forward_dynamics_crb(model, data0)
1.99 ms ± 69.1 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

After

In [2]: %timeit -r1 -n1 _ = js.model.generalized_free_floating_jacobian(model, data0)
9.5 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [3]: %timeit -r 10 -n 1000 _ = js.model.generalized_free_floating_jacobian(model, data0)
185 µs ± 16.4 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

In [4]: %timeit -r1 -n1 _ = js.model.forward_dynamics_crb(model, data0)
55.6 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

In [5]: %timeit -r10 -n1000 _ = js.model.forward_dynamics_crb(model, data0)
896 µs ± 51.5 µs per loop (mean ± std. dev. of 10 runs, 1,000 loops each)

Overview

• Compiling generalized_free_floating_jacobian is now 3x faster.
• Compiling forward_dynamics_crb is now 1.5x faster.
• Running generalized_free_floating_jacobian is now 7x faster.
• Running forward_dynamics_crb is now 2x faster.

script

import jaxsim.api as js

import resolve_robotics_uri_py
import rod

# Find the urdf file.
urdf_path = resolve_robotics_uri_py.resolve_robotics_uri(
    uri="model://ergoCubSN001/model.urdf"
)

# Build the ROD model.
rod_sdf = rod.Sdf.load(sdf=urdf_path)

# Build the model.
model = js.model.JaxSimModel.build_from_model_description(
    model_description=rod_sdf.model,
)

# Create random data.
data0 = js.data.random_model_data(
    model=model,
    base_pos_bounds=((0, 0, 0.85), (0, 0, 0.85)),
    joint_vel_bounds=(0, 0),
    base_vel_lin_bounds=(0, 0),
    base_vel_ang_bounds=(0, 0),
)

%timeit -r1 -n1 _ = js.model.generalized_free_floating_jacobian(model, data0)

%timeit -r 10 -n 1000 _ = js.model.generalized_free_floating_jacobian(model, data0)

%timeit -r1 -n1 _ = js.model.forward_dynamics_crb(model, data0)

%timeit -r10 -n1000 _ = js.model.forward_dynamics_crb(model, data0)

[diegoferigo]

In other words, on such a large model, with this change we can compute FD with CRB in less than 1ms with JAX running on CPU. Not too bad. As a comparison, the equivalent computation with ABA runs 3x faster:

In [15]: %timeit _ = js.model.forward_dynamics_aba(model, data0)
346 µs ± 6.74 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[diegoferigo]

cc @ami-iit/vertical_control-oriented-learning