Train Diffusion Policy đầu tiên với UMI và test trên robot arm

Đây là bài 4 trong series UMI + VLA. Bài này giả sử bạn đã có replay_buffer.zarr.zip từ bài 3.

Mục tiêu: train Diffusion Policy baseline, hiểu training curve, và chạy eval_real_umi.py để test policy trên robot thật.

Tại sao train Diffusion Policy trước, không phải VLA lớn ngay?

Câu hỏi này xuất hiện thường xuyên. Câu trả lời thực dụng:

1. Debug nhanh hơn nhiều. Diffusion Policy train trong vài tiếng trên 1 GPU. VLA 3B+ cần nhiều GPU, mất 1-2 ngày. Nếu data có vấn đề (misaligned timestamps, noisy poses, sai coordinate frame), bạn muốn phát hiện sớm với model nhỏ.

2. Baseline tham chiếu. Nếu VLA sau đó không hoạt động, bạn cần biết: vấn đề nằm ở data hay model? Baseline nhỏ cho bạn điểm so sánh.

3. VLA không phải phép màu. VLA cải thiện language conditioning và generalization, nhưng không cứu được data xấu. Nếu Diffusion Policy baseline cũng không học được gì, VLA cũng vậy.

Cài đặt môi trường training

cd universal_manipulation_interface
conda activate umi

# Kiểm tra CUDA
python -c "import torch; print(torch.cuda.is_available(), torch.version.cuda)"

# Cài thêm nếu cần
pip install wandb  # optional, cho logging

VRAM requirement:

• UMI Diffusion Policy (RGB-only, single cam): 1x 12–24 GB
• UMI Diffusion Policy (2 cameras): 1x 24 GB khuyến nghị
• Nếu thiếu VRAM: giảm batch_size trong config

Xem các config chính thức

# List tất cả configs training có sẵn
ls diffusion_policy/config/
ls diffusion_policy/config/task/

# Configs cho UMI single-arm:
# - diffusion_policy/config/task/umi.yaml
# - diffusion_policy/config/train_diffusion_unet_timm_umi_workspace.yaml
# - diffusion_policy/config/train_diffusion_transformer_umi_workspace.yaml

Đọc diffusion_policy/config/task/umi.yaml để hiểu dataset keys mà config expect. Keys phải khớp với những gì có trong replay_buffer.zarr.zip của bạn.

Train: lệnh chính thức

UMI dùng Hydra config system. Lệnh cơ bản:

python train.py --config-name=train_diffusion_unet_timm_umi_workspace \
  task.dataset.dataset_path=/absolute/path/to/replay_buffer.zarr.zip \
  training.seed=42

Nếu muốn dùng Transformer architecture thay UNet:

python train.py --config-name=train_diffusion_transformer_umi_workspace \
  task.dataset.dataset_path=/absolute/path/to/replay_buffer.zarr.zip \
  training.seed=42

Một số Hydra overrides hữu ích:

# Giảm batch size nếu OOM
python train.py --config-name=train_diffusion_unet_timm_umi_workspace \
  task.dataset.dataset_path=/path/to/replay_buffer.zarr.zip \
  training.batch_size=16 \
  training.num_epochs=500

# Đổi output dir
python train.py --config-name=train_diffusion_unet_timm_umi_workspace \
  task.dataset.dataset_path=/path/to/replay_buffer.zarr.zip \
  hydra.run.dir=outputs/umi_run_001

Xem tất cả options:

python train.py --help
python train.py --config-name=train_diffusion_unet_timm_umi_workspace --cfg job

Đọc training metrics

Khi training chạy, theo dõi các chỉ số sau:

Train loss: phải giảm liên tục. Nếu loss dao động không có xu hướng → học rate cần điều chỉnh hoặc data có vấn đề.

Val loss: nếu config có validation split, val loss phải giảm cùng chiều với train loss (không có big gap). Gap lớn = overfitting.

Epoch/step count: UMI thường cần 500–2000 epoch tùy dataset size. Đừng train quá ít.

