Enhance performance and memory efficiency of hardware parametrization

[flferretti]

This pull request introduces significant changes to simplify transform updates, add support for joint-free models, and reduce memory usage in various computations. The most important updates include refactoring the HwLinkMetadata class, optimizing the scaling logic, and enhancing the computation of mass and inertia. Additionally, some test cases and imports have been updated to align with the new implementation.

Benchmarks

Computational times obtained with time pytest tests/test_api_model_hw_parametrization.py on Intel Core Ultra 7 155H:

< pytest tests/test_api_model_hw_parametrization.py  40.63s user 1.44s system 148% cpu 28.282 total
> pytest tests/test_api_model_hw_parametrization.py  36.00s user 1.23s system 141% cpu 26.242 total

Refactoring and Enhancements to HwLinkMetadata:

• Changed the shape attribute in HwLinkMetadata to _shape (private) and added a new shape property for accessing it as an np.array. This improves encapsulation and clarity. (src/jaxsim/api/kin_dyn_parameters.py)
• Refactored shape-specific methods (_box, _cylinder, _sphere) to be local functions within compute_mass_and_inertia, simplifying the code structure. (src/jaxsim/api/kin_dyn_parameters.py) [1] [2] [3]

Improved Scaling Logic:

• Introduced a has_joints boolean to handle models with or without joints more effectively in scaling operations. Updated apply_scaling to conditionally apply scaling based on this flag. (src/jaxsim/api/kin_dyn_parameters.py)
• Replaced jax.lax.switch with a more flexible approach for computing scaling vectors and mass/inertia for supported shapes, improving maintainability. (src/jaxsim/api/kin_dyn_parameters.py) [1] [2]

Mass and Inertia Computation:

• Updated compute_mass_and_inertia to use vmap for batch processing of shape types, improving performance and scalability. (src/jaxsim/api/kin_dyn_parameters.py)
• Adjusted test cases to reflect the new compute_mass_and_inertia method signature and ensure correctness. (tests/test_api_model_hw_parametrization.py)

Miscellaneous Updates:

• Removed redundant imports and cleaned up unused code segments for better readability. (src/jaxsim/api/model.py) [1] [2]
• Simplified the logic for handling unsupported links by removing unnecessary conditional branches. (src/jaxsim/api/kin_dyn_parameters.py)

These changes collectively enhance the library's flexibility, maintainability, and performance, particularly in handling complex hardware link metadata and scaling operations.

---

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

[github-actions]

Benchmark

Benchmark suite	Current: f8b58c1	Previous: cb36385	Ratio
tests/test_benchmark.py::test_forward_dynamics_aba[1]	386.8809652626908 iter/sec (stddev: 0.000029622668238794555)	386.2552251710958 iter/sec (stddev: 0.00020833313469613508)	1.00
tests/test_benchmark.py::test_forward_dynamics_aba[128]	25.061283237656657 iter/sec (stddev: 0.00038026174418327074)	25.623251572538486 iter/sec (stddev: 0.000628276488211365)	1.02
tests/test_benchmark.py::test_free_floating_bias_forces[1]	343.738296500269 iter/sec (stddev: 0.00003254247889272427)	329.6208935032452 iter/sec (stddev: 0.00017678390636095912)	0.96
tests/test_benchmark.py::test_free_floating_bias_forces[128]	14.373161352225022 iter/sec (stddev: 0.0002957696142840192)	14.042023282252773 iter/sec (stddev: 0.0004701444679618353)	0.98
tests/test_benchmark.py::test_forward_kinematics[1]	471.7702879880674 iter/sec (stddev: 0.000027922230114891526)	505.33330250290584 iter/sec (stddev: 0.00003595675660629741)	1.07
tests/test_benchmark.py::test_forward_kinematics[128]	28.08812364428693 iter/sec (stddev: 0.00028886312246266397)	29.344980191633002 iter/sec (stddev: 0.00036979712165659257)	1.04
tests/test_benchmark.py::test_free_floating_mass_matrix[1]	238.23485567440022 iter/sec (stddev: 0.00003759650343186609)	204.56693174871816 iter/sec (stddev: 0.0008038774133928276)	0.86
tests/test_benchmark.py::test_free_floating_mass_matrix[128]	215.75027757004133 iter/sec (stddev: 0.000023744625770907658)	185.32420162463208 iter/sec (stddev: 0.001038192142422702)	0.86
tests/test_benchmark.py::test_free_floating_jacobian[1]	565.6756483351708 iter/sec (stddev: 0.000029123323823201942)	584.370298139834 iter/sec (stddev: 0.00001860199318220017)	1.03
tests/test_benchmark.py::test_free_floating_jacobian[128]	580.263330173672 iter/sec (stddev: 0.000015738312245575354)	544.3999181675149 iter/sec (stddev: 0.00023044493810550008)	0.94
tests/test_benchmark.py::test_free_floating_jacobian_derivative[1]	397.8366020206721 iter/sec (stddev: 0.000028864282486279092)	410.73560028011394 iter/sec (stddev: 0.00022687143203532596)	1.03
tests/test_benchmark.py::test_free_floating_jacobian_derivative[128]	242.3766063985751 iter/sec (stddev: 0.00005299746378322223)	199.38111540598362 iter/sec (stddev: 0.001457819036270435)	0.82
tests/test_benchmark.py::test_soft_contact_model[1]	324.7514144191113 iter/sec (stddev: 0.000034026680675898034)	337.6366937284971 iter/sec (stddev: 0.00011127692295687577)	1.04
tests/test_benchmark.py::test_soft_contact_model[128]	25.569510972775788 iter/sec (stddev: 0.0003308336185568806)	26.884777773429075 iter/sec (stddev: 0.0006199761164350777)	1.05
tests/test_benchmark.py::test_rigid_contact_model[1]	34.62702761098002 iter/sec (stddev: 0.00007059823503376248)	31.251271243159007 iter/sec (stddev: 0.007996216837436577)	0.90
tests/test_benchmark.py::test_rigid_contact_model[128]	0.3409199110019724 iter/sec (stddev: 0.007758096807032114)	0.3460028427282277 iter/sec (stddev: 0.0005242995729556135)	1.01
tests/test_benchmark.py::test_relaxed_rigid_contact_model[1]	84.5365097757624 iter/sec (stddev: 0.00005558980756165383)	79.40038194340299 iter/sec (stddev: 0.002071043703591531)	0.94
tests/test_benchmark.py::test_relaxed_rigid_contact_model[128]	6.150201033782996 iter/sec (stddev: 0.0004912174475872905)	6.11243560362339 iter/sec (stddev: 0.0013425887323580527)	0.99
tests/test_benchmark.py::test_simulation_step[1]	75.88875100170199 iter/sec (stddev: 0.00006640162710148092)	69.93411634859267 iter/sec (stddev: 0.0027247172648478526)	0.92
tests/test_benchmark.py::test_simulation_step[128]	4.9462226590880505 iter/sec (stddev: 0.0006328661360358932)	5.076712249345317 iter/sec (stddev: 0.00163782888871275)	1.03

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

[younik]

good, thank you!

[xela-95]

Thanks @flferretti ! LGTM!