diff-t

This is a minimal spatial diffusion (heat diffusion) terrain generator. The terrain is used to train a PPO policy for spot to walk on varying terrain.

Important

The goal of this project is to implement a Latent Diffusion model (LDMs) or Rectified Flow model (RFs) to generate more realistic, mountainous terrain and train a PPO policy for spot to walk through the terrain to a point.

Setup

To install dependencies and activate environment:

uv sync && source .venv/bin/activate

There can be issues with running mjpython and Mujoco on macOS. Run with pip instead of uv.

python3 -m venv .venv && source .venv/bin/activate && pip install -r requirements.txt

Generate terrain

This creates a terrain file in data/terrains/terrain_sim.xml that can be used to test the robot in Mujoco.

python utils/generate_terrain.py

To see terrain in matplotlib:

python utils/viz.py

Test terrain & robot (viewing)

View spot in generated terrain, in passive mode to test the robot.

mjpython train/test_spot_sim.py

To test with custom terrain, use the --terrain argument.

mjpython train/test_spot_sim.py --terrain data/terrains/terrain_sim.xml

Train PPO policy

This will train a PPO policy on the terrain. Training can be set to different difficulties, and can be run on the cpu or a Nvidia GPU with cuda.

# easy difficulty
mjpython train/train_spot.py --difficulty easy

# standard difficulty
mjpython train/train_spot.py --difficulty standard

# challenge difficulty
mjpython train/train_spot.py --difficulty challenge

Test trained model (sim)

This will test the trained model on the terrain and view the robot in Mujoco.

mjpython train/test_spot_sim.py --model models/spot_*/final_model.zip

# test model with step limit
mjpython train/test_spot_sim.py --model models/spot_*/final_model.zip --steps 5000

File Structure

├── config/                # training configs
├── data/                  # terrain and robot assets
├── logs/                  # training logs
└── models/                # trained models
├── train/                 # training and testing
│   ├── test_spot_sim.py   # testing spot in sim
│   └── train_spot.py      # training spot
├── utils/                 # terrain gen utilities

Training difficulties

Note that standard is the default difficulty.

• easy
• standard
• challenge

Credits

TODO

• Make spot PPO train for Mujoco
• Make XML converter for terrain gen
• Fix no sensor data in Mujoco for spot training
• Fix scale issue with viz to Mujjoco for terrain gen
• Figure out better XML converter for terrain generation (look at Issac Sim or Gazebo?)
• Use modified diffusion or newer rectified flow implementation to create more realistic (mountainous) terrain
  • https://arxiv.org/pdf/2410.10766v1
• Train PPO for spot to properly walk on simple block based terrain
  • https://developer.nvidia.com/blog/closing-the-sim-to-real-gap-training-spot-quadruped-locomotion-with-nvidia-isaac-lab/
• Train PPO for spot to walk on more realistic terrain