RobotLocomotion / drake

Model-based design and verification for robotics.
https://drake.mit.edu
Other
3.18k stars 1.24k forks source link

Add tutorial on specifying custom gradients #21502

Closed sea-bass closed 4 weeks ago

sea-bass commented 1 month ago

@EricCousineau-TRI recently helped me figure out how to specify your own custom gradients, both analytically and with finite differencing, so I owe you guys a tutorial.

See here for the nbviewer link: https://nbviewer.org/github/sea-bass/drake/blob/custom-gradients-tutorial/tutorials/custom_gradients.ipynb


This change is Reviewable

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 113 at r1 (raw file):

Previously, RussTedrake (Russ Tedrake) wrote…
note: I think we should recommend calling `InitializeAutoDiff(value, gradient)` here; in general constraints are vectors and gradients are (Jacobian) matrices, which could be difficult to get right. For that reason, I would also recommend making the example constraint have num_constraints > 1.

Re: num_constraints > 1, I was debating something like this.

Basically a tradeoff between pedagogical clarity (1 constraint) and actually motivating the need for an array of things.

I can probably do 2 "stay out of this circle" constraints which strikes a good balance.

sea-bass commented 1 month ago

Thanks for the review, @RussTedrake! I implemented your feedback and also added plots and equations to better illustrate the problem. I'm much happier with it now.

jwnimmer-tri commented 1 month ago

Would any volunteers like to sign up as feature reviewer @EricCousineau-TRI @hongkai-dai @RussTedrake?

sea-bass commented 1 month ago

Great, thanks! Yes, I'm happy to review but I'll wait until after the output cells are cleared before I review. It's too cluttered for me to review in the current shape.

Fantastic timing -- just did this, plus addressed the remaining minor comments from @EricCousineau-TRI that he didn't do himself.

jwnimmer-tri commented 1 month ago

@drake-jenkins-bot ok to test

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 99 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
When I run this locally, I get the following answer: ``` Success? True SolutionResult.kSolutionFound x* = [-9.5e-11 5.0e-01] cost = 0.7499999997575 ``` This doesn't match the `0.48192, 0.13325` answer from the prior figure, but it seems to me like it should have? Ditto for the custom gradients cell, below. They all have the `(0, 0.5)` answer.

I'm not able to reproduce this -- it seems to run fine consistently on my end with drake==1.29.0. What environment are you using?

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 10 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
BTW Somewhere either in the title cell here, or the next cell below, it might help the audience to outline what they are going to learn by going through this document. Currently it kinda dives right into the lower level details without pointing out where we're headed.

Good idea -- added a small preamble.

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 99 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
I reverted #21499 and #21492 and #21486 and tried again; same result. So you can probably ignore my ping Hongkai; sorry. I think it's just luck for which result I see vs you're seeing locally. If I change the initial guess to `[0.5, 0.5]`, then I get the same answer as you. WDYT about moving the guess closer to the solution, to make the NLP more likely to find the answer we want?

I'll move the ICs to [0.5, 0.5] -- it makes sense to start in the center of the solution space anyhow.

Thanks for digging into this!

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 150 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
nit > NOTE: x_value is a column vector so we specify two indices to avoid warnings Wouldn't it be more clear to reshape it after the `ExtractValue`, so that we can use 1d indexing pervasively for it? Especially in the constraint, it would probably help a lot.

Nice suggestion! One can indeed throw a .flatten() at this problem to squash the NumPy warnings.

sea-bass commented 1 month ago

tutorials/custom_gradients.ipynb line 171 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
BTW Here (and possibly also in the cost), I could imagine calling these `y_value` and `y_gradient` for symmetry with `x_value` and `x_gradient`. That would make it crystal clear that x is the input, y is the output, and each has a value part and a gradient part of similar fates -- they can be unbundled and rebundled using Extract and Initialize.

I'm not sure I agree here -- I think since the decision variables are named x = [x0, x1], expressing things in terms of x is clearer than introducing another letter.

Plus for the chain rule expressions, it reads as grad(f(x)) = g(x) * grad(x), where g(x) is your analytically provided gradient.

sea-bass commented 1 month ago

.gitignore line 34 at r7 (raw file):

Previously, jwnimmer-tri (Jeremy Nimmer) wrote…
nit Please revert the changes to this file. The central purpose of this PR is to add a new tutorial. Customizing local developers' dotfiles is a fine proposal, but is independent and distinct from the main end-user purpose here, and might require a different set of reviewers / approvals.

Sounds good to me!

jwnimmer-tri commented 2 weeks ago

Thanks for this!

FYI It's reached Drake's public website now, as part of the v1.30.0 release: https://deepnote.com/workspace/Drake-0b3b2c53-a7ad-441b-80f8-bf8350752305/project/Tutorials-2b4fc509-aef2-417d-a40d-6071dfed9199/notebook/custom_gradients-7fc33ebf63674e9faa9f8aa0b98d1091