Fine-tuning

Fine-tuning GR00T on custom datasets

This guide demonstrates how to fine-tune GR00T on your own robot data and configuration. We provide a complete example for the SO-100 robot under examples/SO100, which uses demo_data/cube_to_bowl_5 as the demo dataset.

Key parameters

Parameter	Description
--base-model-path	Path to the pre-trained base model checkpoint
--dataset-path	Path to your training dataset
--embodiment-tag	Tag to identify your robot embodiment
--modality-config-path	Path to user-specified modality config (required only for NEW_EMBODIMENT tag)
--output-dir	Directory where checkpoints will be saved
--save-steps	Save checkpoint every N steps
--max-steps	Total number of training steps
--use-wandb	Enable Weights & Biases logging for experiment tracking
--global-batch-size	Global batch size across all GPUs
--color-jitter-params	Color jitter augmentation parameters
--dataloader-num-workers	Number of data loading workers

Recommended configuration

For optimal results, maximize your batch size based on available hardware and train for a few thousand steps.

Hardware performance

• We recommend using 1 H100 node or L40 node for optimal fine-tuning performance
• Other hardware configurations (e.g., A6000) will also work but may require longer training time
• Optimal batch size depends on your hardware and which model components are being tuned

Training variance

Users may observe some variance in post-training results across runs, even when using the same configuration, seed, and dropout settings. In our experiments, we have observed performance differences as large as 5-6% between runs. This variance may be attributed to non-deterministic operations in image augmentations or other stochastic components.

Dataloader optimization

When training a model, you can optimize the dataloading speed vs memory usage via various command line arguments:

python gr00t/experiment/launch_finetune.py \
    ... \
    --num-shards-per-epoch 100 \
    --dataloader-num-workers 2 \
    --shard-size 512

If VRAM is limited, you can reduce all the numbers above to reduce memory usage.

To ensure more IID during sampling of shards, you can reduce the episode_sampling_rate to 0.05 or lower.

Advanced configuration

For more extensive fine-tuning configuration, use gr00t/experiment/launch_train.py instead to launch the training process with full control over all training parameters.