google-deepmind / google-deepmind/mujoco_playground

Early termination for Deepmind control suite's Humanoid

Open
#242 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
2.2k
Forks
359
Avg merge
7d 3h
Merged PRs (30d)
1

Description

I am doing some Reinforcement learning experiment (reproducing CrossQ) and trying to train in HumanoidRun environment.
I use SAC as my baseline and use the hyperparameter the same as the parameters in https://github.com/google-deepmind/mujoco_playground/blob/main/learning/notebooks/dm_control_suite.ipynb. It can work. However my algorithm's UTD is 1 so that I set the SAC's gradient steps to 1. Then it didn't learn and the reward didn't increase.
So I am confused, I thought that the humanoid environment should be a widely used environment and to my knowledge even the gradient steps is 1 it should also work. Then I compare the environment with https://github.com/Farama-Foundation/Gymnasium/blob/main/gymnasium/envs/mujoco/humanoid_v4.py . I found that there is early termination in gymnasium's humanoid and it can help a lot to increase data efficiency and prevent the neural network to update in too much useless data. Also I found the early termination in Mujoco Playground's locomotion environment in https://github.com/google-deepmind/mujoco_playground/blob/1bd244fa9c51fa7074c7a449934f556d898f7a15/mujoco_playground/_src/locomotion/h1/joystick_gait_tracking.py#L285 . Then I check the original code for deepmind control suite and I found the original code doesn't have early termination also.
And I set the num_env to 1. And both SAC and CrossQ can work and the performance of CrossQ > SAC as the original paper. So we can see that, the effective UTD = gradient steps / num_envs
So I compare these environment:
Mujoco playground running on TPU/GPUs with parallel environment, so the effective UTD is much smaller than the RL algorithm using CPU Deepmind control. So that some algorithm may perform good using original Deepmind control suite even though it has no early termination.
Gymnasium Humanoid has early termination, so many RL framework like Stable Baselines 3 using it with vectored environment. Even though the effective UTD is smaller but we use early termination it can still learn well.
So I think I am in a dilemma that my algorithm CrossQ is write in Brax and trying to integrate with MuJoCo Playground's ecosystem. But the environment in MuJoCo Playground won't perform well. However Brax need a JAX environment because its rollout function will be compiled by JIT.
So I decide to transfer the Gymnasium's Humanoid environment to the MJX env in mujoco playground. I raise this issue to give you the problem I faced and I was wondering if it's on your roadmap. I think the Humanoid environment is more convenient to use in RL algorithm experiment and hope we can make a better environment for RL algorithm training.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.