google-deepmind / google-deepmind/mujoco_playground

training result is different in dm_control_suite and train_jax_ppo

Open
#177 5 comments 0 reactions 1 assignee Claimed by @btaba View on GitHub
wontfix
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:

Image

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

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.