Dấu hiệu tốt:

Epoch 50:  train_loss=0.08, val_loss=0.09
Epoch 200: train_loss=0.04, val_loss=0.05
Epoch 500: train_loss=0.02, val_loss=0.03

Dấu hiệu xấu:

Epoch 50:  train_loss=0.08
Epoch 200: train_loss=0.08  ← không giảm = learning rate sai hoặc data vấn đề
Epoch 500: train_loss=0.09  ← tăng lên = unstable training

Nếu loss không giảm sau 100 epoch:

1. Kiểm tra path đến replay buffer có đúng không
2. Kiểm tra dataset keys khớp với config
3. Thử visualize 1-2 samples từ dataset để confirm data format

Kiểm tra checkpoint

Sau training, checkpoints nằm trong output directory:

ls outputs/umi_run_001/checkpoints/
# Thường có: latest.ckpt, epoch_200.ckpt, epoch_best.ckpt, ...

# Load và inspect checkpoint
python -c "
import torch
ckpt = torch.load('outputs/umi_run_001/checkpoints/latest.ckpt', map_location='cpu')
print('Epoch:', ckpt.get('epoch', 'N/A'))
print('Train loss:', ckpt.get('train_loss', 'N/A'))
"

Dùng latest.ckpt cho inference ban đầu. Nếu repo có script evaluate offline, chạy nó trước khi test trên robot thật.

Deploy lên robot: eval_real_umi.py

UMI có script chính thức để chạy policy trên robot thật:

# Xem tất cả options
python scripts_real/eval_real_umi.py --help

# Chạy evaluation (cần robot adapter setup)
python scripts_real/eval_real_umi.py \
  --input outputs/umi_run_001/checkpoints/latest.ckpt \
  --output data/eval/run_001 \
  --frequency 10

Trước khi chạy eval thật, BẮT BUỘC:

[ ] E-stop hardware kết nối và test
[ ] Joint limits đã set trong robot config
[ ] Workspace box constraint đã khai báo
[ ] Chạy dry-run nếu robot SDK có option này
[ ] Người đứng cạnh, tay gần E-stop
[ ] Tốc độ ban đầu đặt ở mức thấp (50% hoặc ít hơn)

Robot adapter

eval_real_umi.py cần robot adapter — class Python kết nối giữa policy output và robot SDK của bạn. UMI repo có sẵn adapters cho Franka và một số robot khác trong scripts_real/. Nếu robot của bạn không có sẵn adapter, bạn cần implement interface tương ứng.

Xem các adapter có sẵn:

ls scripts_real/
# control_franka.py, control_robots.py, control_wsg_spacemouse.py, ...

Test scenarios: từ đơn giản đến phức tạp

Bắt đầu với scenario dễ nhất:

Nếu scenario đầu tiên fail (đúng vị trí vẫn fail) → vấn đề ở policy. Kiểm tra lại:

1. Camera calibration có đúng không?
2. Action convention (relative vs absolute) có khớp không?
3. Coordinate frame có nhất quán không?

Lỗi thường gặp

Nếu baseline hoạt động: bước tiếp theo

Nếu Diffusion Policy baseline hoạt động tốt (>60% success rate trên scenario dễ), bạn có thể:

1. Thu thêm data — 100-200 demo cho task hiện tại
2. Thêm diverse scenarios — nhiều lighting, position, object variation
3. Scale lên bimanual — bài 5
4. Fine-tune VLA — nếu muốn language conditioning (dùng GR00T hoặc model khác)

Đừng vội fine-tune VLA nếu baseline chưa ổn định — fix data trước.

Nguồn tham khảo

• real-stanford/universal_manipulation_interface
• Diffusion Policy paper (Chi et al., 2023)
• UMI paper (Chi et al., 2024)
• Hydra documentation

Bài viết liên quan

• Bài 3: Thu demo và chạy SLAM pipeline
• Bài 5: Lên hai tay — UMI bimanual pipeline
• Diffusion Policy: hiểu kiến trúc cho kỹ sư robotics