Unable to replicate DM Humanoid results

[mttga]

I’m trying to replicate the results for PPO presented in the paper using the learning/train_jax_ppo.py script. The results for most environments appear correct, but the returns for the humanoid tasks are significantly off (see attached image).

I’m running the scripts using the NVIDIA JAX Container 25.01 on H100 GPUs.

The commands I’m using are:

python learning/train_jax_ppo.py --env_name HumanoidStand --num_timesteps 100000000
python learning/train_jax_ppo.py --env_name HumanoidWalk --num_timesteps 100000000
python learning/train_jax_ppo.py --env_name HumanoidRun --num_timesteps 100000000

Each script was run with 5 different random seeds: 0, 1, 2, 3, and 4.

[btaba]

Ok that's interesting. I ran this a few weeks ago, and the results still repro'd the paper. I didn't use learning/train_jax_ppo.py for the paper FWIW (nor do I use it now). I'll take a closer look. I use this

[mttga]

Ok, could you share the exact command you use to get the paper's results with that script?

[btaba]

Ok I just OSS'd the parts of the script that run playground, and it repros locally (on RTX 5090).

python brax/training/learner.py --playground_dm_control_suite --learner=ppo --env=HumanoidRun --playground_config_overrides='{"impl": "jax"}' --total_env_steps=60000000 --num_evals=10 --reward_scaling=10.0 --episode_length=1000 --normalize_observations --action_repeat=1 --entropy_cost=0.01 --learning_rate=0.001 --discounting=0.995 --num_envs=2048 --unroll_length=30 --batch_size=1024 --num_minibatches=32 --num_updates_per_batch=16 --backend=mjx --seed=0

[btaba]

Ok I'm not really able to spot the difference/bug for python learning/train_jax_ppo.py --env_name HumanoidStand --num_timesteps 60000000

I'm going to mark this as won't fix and update the README (just due to time constraints), but if you want to debug the diff with learning/train_jax_ppo.py that would be appreciated!