ADD Gym

An ML training repo for imitation training using the ADD learning robot, G1 robot, and physics simulators: Genesis and Mujoco Warp.

Overview • Getting Started • Deployment • SageMaker Training • Acknowledgements

Overview

An implementation of the ADD learning algorithm with additional optimizations, trained on high-performance simulators: Genesis and Mujoco Warp, and published freely available model weights. This repo makes imitation training scalable, supporting both local development cycles and large-scale distributed training on AWS SageMaker.

Improvements from Original Code

This code makes the following optimizations:

• Removed feature engineering of link positions. Training can rely on joint rotations and root position/rotation and optionally, velocities.
• Pre-calculating ground truth, so rotations and positions don't have to be re-calculated at every timestep
• Removed usage of Nvidia's propietary Isaac Sim and deprecated Isaac Gym. Replaced it with Genesis and Mujoco Warp, which are much faster at startup.
• Adaptive Sampler to focus on harder movements in the clip

Available model weights:

More coming soon...

• Walk | Motion file: assets/motions/walk1_subject1_trimmed.motion
• Dance | Motion file: assets/motions/dance1_subject3.motion
• Fight | Motion file: assets/motions/fight1_subject2_trimmed.motion

To test out one of the available model weights

1. Clone the repo

2. Download the model weights from Hugging Face

3. git checkout [tag]

   Ex: git checkout walk

4. uv run python -m add_gym.main --config-name test +resume_path=<path/to/model.pt> task.motion_file=<path/to/file.motion>

Getting Started

Prerequisites

• Python 3.10 is required.
• uv: We use uv for dependency management.

Installation

1. Install uv (if not already installed):

   curl -LsSf https://astral.sh/uv/install.sh | sh

2. Sync dependencies:

   uv sync

   This will create a virtual environment and install all necessary dependencies, including genesis-world.

Local Development

To run the training loop locally:

uv run python -m add_gym.main

You can customize the run using Hydra configuration overrides:

uv run python -m add_gym.main engine.num_envs=4096

Example for Multi-GPU training:

PYTHONPATH=$PWD uv run torchrun --standalone --nproc_per_node=8 add_gym/main.py -
-config-name train.8xh100

Mujoco Warp

Caution

Mujoco Warp integration is currently under development and is not yet rigorously tested!

To use Mujoco Warp, you'll first need to install the extra dependencies:

uv sync --all-extras

And then, simply change the engine config group to mjwarp. Example:

uv run python -m add-gym.main engine=mjwarp

Deployment

CI/CD Pipeline

The project includes GitHub Actions workflows that automatically:

• build and push the Docker image suitable for training
• optionally, tag to push your trained model weights to Hugging Face

Crucial Setup: To ensure the build pipeline works, you must set the BASE_IMAGE secret in your GitHub repository settings.

• Value: This should point to a valid Genesis or MjWarp base image (e.g., account_id.dkr.ecr.region.amazonaws.com/genesis:latest).
• The workflow checks out the code, builds the training image using this base, and pushes it to your configured ECR repository.

In addition to BASE_IMAGE, you also need to set the following secrets:

• AWS_REGION
• S3_BUCKET
• HF_TOKEN (optional, if you want to push your trained model to Hugging Face)

SageMaker Training

I provide a specialized workflow for submitting training jobs to SageMaker without manually handling config files. The submission script converts your local Hydra configuration into command-line arguments for the remote job. This requires Cloud deployable resources, which are specified in terraform.

1. Configure

• Update deploy/sagemaker-job-config.yaml with job specs.
• Update deploy/train-config.yaml with training overrides.

2. Customization

Update your terraform config in terraform/terraform.tfvars

You can override job parameters in the Hydra config.

• For cloud-executed training runs, you can override it in deploy/train-config.
• You can also specify specific parameters the hydra files in add-gym/configs.

3. Train

Any commit pushed to a branch named models or models/* will trigger a job to build and Sagemaker job.

Acknowledgements

This project was not possible if it weren't for:

• Adversarial Differential Discriminators (Original Implementation)
• LAFAN1 Retargeting Dataset
• Genesis
• Mujoco Warp