mujoco
mujoco copied to clipboard
Speed up simulation and training in the MJX's tutorial
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
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
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:
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.
Hello, wondering if there was any feedback on this? Thank you!
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?
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?
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?
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:
New:
LGTM