In this directory, we demonstrate learning RL agents from MuJoCo Playground environments using Brax and RSL-RL. We provide two entrypoints from the command line: python train_jax_ppo.py and python train_rsl_rl.py.

For more detailed tutorials on using MuJoCo Playground for RL, see:

1. Intro. to the Playground with DM Control Suite
2. Locomotion Environments
3. Manipulation Environments
4. Training CartPole from Vision
5. Robotic Manipulation from Vision

To train with brax PPO, you can use the train_jax_ppo.py script. This script uses the brax PPO algorithm to train an agent on a given environment.

python train_jax_ppo.py --env_name=CartpoleBalance

To train a vision-based policy using pixel observations:

python train_jax_ppo.py --env_name=CartpoleBalance --vision

Use python train_jax_ppo.py --help to see possible options and usage. Logs and checkpoints are saved in logs directory.

To train with RSL-RL, you can use the train_rsl_rl.py script. This script uses the RSL-RL algorithm to train an agent on a given environment.

python train_rsl_rl.py --env_name=LeapCubeReorient

To render the behaviour from the resulting policy:

python learning/train_rsl_rl.py --env_name LeapCubeReorient --play_only --load_run_name <run_name>

where run_name is the name of the run you want to load (will be printed in the console when the training run is started).

Logs and checkpoints are saved in logs directory.