DIAL-MPC: Diffusion-Inspired Annealing For Legged MPC

This repository contains the code (simulation and real-world experiments with minimum setup) for the paper "Full-Order Sampling-Based MPC for Torque-Level Locomotion Control via Diffusion-Style Annealing".

DIAL-MPC is a sampling-based MPC framework for legged robot full-order torque-level control with both precision and agility in a training-free manner. DIAL-MPC is designed to be simple and flexible, with minimal requirements for specific reward design and dynamics model. It directly samples and rolls out in physics-based simulations (Brax) and does not require reduced-order modeling, linearization, convexification, or predefined contact sequences. That means you can test out the controller in a plug-and-play manner with minimum setup.

News

09/25/2024: 🎉 DIAL-MPC is released with open-source codes! Sim2Real pipeline coming soon!

dial-mpc-demos_compressed.mp4

Simulation Setup

Install `dial-mpc`

Important

We recommend Ubuntu >= 20.04 + Python >= 3.10 + CUDA >= 12.3. You can create a mamba (or conda) environment before proceeding.

Our environment is Ubuntu 22.04 + Python 3.10 + CUDA 12.6.

git clone https://github.com/LeCar-Lab/dial-mpc.git --depth 1
cd dial-mpc
pip3 install -e .

Synchronous Simulation

Run Examples

List available examples:

dial-mpc --list-examples

Run an example:

dial-mpc --example unitree_h1_jog

After rollout completes, go to 127.0.0.1:5000 to visualize the rollouts.

Asynchronous Simulation

The asynchronous simulation is meant to test the algorithm before Sim2Real.

List available examples:

dial-mpc-sim --list-examples

Run an example:

In terminal 1, run

dial-mpc-sim --example unitree_go2_seq_jump_deploy

This will open a mujoco visualization window.

In terminal 2, run

dial-mpc-plan --example unitree_go2_seq_jump_deploy

Deploy in Real

🚧 Check back in late Sep. - early Oct. 2024 for real-world deployment pipeline on Unitree Go2.

Writing Custom Environment

If custom robot model is needed, Store it in dial_mpc/models/my_model/my_model.xml.
Import the base environment and config.
Implement required functions.
Register environment.
Configure config file.

Example environment file (my_env.py):

from dataclasses import dataclass

from brax import envs as brax_envs
from brax.envs.base import State

from dial_mpc.envs.base_env import BaseEnv, BaseEnvConfig
import dial_mpc.envs as dial_envs

@dataclass
class MyEnvConfig(BaseEnvConfig):
    arg1: 1.0
    arg2: "test"

class MyEnv(BaseEnv):
    def __init__(self, config: MyEnvConfig):
        super().__init__(config)
        # custom initializations below...

    def make_system(self, config: MyEnvConfig) -> System:
        model_path = ("my_model/my_model.xml")
        sys = mjcf.load(model_path)
        sys = sys.tree_replace({"opt.timestep": config.timestep})
        return sys

    def reset(self, rng: jax.Array) -> State:
        # TODO: implement reset

    def step(self, state: State, action: jax.Array) -> State:
        # TODO: implement step

brax_envs.register_environment("my_env_name", MyEnv)
dial_envs.register_config("my_env_name", MyEnvConfig)

Example configuration file (my_env.yaml):

# DIAL-MPC
seed: 0
output_dir: dial_mpc_ws/my_model
n_steps: 400

env_name: my_env_name
Nsample: 2048
Hsample: 25
Hnode: 5
Ndiffuse: 4
Ndiffuse_init: 10
temp_sample: 0.05
horizon_diffuse_factor: 1.0
traj_diffuse_factor: 0.5
update_method: mppi


# Base environment
dt: 0.02
timestep: 0.02
leg_control: torque
action_scale: 1.0

# My Env
arg1: 2.0
arg2: "test_2"

Run the following command to use the custom environment in synchronous simulation. Make sure that my_env.py is in the same directory from which the command is run.

dial-mpc --config my_env.yaml --custom-env my_env

You can also run asynchronous simulation with the custom environment:

# Terminal 1
dial-mpc-sim --config my_env.yaml --custom-env my_env

# Terminal 2
dial-mpc-plan --config my_env.yaml --custom-env my_env

Rendering Rollouts in Blender

If you want better visualization, you can check out the render branch for the Blender visualization examples.

Acknowledgements

This codebase's environment and RL implementation is built on top of Brax.
We use Mujoco MJX for the physics engine.
Controller design and implementation is inspired by Model-based Diffusion.

BibTeX

If you find this code useful for your research, please consider citing:

@misc{xue2024fullordersamplingbasedmpctorquelevel,
      title={Full-Order Sampling-Based MPC for Torque-Level Locomotion Control via Diffusion-Style Annealing}, 
      author={Haoru Xue and Chaoyi Pan and Zeji Yi and Guannan Qu and Guanya Shi},
      year={2024},
      eprint={2409.15610},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2409.15610}, 
}

EGalahad/dial-mpc