Add recursive algorithm to compute the inverse of the mass matrix

[flferretti]

This PR adds a RBDA to compute the inverse of the floating base mass matrix analytically.

Note

Reference pseudocode: https://laas.hal.science/hal-01790934v2/document

Solves #107

Benchmarks:

Main method compilation time: 0.4445 s
RBDA method compilation time: 0.4309 s
Main method FPS: 1176.09
RBDA method FPS: 1952.34

Test Script

import pathlib

import jax
import jax.numpy as jnp

import jaxsim.api as js

model_urdf_path = pathlib.Path(
    "../element_gymnaxsim/gymnaxsim/src/gymnaxsim/assets/models/ergocub.urdf"
)

model = js.model.JaxSimModel.build_from_model_description(model_urdf_path)
data = js.data.JaxSimModelData.build(model=model)


@jax.jit
def compute_mass_inverse_main(model, data):
    return jnp.linalg.inv(js.model.free_floating_mass_matrix(model=model, data=data))


@jax.jit
def compute_mass_inverse_rbda(model, data):
    return js.model.free_floating_mass_matrix_inverse(model=model, data=data)


if __name__ == "__main__":
    import time
    import timeit

    # Compilation
    now = time.perf_counter()
    M_inv_main = compute_mass_inverse_main(model, data)
    print(f"Main method compilation time: {time.perf_counter() - now:.4f} s")

    now = time.perf_counter()
    M_inv_rbda = compute_mass_inverse_rbda(model, data)
    print(f"RBDA method compilation time: {time.perf_counter() - now:.4f} s")

    # Execution
    n_iters = 1000
    main_time = timeit.timeit(
        stmt="compute_mass_inverse_main(model, data).block_until_ready()",
        setup="from __main__ import compute_mass_inverse_main, model, data",
        number=n_iters,
    )
    rbda_time = timeit.timeit(
        stmt="compute_mass_inverse_rbda(model, data).block_until_ready()",
        setup="from __main__ import compute_mass_inverse_rbda, model, data",
        number=n_iters,
    )

    main_fps = n_iters / main_time
    rbda_fps = n_iters / rbda_time

    print(f"Main method FPS: {main_fps:.2f}")
    print(f"RBDA method FPS: {rbda_fps:.2f}")

    # Verify that the results are close
    assert jnp.allclose(M_inv_main, M_inv_rbda)

---

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

[github-actions]

Benchmark

Details

