docs/TRAINING.md

Training & Evaluation

Dataset

GEM is trained on Bones RigPlay-1, an internal NVIDIA synthetic dataset. Bones RigPlay-1 is not publicly released.

Expected directory layout:

inputs/
└── metrosim_data_A2G_Bones2_DH/
    ├── train/
    ├── val/
    └── mocap/

Update data_root in configs/train_datasets/metrosim_dh_train.yaml if your data is located elsewhere.

Training

Single-GPU

python scripts/train.py exp=gem_soma_regression

Multi-GPU (DDP)

python scripts/train.py exp=gem_soma_regression pl_trainer.devices=4

Key Config Settings

The main experiment config is configs/exp/gem_soma_regression.yaml:

Setting	Value
Body model	SOMA
Max steps	500,000
Precision	16-mixed
Optimizer	AdamW (lr=2e-4)
Gradient clipping	0.5
Validation interval	Every 3,000 steps

W&B Logging

Logging uses Weights & Biases by default. To disable:

python scripts/train.py exp=gem_soma_regression use_wandb=false

Evaluation

python scripts/train.py exp=gem_soma_regression task=test

This runs evaluation on the MetroSim validation split and reports per-frame SOMA body pose and global translation accuracy metrics.

Hydra Config System

GEM uses Hydra for configuration management. Key config groups:

Group	Description
exp/	Experiment configs (e.g., gem_soma_regression)
model/	Model architecture
network/	Network details (denoiser, regression)
pipeline/	Training pipeline (loss weights, features)
train_datasets/	Training data configs
test_datasets/	Evaluation data configs

Override any config value from the command line:

python scripts/train.py exp=gem_soma_regression pl_trainer.max_steps=100000 optimizer.lr=1e-4