Awesome

Maximum Diffusion Reinforcement Learning Repository

System Requirements

Simulations were run on Ubuntu 18.04 and Ubuntu 20.04 operation systems
Simulations must be run with Python 3 (tested Python 3.6, Python 3.8, and Python 3.9)
PyTorch can be installed with CUDA or CPU only (tested versions 1.7, 1.11, and 2.0)
A dockerfile is provided as an alternate to the installation steps below. After installing and running docker desktop on your system, you can instantiate the maxdiff test enviroment with source run_docker.sh on Ubuntu or .\run_docker.bat on Windows.

Installation

MuJoCo must be installed prior to installing python mujoco_py package.
1. Download free license from MuJoCo website (visit https://www.roboti.us/license.html and click on "Activation Key")
2. Download the MuJoCo version 2.0 binaries for Linux (https://www.roboti.us/download/mujoco200_linux.zip)
3. Unzip the downloaded mujoco200 directory into ~/.mujoco/mujoco200, and place your license key (the mjkey.txt file) at ~/.mujoco/mjkey.txt.
4. Add MuJoCo binary path to your environment path export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.mujoco/mujoco200/bin
Make sure you have the following libaries installed sudo apt install libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
Python dependencies can be installed with pip install -r requirements.txt
- Note: specified mujoco_py version must be used for code to work. Other versions can be modified at your own discretion (e.g. PyTorch with CPU only or with a specific CUDA version)

Running the code

To train a policy, modify the relevant yaml file for your test scenario (config/<method>.yaml). Then, run the following with the desired arguments (possible arguments can be viewed with python train.py --help)

python3 train.py --env PointMass2D_DoubleIntEnv --method maxdiff --seed 13 --beta 0.1
python3 train.py --env SwimmerEnv_v3 --method maxdiff --seed 13

To playback the learned policy, run the following with the arguments matching your training configuration (possible arguments can be viewed with python train.py --help)

python3 enjoy.py --env PointMass2D_DoubleIntEnv --method maxdiff --mod _H30_alpha5 --beta 0.1
python3 enjoy.py --env SwimmerEnv_v3 --method maxdiff --seed 13 --mod _H40_alpha100

Results can also be visualized in the jupyter-notebooks provided in the notebooks folder.

Run time (CPU)

Environment	Steps	Horizon	Method	Approx. Runtime
SwimmerEnv_v3	1,000,000	40	MaxDiff & MPPI	1 day
SwimmerEnv_v3	1,000,000	N/A	SAC	4 hours
HalfCheetahEnv_v3	1,000,000	10	MaxDiff & MPPI	19 hours
HalfCheetahEnv_v3	1,000,000	N/A	SAC	4 hours
AntEnv_v3	1,000,000	20	MaxDiff & MPPI	1 day 20 hours
AntEnv_v3	1,000,000	N/A	SAC	10 hours

Repository Contents

.
├── config/                       # contains environment-specific test setup yaml files
├── data/                         # contains models and data to run notebooks
├── envs/
│   ├── __init__.py
│   ├── build_env.py              # helper function to build environments for training and testing
│   ├── light_swimmer.xml         # specifies alternate physics model (light tail)
│   ├── normalized_actions.py     # wrapper for gym environments to normalize action space
│   ├── pointmass_lib/            # folder contains functions to generate point mass environment and for rendering pointmasss figures
│   ├── replay_buffer.py          # memory buffer for all training all policies
│   └── wrappers.py               # wrappers for gym environments to modify reward functions
├── mpc_lib/                      # model predictive control policy files
│   ├── __init__.py
│   ├── entropy.py                # helper functions for MaxDiff
│   ├── max_diff.py               # MaxDiff policy module        
│   ├── model.py                  # model for MPPI and MaxDiff
│   ├── mppi.py                   # MPPI policy module
│   ├── optimizer.py              # model optimizer for MPPI and MaxDiff
│   └── utils.py                  # helper functions for max_diff and mppi
├── notebooks/                    # jupyter-notebooks for plotting reward curves
├── sac_lib/                      # model-free policy files
│   ├── __init__.py
│   ├── policynetwork.py          # policy network
│   ├── sac_orig.py               # SAC training module (with optimizer)
│   └── sac_networks.py           # SAC training networks
├── Dockerfile                    # Contains setup instructions for docker
├── enjoy.py                      # replay of trained policies (visualization only)
├── eval.py                       # test trained policies (collect data for analysis)
├── LICENSE
├── README.md
├── requirements.txt              # python dependencies
├── run_docker.bat                # run file to start docker container on Ubuntu systems
├── run_docker.sh                 # run file to start docker container on Windows systems
├── train.py                      # main training function
└── utils.py                      # helper functions

Copyright and License

The implementations of MaxDiff contained herein are copyright (C) 2024 - 2025 by Allison Pinosky and Todd Murphey and are distributed under the terms of the GNU General Public License (GPL) version 3 (or later). Please see the LICENSE for more information.

Contact: apinosky@u.northwestern.edu

Lab Info:
Todd D. Murphey
https://murpheylab.github.io/
Northwestern University