mujoco icon indicating copy to clipboard operation
mujoco copied to clipboard

Speed up simulation and training in the MJX's tutorial

Open goncalog opened this issue 1 year ago • 1 comments

Summary

This change can result in a 1.7x speed up in the humanoid's training time

Example runs

Original - num_envs=2048; num_minibatches=32; batch_size=1024; mj_model.opt.iterations=6; mj_model.opt.ls_iterations=6; physics_steps_per_control_step=5 Screenshot 2024-08-30 at 15 55 02

New - num_envs=3072; num_minibatches=24; batch_size=512; mj_model.opt.iterations=4; mj_model.opt.ls_iterations=4; physics_steps_per_control_step=4 Screenshot 2024-08-30 at 15 00 04

Experiment results

Ran the tutorial notebook 4 times with each set of params on Kaggle with a T4 x2 GPU runtime, and this reduced the training time from 9.5min to 5.5min with a similar reward per episode: Screenshot 2024-08-30 at 16 16 36

goncalog avatar Aug 30 '24 15:08 goncalog

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

google-cla[bot] avatar Aug 30 '24 15:08 google-cla[bot]

Hello, wondering if there was any feedback on this? Thank you!

goncalog avatar Sep 25 '24 10:09 goncalog

I'm finding that the new parameters, in particular the solver parameters, deteriorate the "MJX Policy in MuJoCo" example, since the simulation is less stable. Can you find better hypers for PPO without modifying the solver parameters?

btaba avatar Oct 04 '24 20:10 btaba

hi @btaba, thank you for looking into this! of course, just reverted the solver param changes. please let me know if that's what you had in mind?

goncalog avatar Oct 05 '24 12:10 goncalog

hi @btaba, thank you for looking into this! of course, just reverted the solver param changes. please let me know if that's what you had in mind?

Thanks @goncalog , do you still get the same speed-up?

btaba avatar Oct 05 '24 23:10 btaba

hi @btaba, thank you for looking into this! of course, just reverted the solver param changes. please let me know if that's what you had in mind?

Thanks @goncalog , do you still get the same speed-up?

hey @btaba, no. now i only get a 1.4x speedup.

Original: Screenshot 2024-10-07 at 15 52 12

New: Screenshot 2024-10-07 at 15 31 52

goncalog avatar Oct 07 '24 14:10 goncalog

LGTM

btaba avatar Oct 07 '24 20:10 btaba