Tutorial for how to use the jacobian of mjx.step

[Andrew-Luo1]

A notebook showing how to use vanilla policy gradients to accelerate policy learning, thanks to MJX's differentiability. Demonstrations are given for imitation learning and locomotion, for a quadruped. Requires Brax PR 476.

Supporting changes

1. Add an Anymal C model, cleaned for faster simulation on MJX similarly to MJX Barkour
2. Add an image for the notebook

[Andrew-Luo1]

Just as a heads up:

1. Right now I hard-code the xml path, rather than use epath like in the other MJX tutorial - I could not get epath working locally.
2. For the PPO benchmarks, I ran into issues with nan rewards which dissapeared when I turned on the debug nans flag.

[erikfrey]

Exciting! How many changes did you make to anymal? Can you try this:

• Make a PR to mujoco menagerie with anymal_c_mjx.xml and scene_mjx.xml with your anymal modifications
• Somewhere at the top of your colab, do:

!git clone https://github.com/google-deepmind/mujoco_menagerie.git

Then you should be able to remove the anymal assets from this PR. Let me know if that works!

[Andrew-Luo1]

Hi Erik, I've proposed Menagerie PR 50 with the Anymal C updates, and have documented the modifications in the readme. I've updated this notebook as well.

The formatting of this notebook isn't quite the same as the current mujoco notebooks. In the case that this PR goes through, of course feel free to modify it as you see fit: adding the Mujoco banner, adding a colab instance, etc!

And let me know if there's anything I can do to bring the notebook more up to standards.

[erikfrey]

OK, I see your change is incorporated into menagerie. Looking great! One last nit and we will merge.

[erikfrey]

I suspect this is an unneeded change - can you revert? Otherwise, let me know what this is for.

[erikfrey]

Wonderful, thank you!