Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Speed up computation of contact jacobians #188

Merged
merged 2 commits into from
Jun 24, 2024

Conversation

diegoferigo
Copy link
Member

@diegoferigo diegoferigo commented Jun 21, 2024

The computation of the Jacobians corresponding to the implicit frames associated to the collidable points was introduced in #163. At that time, there was a difference in the body frame considered in the Jacobians computed with jaxsim.api.model.generalized_free_floating_jacobian, that was fixed in #167.

For this reason, originally the code to compute the contact Jacobian was calling jaxsim.link.jacobian inside a jax.vmap transformation. This PR modifies this logic to compute the Jacobians of all links only once, and then with jax.vmap extract the one corresponding to the link to which the collidable point is attached.

On the ErgoCub model, that has 32 collidable points, the speedup is significant:

# Before
8.46 ms ± 574 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# After
262 µs ± 6.87 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

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

@diegoferigo diegoferigo self-assigned this Jun 21, 2024
@diegoferigo diegoferigo marked this pull request as ready for review June 21, 2024 14:04
Copy link
Collaborator

@flferretti flferretti left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great, thanks Diego! LGTM

@diegoferigo diegoferigo merged commit 69d3f40 into main Jun 24, 2024
30 checks passed
@diegoferigo diegoferigo deleted the optimize_contact_jacobians branch June 24, 2024 09:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants