Skip to content

Tsinghua-MARS-Lab/OMG

Repository files navigation

OMG: Omni-Modal Motion Generation for Generalist Humanoid Control

Official repository for OMG: Omni-Modal Motion Generation for Generalist Humanoid Control.

arXiv Website

OMG teaser

  • Training Code
  • Evaluation Pipelines
    • Pretrained Evaluator Checkpoints
    • Benchmark Code
  • Inference Code
  • Sim-to-Real Deployment Code
  • OMG-Data
  • Pretrained Checkpoints

Pipeline

The usual end-to-end workflow is:

  1. Install the environment.
  2. Download OMG-Data, model checkpoints, and HoloMotion artifacts.
  3. Materialize OMG-Data for faster training.
  4. Compute normalization stats.
  5. Train a diffusion model.
  6. Export ONNX for TensorRT/CUDA inference.
  7. Run generation or full pipeline modes.
  8. Run benchmarks.
  9. Deploy to a G1 robot when needed.

1. Install

cd /path/to/OMG
make venv
source .venv/bin/activate
make install

export PYTHONPATH=src
export TOKENIZERS_PARALLELISM=false

For China mainland networks:

make install-cn

See Installation for manual uv commands and optional task-specific extras.

2. Download Data and Artifacts

OMG-Data and pretrained OMG checkpoints will be released on Hugging Face:

Text-conditioned training and generation require the Hugging Face t5-base text encoder. By default, configs load it from ${OMG_MODELS_ROOT}/t5-base-local. Download t5-base for offline runs, or override model.text_encoder.model_name with another local path or Hugging Face model id.

HoloMotion weights are not redistributed by OMG. Download HoloMotion models from the official HoloMotion repository or HoloMotion Hugging Face artifacts.

Recommended local layout:

data/OMG-Data/
  omg_data/
  materialized/
models/
  generation/
  evaluator/
  t5-base-local/
  holomotion/
    motion_tracking/model.onnx
    velocity_tracking/model.onnx

Set explicit roots when using external storage:

export OMG_DATA_ROOT=/path/to/OMG-Data
export OMG_MATERIALIZED_ROOT=/path/to/OMG-Data/materialized
export OMG_MODELS_ROOT=/path/to/OMG-models

3. Materialize Data

Materialization precomputes fixed-window training shards. It is recommended for full training because it removes repeated source parsing and feature assembly from the training loop.

You can download precomputed materialized OMG-Data into OMG_MATERIALIZED_ROOT, or generate the same layout locally from source OMG-Data:

scripts/materialize_omg_data.sh --overwrite

Train with materialized data by using:

data=omg_data_materialized

For small debugging runs or custom tiny datasets, source data can be used directly with:

data=omg_data

4. Compute Stats

Compute normalization statistics before training. The default representation config expects the generated stats file at:

assets/stats/g1_125d_stats.json
PYTHONPATH=src python -m omg.cli.generation.compute_stats \
  --data-config configs/generation/data/omg_data.yaml \
  --representation-config configs/generation/representation/125d.yaml \
  --paths-config configs/generation/paths/default.yaml \
  --output assets/stats/g1_125d_stats.json

Recompute this file whenever the training data, representation, sequence length, or preprocessing changes.

5. Train

Example 50M training run:

CUDA_VISIBLE_DEVICES=0,1,2,3 \
PYTHONPATH=src python -m omg.cli.generation.train \
  exp=50m \
  data=omg_data_materialized \
  trainer=4gpu \
  logger=wandb \
  exp_name=50m_release_train

Model-size configs are available under configs/generation/exp/:

50m.yaml  100m.yaml  300m.yaml  500m.yaml  1b.yaml

See Training for resume, initialization, and config details.

6. Export ONNX

Export a TensorRT-compatible denoiser step from a checkpoint:

PYTHONPATH=src python -m omg.cli.generation.export_onnx \
  --exp 50m \
  --ckpt_path outputs/50m_release_train/checkpoints/last.ckpt \
  --output models/generation/onnx/50m/last_denoiser_step.onnx \
  --batch_size 2 \
  --device cuda

The exporter writes a metadata sidecar next to the ONNX file. Runtime planners use it to recover sequence length, condition dimensions, representation, and diffusion settings.

7. Generate and Track

Offline async generation with HoloMotion tracking:

PYTHONPATH=src python -m omg.cli.pipeline.main \
  --mode async \
  --diffusion-onnx models/generation/onnx/50m/last_denoiser_step.onnx \
  --holomotion-onnx models/holomotion/motion_tracking/model.onnx \
  --seed-motion /path/to/seed_motion.npz \
  --condition-sequence "text: walk forward" \
  --num-frames 300 \
  --video \
  --output-root outputs_pipeline

Supported pipeline modes:

  • diffusion-only
  • tracker-only
  • sync
  • async
  • offline-track

See Generation and Tracking.

8. Benchmark

Benchmarks that report evaluator-based distribution and retrieval metrics use a pretrained evaluator checkpoint. It will be released at:

https://huggingface.co/<org>/OMG-Evaluator

Recommended local path:

models/evaluator/pretrained.ckpt

Prepare fixed benchmark sample manifests:

PYTHONPATH=src python -m omg.cli.evaluation.prepare_samples \
  --data omg_data \
  --exp 50m \
  --output_dir outputs/benchmark_samples

Run text, audio, human-reference, or artifact benchmarks:

PYTHONPATH=src python -m omg.cli.generation.benchmark text \
  --exp 50m \
  --ckpt_path outputs/50m_release_train/checkpoints/last.ckpt \
  --evaluator_checkpoint models/evaluator/pretrained.ckpt \
  --output_dir outputs/benchmarks/50m_text

See Benchmark for modality-specific commands and tracker-executed evaluation.

9. Deploy

Realtime deployment uses:

  • HoloMotion deployment on the G1 Orin.
  • OMG realtime planner server on a GPU workstation.
  • OMG real bridge on the G1 Orin.

For real-robot deployment, prefer the HoloMotion velocity-tracking model:

models/holomotion/velocity_tracking/model.onnx

See Realtime G1 Deployment for the full launch sequence.

Documentation

License

This project is released under the MIT License.

Releases

No releases published

Packages

 
 
 

Contributors

Languages