Benchmark suite	Current: ed3214e	Previous: c28e89a	Ratio
tests/test_benchmark.py::test_forward_dynamics_aba[1]	332.1375691488637 iter/sec (stddev: 0.000009114824815381149)	385.94492720843766 iter/sec (stddev: 0.000009463978323506126)	1.16
tests/test_benchmark.py::test_forward_dynamics_aba[128]	26.071534816387622 iter/sec (stddev: 0.00012060978877376893)	29.216966653272625 iter/sec (stddev: 0.00011455086573428671)	1.12
tests/test_benchmark.py::test_free_floating_bias_forces[1]	306.90678503574213 iter/sec (stddev: 0.00007429628500080161)	361.396251061539 iter/sec (stddev: 0.00000910574964681653)	1.18
tests/test_benchmark.py::test_free_floating_bias_forces[128]	14.613588504133034 iter/sec (stddev: 0.00018935031210388077)	16.218942194940492 iter/sec (stddev: 0.00025096534945186314)	1.11
tests/test_benchmark.py::test_forward_kinematics[1]	378.00488549455054 iter/sec (stddev: 0.000022540270976848718)	452.4863643105723 iter/sec (stddev: 0.000008022285629396292)	1.20
tests/test_benchmark.py::test_forward_kinematics[128]	28.092721177604126 iter/sec (stddev: 0.00014483414454265675)	31.765642438814933 iter/sec (stddev: 0.0002088457881717258)	1.13
tests/test_benchmark.py::test_free_floating_mass_matrix[1]	153.96005923658066 iter/sec (stddev: 0.000012540423101131511)	181.58130216807277 iter/sec (stddev: 0.00005072371260792714)	1.18
tests/test_benchmark.py::test_free_floating_mass_matrix[128]	152.23098438720007 iter/sec (stddev: 0.000015525758360546094)	175.24079377087364 iter/sec (stddev: 0.000018760096947039575)	1.15
tests/test_benchmark.py::test_free_floating_jacobian[1]	433.38360053633124 iter/sec (stddev: 0.000007436200619924758)	525.7508803201982 iter/sec (stddev: 0.000006377672924636457)	1.21
tests/test_benchmark.py::test_free_floating_jacobian[128]	438.71148687199735 iter/sec (stddev: 0.000009253173342369444)	534.2483568468638 iter/sec (stddev: 0.000006942445781767468)	1.22
tests/test_benchmark.py::test_free_floating_jacobian_derivative[1]	353.06409773751716 iter/sec (stddev: 0.000010054695132100777)	410.0473494948492 iter/sec (stddev: 0.000012243950755069747)	1.16
tests/test_benchmark.py::test_free_floating_jacobian_derivative[128]	302.0616069887201 iter/sec (stddev: 0.00002805283036423913)	302.58988317111164 iter/sec (stddev: 0.000015109621352982298)	1.00
tests/test_benchmark.py::test_soft_contact_model[1]	301.67479895263335 iter/sec (stddev: 0.000013448270613339103)	349.9366267743607 iter/sec (stddev: 0.00000947480013200776)	1.16
tests/test_benchmark.py::test_soft_contact_model[128]	26.47180196248617 iter/sec (stddev: 0.00014897477396433738)	29.951416599383442 iter/sec (stddev: 0.0001405258247109449)	1.13
tests/test_benchmark.py::test_rigid_contact_model[1]	47.48909828448196 iter/sec (stddev: 0.000013378418542120032)	40.276241063183 iter/sec (stddev: 0.000024302278784552764)	0.85
tests/test_benchmark.py::test_rigid_contact_model[128]	0.7094938449588688 iter/sec (stddev: 0.00038758875379999954)	0.7258859819172665 iter/sec (stddev: 0.0003607214923404956)	1.02
tests/test_benchmark.py::test_relaxed_rigid_contact_model[1]	104.73092822478951 iter/sec (stddev: 0.000018817220702752182)	74.75017951811331 iter/sec (stddev: 0.00003167479964364776)	0.71
tests/test_benchmark.py::test_relaxed_rigid_contact_model[128]	5.072356605038323 iter/sec (stddev: 0.0002489849756172521)	6.400062476314604 iter/sec (stddev: 0.00024891629896200675)	1.26
tests/test_benchmark.py::test_simulation_step[1]	92.41770985603854 iter/sec (stddev: 0.00002955414463102609)	69.29459728383536 iter/sec (stddev: 0.00002754549281911007)	0.75
tests/test_benchmark.py::test_simulation_step[128]	4.474772208032762 iter/sec (stddev: 0.00021237547320095867)	5.39599749835251 iter/sec (stddev: 0.000279412049524645)	1.21
tests/test_benchmark.py::test_update_hw_parameters[1]	115.8217313659419 iter/sec (stddev: 0.00004646220517677213)
tests/test_benchmark.py::test_update_hw_parameters[128]	117.28829857718046 iter/sec (stddev: 0.00002883992689621621)
tests/test_benchmark.py::test_export_updated_model[1]	3.494226348133711 iter/sec (stddev: 0.0020932500210026845)
tests/test_benchmark.py::test_export_updated_model[128]	0.01896556570452362 iter/sec (stddev: 0.4055584729954908)

This comment was automatically generated by workflow using github-action-benchmark.

[flferretti]

In edf62ad I added a helper function to avoid the use of jax.scipy.linalg.block_diag, which allocates a (6 + NDoFs) matrix and directly multiply the blocks. This saves some memory and provides an additional ~10% speedup:

Main method compilation time: 0.4495 s
RBDA method compilation time: 0.4388 s
Main method FPS (with the helper): 1195.85
RBDA method FPS (with the helper): 2095.91

[xela-95]

Thanks @flferretti that's great! LGTM