google-deepmind / google-deepmind/mujoco_playground
training result is different in dm_control_suite and train_jax_ppo
- Dominant language
- Python
- Stars
- 2.2k
- Forks
- 359
- Avg merge
- 7d 3h
- Merged PRs (30d)
- 1
Description
I trained `python train_jax_ppo.py --env_name=PendulumSwingup` and the result is:
> Experiment name: PendulumSwingup-20250814-144053
Logs are being stored in: /home/balance/fork_mujoco_playground/mujoco_playground/learning/logs/PendulumSwingup-20250814-144053
No checkpoint path provided, not restoring from checkpoint
Checkpoint path: /home/balance/fork_mujoco_playground/mujoco_playground/learning/logs/PendulumSwingup-20250814-144053/checkpoints
0: reward=-675.831
34078720: reward=368.568
68157440: reward=383.898
102236160: reward=359.365
136314880: reward=384.249
170393600: reward=347.015
204472320: reward=364.253
238551040: reward=347.794
272629760: reward=375.346
306708480: reward=371.025
Done training.
Time to JIT compile: 14.696627756999078
Time to train: 674.8724139680016
Starting inference...
FPS for rendering: 50.0
100%|███████████████████████████████████████████████████████████████████████████████████████████| 1001/1001 [00:01<00:00, 722.05it/s]
Rollout video saved as 'rollout.mp4'.
Also, I trained in dm_control_suite.ipynb and used registry.load('PendulumSwingup') to make sure they have the same config:
> Environment Config:
action_repeat: 1
ctrl_dt: 0.01
episode_length: 2000
impl: jax
nconmax: 0
njmax: 0
sim_dt: 0.005
vision: false
PPO Training Parameters:
action_repeat: 1
batch_size: 1024
discounting: 0.998
entropy_cost: 0.008
episode_length: 2000
learning_rate: 0.0002
normalize_observations: true
num_envs: 2048
num_evals: 10
num_minibatches: 32
num_timesteps: 300000000
num_updates_per_batch: 8
reward_scaling: 0.5
unroll_length: 40
But in the dm_control_suite, the result is much better:
I used the same conda environment and only changed CartpoleBalance to PendulumSwingup.
I guess maybe eval_env is different, but I commented out `eval_env=None if _VISION.value else eval_env` and the results didn't change:
> make_inference_fn, params, _ = train_fn( # pylint: disable=no-value-for-parameter
environment=env,
progress_fn=progress,
policy_params_fn=policy_params_fn,
eval_env=None if _VISION.value else eval_env,
)
In addition, I have tried `export JAX_DEFAULT_MATMUL_PRECISION=highest` but it didn't work.
I wonder why the results are so different and hope for some advice. Thanks!
Contributor guide
Assessment
This issue has not been assessed yet.