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

First steps towards automatic differentiation of RBDAs #54

Merged
merged 3 commits into from
Oct 11, 2023

Conversation

diegoferigo
Copy link
Member

@diegoferigo diegoferigo commented Oct 10, 2023

This PR works around some problems I encountered when trying to make automatic differentiation (AD) run on our rigid-body dynamics algorithms (RBDAs).

This is not the first time I've tried this in the past, but every time there were NaNs generated somewhere in the code that were quite difficult to spot. I never had enough time to dedicate to dissect small sections and proceed with a pragmatic debugging campaign.

Luckily, JAX development progressed considerably, and now thanks to the following features:

  • better errors in case of problems,
  • JAX_DEBUG_NANS=1 that points directly to the origin of any NaN ,
  • JAX_DISABLE_JIT=1 that disabled JIT entirely also inside e.g. jax.lax.scan, allowing to spot NaNs occurring there,

I finally managed to make the tests against finite difference (with jax.test_util.check_grad) work, both for the first-order and second-order derivatives! 🎉 I guess we are one step closer to a fully differentiable simulator, and definitely towards #4.

For the moment, this PR checks that JAX is able to differentiate through the following algorithms:

  • ABA
  • RNEA
  • CRBA
  • Forward kinematics
  • Jacobians
  • Soft contacts

I still have to check properly our integrators in a future PR. I suspect that the integration of quaternion with Baumgarte stabilization might be suboptimal in these cases, but it can be simply solved by implementing the integration on SO(3). I'm considering giving it a try.

@diegoferigo diegoferigo self-assigned this Oct 10, 2023
@diegoferigo diegoferigo changed the base branch from main to new_api October 10, 2023 08:05
@DanielePucci
Copy link
Member

CC @Giulero

@diegoferigo diegoferigo marked this pull request as ready for review October 10, 2023 15: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.

That's great! 🎉

@diegoferigo diegoferigo merged commit e9bb166 into new_api Oct 11, 2023
16 checks passed
@diegoferigo diegoferigo deleted the automatic_differentiation branch October 11, 2023 06:58
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.

3 participants