IOAI ACT Trainer

简体中文

Train ACT from LeRobot v2/v3 datasets via Docker. Images: Docker Hub ioaitech/train_act. Source: ioai-tech/train_act.

Requires Linux, an NVIDIA GPU, Docker, and the NVIDIA Container Toolkit.

Images

Tag	PyTorch	CUDA
cuda / cuda121	2.5.1	12.1
cuda126	2.10.0	12.6
cuda130	2.10.0	13.0

Platform: linux/amd64. Prefer cuda121 for wider driver compatibility.

Quick start

docker pull ioaitech/train_act:cuda121

mkdir -p ./act-output
docker run --rm --gpus all \
  -v /path/to/lerobot_dataset:/data/input:ro \
  -v "$(pwd)/act-output":/data/output \
  ioaitech/train_act:cuda121 \
  --run_name act_smoke \
  --num_epochs 1 \
  --batch_size 2 \
  --max_episodes 2

Mount a LeRobot dataset at /data/input (meta/info.json required). State/action dimensions are inferred from metadata. Checkpoints go to /data/output.

More examples

# Single GPU
docker run --rm --gpus '"device=0"' --shm-size=8g \
  -v /data/my_dataset:/data/input:ro \
  -v /data/my_output:/data/output \
  ioaitech/train_act:cuda121 \
  --run_name experiment_01 \
  --num_epochs 5000 \
  --batch_size 64

# Multi-GPU
docker run --rm --gpus all --ipc=host \
  -v /data/my_dataset:/data/input:ro \
  -v /data/my_output:/data/output \
  ioaitech/train_act:cuda126 \
  --run_name experiment_ddp \
  --gpus 0,1 \
  --batch_size 64

# Select cameras
docker run --rm --gpus all \
  -v /data/my_dataset:/data/input:ro \
  -v /data/my_output:/data/output \
  ioaitech/train_act:cuda121 \
  --camera_keys observation.images.camera_high,observation.images.camera_left_wrist \
  --camera_names camera_high,camera_left_wrist

Common flags: --num_epochs, --batch_size, --learning_rate, --chunk_size, --save_interval, --max_episodes, --keep_converted_hdf5. Full list: docker run --rm ioaitech/train_act:cuda121 --help.

License

MIT. See THIRD_PARTY_NOTICES.md for upstream ACT/DETR.