minimal-mjx 0.1.2

Short description

MINIMAL-MJX

WARNING:jax._src.xla_bridge:An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.

This repository represents starter code for MJX-based RL from the Dynamic Mobility Lab. The code has been built of off MuJoCo Playground and tailored to make policy training, saving, and evaluation easy and implements some nice-to-haves.

1. Swappable backend, allowing for fast evaluation using numpy/C++ MuJoCo and fast training using JaX/MJX
2. Base environment class that contains generic reward functions and useful functions
3. Configuration files that allow you to set parameters for your envs + for PPO

There exists a requirements.txt in the home directory which includes the packages needed to run this code.

Steps to take to 'do RL'!

We will work of off the Cheetah environment from Deepmind's "DM Control". This guide is not meant to be comprehensive, but is hopefully enough information to fill in the gaps as you develop and look around the codebase.

0. The setup basics

   Each environment has a reset and step function, which primarily operate off of a Markov Decision Process (MDP) state. In our case, this is represented as a class that looks like this

       @dataclass
       class MujocoState:
           data: mujoco.MjData
           obs: np.ndarray
           reward: float
           done: bool
           metrics: dict
           info: dict
   

   Note that you should treat info as a 'carry' variable between consecutive step functions. Store things here you need to use in future timesteps that are not ordinarily stored in data (i.e. store a history of system states x).

1. Simulate your environment to make sure things look good

   Environments should be simulated and inspected before training. Once you are done developing your env, make a configuration file to specify variables specific to training runs, like PPO or reward weight parameters.

   For the Cheetah, one is already created. To simulate the environment, run

   python3 -m envs.simulate envs/dmcontrol/config/cheetah.yaml
   

   Under visualization, you should see a metrics plot as well as a video of your environment. Note that the policy here is simulated to output all zeros (see envs/simulate.py). Inspect these files to make sure things look good.

   For extra good measure, change the backend parameter in your config file to jnp to make sure your environment is JaX compatible.

2. Train a policy

   Once your environment is ready and JaX compatible, you can train by running

   python3 -m learning.training.begin_run envs/dmcontrol/config/cheetah.yaml
   

   Note that this will train directly in your terminal session. A bash script has been provided at learning/training/train.sh that opens a tmux terminal for this (useful if you want to train for long periods of time). Make sure to change your conda environment name in this script.

   Within your save directory mentioned in your config file, a new directory will be created in which intermediate and final training results will be saved.

3. Rollout your trained policy

   Once your policy is finished training, simply run

   python3 -m eval.rollout_policy envs/dmcontrol/config/cheetah.yaml
   

   This will rollout your policy, plot metrics of your reward function, and save a video of your result. Make sure to use an np backend for quick evaluation!

   If you would like to rollout older policies, check the directory of your past training runs. There should be a config there, and you can simply replace the config above with that one (running that specific older policy). Note that your code might have changed, so this config holds a git commit that references when the training was run.