Optimize computations via caching

[Giulero]

This PR aims to optimize crba, rnea and aba computations by caching and precomputing quantities.

Some benchmarking:

Now

URDF: stickbot.urdf | NDoF: 34 | batch: 4096 | dtype: torch.float32 | device: cuda | rep: mixed
RNEA: 112.229 ms/iter
ABA : 149.751 ms/iter
CRBA: 94.922 ms/iter

Previous

URDF: stickbot.urdf | NDoF: 34 | batch: 4096 | dtype: torch.float32 | device: cuda | rep: mixed
RNEA: 145.100 ms/iter
ABA: 190.172 ms/iter
CRBA: 143.114 ms/iter

---

📚 Documentation preview 📚: https://adam-docs--143.org.readthedocs.build/en/143/

[Copilot]

Pull Request Overview

This PR optimizes the CRBA, RNEA, and ABA computations by introducing caching mechanisms and precomputing static tree data. The key optimization is the addition of a _prepare_tree_cache() method that pre-computes tree structure information at initialization time, avoiding repeated Python lookups during algorithm execution.

• Introduces _prepare_tree_cache() to precompute tree indices, parent relationships, joint indices, motion subspaces, and spatial inertias
• Refactors CRBA, RNEA, and ABA algorithms to use cached data structures instead of repeatedly traversing the model tree
• Changes default dtype from torch.float64 to torch.float32 in PyTorch implementations

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/adam/core/rbd_algorithms.py	Adds caching infrastructure and refactors CRBA/RNEA/ABA algorithms to use pre-computed tree data, adds Joint import
src/adam/pytorch/computations.py	Changes default dtype parameter from torch.float64 to torch.float32
src/adam/pytorch/computation_batch.py	Changes default dtype parameter from torch.float64 to torch.float32
tests/test_pytorch.py	Explicitly sets dtype=torch.float64 in test setup
tests/test_pytorch_batch.py	Explicitly sets dtype=torch.float64 in test setup

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

The default dtype has been changed from torch.float64 to torch.float32, but the tests explicitly set dtype=torch.float64. This inconsistency suggests that either the tests need to be updated to test with the new default (float32), or the default should remain float64. Changing precision defaults can affect numerical accuracy in robotics computations, especially for inverse dynamics and mass matrix calculations. If the change is intentional for performance reasons, it should be documented why float32 is sufficient for these computations.

Suggested change

	dtypes: torch.dtype = torch.float32,
	dtypes: torch.dtype = torch.float64,

[Giulero]

I'd like to have increased precision for the tests.

[Copilot]

The default dtype has been changed from torch.float64 to torch.float32, but the tests explicitly set dtype=torch.float64. This inconsistency suggests that either the tests need to be updated to test with the new default (float32), or the default should remain float64. Changing precision defaults can affect numerical accuracy in robotics computations, especially for inverse dynamics and mass matrix calculations. If the change is intentional for performance reasons, it should be documented why float32 is sufficient for these computations.

Suggested change

	dtype: torch.dtype = torch.float32,
	dtype: torch.dtype = torch.float64,

[Giulero]

I'd like to have increased precision for the tests.

[Copilot]

Import of 'Node' is not used.

Suggested change

	from adam.model import Model, Node, Joint
	from adam.model import Model, Joint

[Giulero]

Thanks @traversaro!

CI's happy. Merging!