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:

https://github.com/ami-iit/jaxsim/blob/4fd2032c74de495cd59e1f6999e1f7c3c9ec2ff8/src/jaxsim/api/model.py#L667

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
• https://github.com/google/jax/issues/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.
• https://github.com/stack-of-tasks/pinocchio/issues/1215