Exploit branch-induced sparsity to invert the mass matrix

[diegoferigo]

Inverting the mass matrix $M \in \mathbb{R}^{(6+n)\times(6+n)}$ is a operation commonly used in model-based control. Also for simulation purpose, it can be used to compute the forward dynamics without relying on ABA, like we do in:

jaxsim/src/jaxsim/api/model.py

Line 667 in 4fd2032

ν̇ = jnp.linalg.solve(M, S @ τ - h + JTf)

If we decompose the free-floating as follows:

$$M = \begin{pmatrix} M_{bb} & M_{bs} \\\ M_{bs}^T & M_{ss} \end{pmatrix}$$

we can exploit the known topology of the kinematic tree defined by the parent array $\lambda(i)$ to speed up the inversion of $M_{ss} \in \mathbb{R}^{n \times n}$.

Enhancing the performance of this inversion could enable downstream users to implement alternative forward dynamics beyond our ABA and CRB implementations, for example including second-order dynamics like advanced motor dynamics (e.g. #62) or musculoskeletal models. If performance are not too far from ABA, it could be a great alternative of include these effects in ABA since it might be a daunting task.

Some references:

• Slide 10 of CRBA.pdf from https://royfeatherstone.org/teaching/CompuRobDyn2022.zip
• https://royfeatherstone.org/spatial/v2/index.html#LTL
• https://royfeatherstone.org/spatial/v2/sourceText/LTL.txt

Note that the code from Featherstone only inverts $M_{ss}$. We can exploit the following property to extend the result to the free-floating mass matrix:

https://www.wikiwand.com/en/Block_matrix#Block_matrix_inversion

[flferretti]

More resources:

• jax.experimental.sparse module
• Semi-structured sparsity jax-ml/jax#19212
• jax.scipy.sparse.linalg.cg

[diegoferigo]

A while ago I started playing a bit around with this type of inversion, but I couldn't get any better performance than calling jax.numpy.linalg.inv. I came up to the conclusion that the best approach to get the inverse of the mass matrix is to implement a new RBDA that directly provides $M(\mathbf{q})^{-1}$:

• Carpentier J., Analytical Inverse of the Joint Space Inertia Matrix, URL.
• Best way to compute forward dynamics and M inverse stack-of-tasks/pinocchio#1215

[flferretti]

WIP at https://github.com/ami-iit/jaxsim/tree/feature/mass_inverse_rbd