Speed up humanoid's training in MJX notebook

[goncalog]

Summary

This change can result in:

• 1.4x speed up in the humanoid's training in the MJX notebook
• 14% better average reward per episode at the end of the training (improvement is only 7% if ignoring the worst run with each batch_size)

Example runs

Original - batch_size=1024

New - batch_size=512

Experiment results

Ran the tutorial notebook 5 times with each batch_size on Colab with a T4 GPU runtime, and this reduced the training time from 18min to 13min with better reward per episode:

[google-cla[bot]]

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.

[btaba]

Thanks for the contribution! Can you please sign the CLA? https://github.com/google-deepmind/mujoco/pull/1465/checks?check_run_id=22203814795