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_sizetrong 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:
- Kiểm tra path đến replay buffer có đúng không
- Kiểm tra dataset keys khớp với config
- 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:
| Scenario | Mục đích |
|---|---|
| Object ở đúng vị trí demo | Sanity check cơ bản |
| Object lệch 5cm so với demo | Test spatial generalization nhỏ |
| Object lệch 10-15cm | Test spatial generalization lớn hơn |
| Lighting khác nhẹ | Test visual robustness |
| Object màu khác cùng shape | Test object generalization |
Nếu scenario đầu tiên fail (đúng vị trí vẫn fail) → vấn đề ở policy. Kiểm tra lại:
- Camera calibration có đúng không?
- Action convention (relative vs absolute) có khớp không?
- Coordinate frame có nhất quán không?
Lỗi thường gặp
| Lỗi | Nguyên nhân | Fix |
|---|---|---|
| Loss không giảm | Path sai hoặc data format sai | Print 1 sample từ dataset, compare với config expected keys |
| OOM khi train | Batch size quá lớn | Giảm training.batch_size xuống 8-16 |
| Robot di chuyển loạn | Action scale sai | Kiểm tra normalization trong config; thêm workspace constraint |
| Policy không reproduce demo | Gripper width sai trong data | Re-run ArUco detection với better settings |
| Training unstable (loss NaN) | Learning rate quá cao hoặc gradient explode | Giảm lr, thêm gradient clipping |
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ể:
- Thu thêm data — 100-200 demo cho task hiện tại
- Thêm diverse scenarios — nhiều lighting, position, object variation
- Scale lên bimanual — bài 5
- 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
