Multitask DiT Policy trên LeRobot v0.5: 1 Model cho Nhiều Task với CLIP Text-Conditioning

Tháng 3/2026, HuggingFace phát hành LeRobot v0.5.0 — bản nâng cấp lớn nhất từ trước đến nay với hơn 200 PR được merge. Một trong những bổ sung quan trọng nhất, nhưng lại ít được chú ý so với Pi0-FAST hay Wall-X, là Multitask Diffusion Transformer (DiT) Policy. Đây là một policy diffusion thế hệ mới, chỉ ~450M tham số nhưng cho dexterity cạnh tranh được với các VLA multi-billion, và quan trọng hơn — nó có thể train 1 model duy nhất cho nhiều task khác nhau, chuyển đổi giữa các task bằng câu lệnh ngôn ngữ tự nhiên thông qua CLIP text encoder.

Bài viết này đi từ ý tưởng cốt lõi (tại sao DiT, tại sao multitask, tại sao CLIP), qua kiến trúc, đến cài đặt thực tế trên SO-100/SO-101, training command đầy đủ, debugging các failure mode phổ biến, và benchmark trên LIBERO.

Tham khảo: tài liệu chính thức Multitask DiT, LeRobot v0.5 release blog, paper TRI LBM gốc, và phân tích chi tiết của Bryson Jones.

1. Tại sao cần Multitask DiT Policy?

Vấn đề của Diffusion Policy nguyên bản

Diffusion Policy (ai-series-4) năm 2023 đã chứng minh rằng việc xem hành động robot như một chuỗi cần "denoise" cho kết quả mượt, multi-modal, và robust hơn behavior cloning truyền thống. Tuy nhiên implementation gốc dùng U-Net 1D — kiến trúc thiết kế cho convolution local, không dễ scale lên hàng trăm triệu tham số, và không có cách tự nhiên để inject ngôn ngữ vào policy.

Khi bạn muốn 1 robot biết làm 5–10 task khác nhau trên cùng phần cứng (pick cube đỏ, đặt vào bowl, lau bàn, đóng tủ, đổ nước…), cách cũ là train 5–10 model riêng biệt. Tốn data, tốn VRAM, và mỗi lần đổi task phải reload weights — không thể deploy production được.

Multitask DiT giải quyết thế nào?

Multitask DiT thay U-Net bằng Diffusion Transformer (DiT — kiến trúc nổi tiếng từ Stable Diffusion 3 và Sora), và thêm 2 conditioning streams:

• CLIP Vision Encoder xử lý ảnh từ nhiều camera (wrist + overhead).
• CLIP Text Encoder mã hoá câu lệnh task ("pick up the red cube"). Text embedding được inject vào mọi transformer block qua cross-attention.

Kết quả: 1 model ~450M tham số có thể học hàng chục task, switch giữa các task chỉ bằng đổi text prompt khi inference. Đây chính là tinh thần của TRI Large Behavior Model và blog mới nhất của Boston Dynamics về Atlas LBM, nay được port chính thức vào LeRobot.

2. Kiến trúc chi tiết

                    ┌──────────────────────────┐
   "Pick up red ─→  │  CLIP Text Encoder       │ ──┐
    cube"           │  (frozen + learnable proj)│   │
                    └──────────────────────────┘   │
                                                    ▼
   RGB images   ─→  CLIP Vision Encoder  ──→  ┌─────────────────┐
   (wrist+top)      (lr_mult = 0.1)           │  DiT Backbone    │ ──→ Predicted
                                              │  (6–8 layers,    │     action chunk
   Action noise ──→ Action token embedder ──→ │   hidden=512–768)│     (horizon=32)
   z_t                                        └─────────────────┘
                          ▲
                          │
                  Diffusion timestep t
                  hoặc Flow matching t ∈ [0,1]

Một số điểm thiết kế đáng chú ý:

• Hỗ trợ cả 2 objective: --policy.objective=diffusion (DDPM/DDIM, mặc định) hoặc --policy.objective=flow_matching (theo Boston Dynamics). Flow matching đôi khi cho action mượt hơn, nhưng không phải silver bullet.
• RoPE positional encoding mặc định cho action sequence — tốt hơn absolute pos enc.
• Vision encoder learning rate chỉ bằng 0.1× backbone LR — vì CLIP đã pretrained rất tốt, fine-tune mạnh sẽ hỏng features.
• Horizon = 32 = ~1.0 giây ở 30Hz, predict ahead 32 step nhưng chỉ thực thi 24 step rồi replan.

3. Cài đặt

Yêu cầu phần cứng cho training tối thiểu là GPU 16 GB VRAM, recommend RTX 4090 / A100 / H100 với batch size 256–320. Inference cần ≥ RTX 5070 Ti để chạy real-time 30Hz.

# 1. Tạo môi trường Python 3.12+ (LeRobot v0.5 yêu cầu)
conda create -n lerobot python=3.12 -y
conda activate lerobot

# 2. Cài LeRobot với extras cho multi_task_dit
pip install "lerobot[multi_task_dit]"

# 3. Verify
python -c "from lerobot.policies.multi_task_dit import MultiTaskDiTPolicy; print('OK')"

Extras [multi_task_dit] sẽ kéo theo transformers>=5.0 để load CLIP models (vision + text).

4. Chuẩn bị dataset cho SO-100/SO-101

LeRobot v0.5 đã consolidate codebase SO-100 và SO-101 — bạn dùng chung interface. Để train multitask, dataset phải chứa task description cho mỗi episode.

Khi record dữ liệu bằng teleop, đặt task description rõ ràng và đa dạng:

lerobot-record \
  --robot.type=so101_follower \
  --robot.port=/dev/ttyACM0 \
  --teleop.type=so101_leader \
  --teleop.port=/dev/ttyACM1 \
  --dataset.repo_id=YOUR_USER/so101_multitask \
  --dataset.single_task="pick up the red cube and place it in the bowl" \
  --dataset.num_episodes=50 \
  --dataset.fps=30

Lặp lại cho từng task khác (đổi --dataset.single_task mỗi 50 episode). Khuyến nghị tối thiểu 30–50 episode/task, 3–5 task để bắt đầu thấy hiệu ứng multitask. Nếu chỉ có 20 episode/task, policy thường idling khi inference.

Một chi tiết quan trọng: task description phải đặc trưng và khác biệt. "Pick the cube" và "Pick the block" sẽ làm model lẫn lộn. Nên dùng "pick up the red cube", "pick up the blue sponge", "open the drawer"… mỗi câu có noun cụ thể, dễ phân biệt.

Nếu bạn chưa quen quy trình teleop SO-101, tham khảo hướng dẫn LeRobot ecosystem cơ bản và tutorial sim2real SO-101 với Isaac Lab.

5. Training command đầy đủ

Đây là command training "vàng" để bắt đầu với SO-101 multitask. Tôi đã chú thích từng tham số để bạn hiểu lý do.

lerobot-train \
  --dataset.repo_id=YOUR_USER/so101_multitask \
  --output_dir=./outputs/multitask_dit_so101 \
  --job_name=so101_multitask_dit_v1 \
  \
  `# === Batch & training schedule ===` \
  --batch_size=256 \
  --steps=50000 \
  --save_freq=2000 \
  --log_freq=100 \
  --num_workers=8 \
  \
  `# === Policy core ===` \
  --policy.type=multi_task_dit \
  --policy.device=cuda \
  --policy.use_amp=true \
  \
  `# === Action horizon (30Hz) ===` \
  --policy.horizon=32 \
  --policy.n_action_steps=24 \
  --policy.n_obs_steps=2 \
  \
  `# === Objective: bắt đầu với diffusion ===` \
  --policy.objective=diffusion \
  --policy.noise_scheduler_type=DDPM \
  --policy.num_train_timesteps=100 \
  --policy.prediction_type=epsilon \
  --policy.clip_sample=true \
  \
  `# === Kiến trúc DiT ===` \
  --policy.num_layers=6 \
  --policy.hidden_dim=512 \
  --policy.num_heads=8 \
  --policy.dropout=0.1 \
  --policy.use_rope=true \
  \
  `# === Vision encoder ===` \
  --policy.vision_encoder_name=openai/clip-vit-base-patch16 \
  --policy.image_resize_shape=[256,256] \
  --policy.image_crop_shape=[224,224] \
  --policy.image_crop_is_random=true \
  --policy.vision_encoder_lr_multiplier=0.1 \
  \
  `# === Text encoder (CLIP) ===` \
  --policy.text_encoder_name=openai/clip-vit-base-patch16 \
  \
  `# === Optimizer ===` \
  --policy.optimizer_lr=2e-5 \
  --policy.optimizer_weight_decay=0 \
  \
  `# === Push checkpoint ===` \
  --policy.push_to_hub=true \
  --policy.repo_id=YOUR_USER/multitask-dit-so101 \
  --wandb.enable=true \
  --wandb.project=multitask_dit_so101

Với 5 task × 50 episode trên RTX 4090, training mất khoảng 8–12 giờ. Loss thường flatlining sau ~20k steps nhưng bạn vẫn nên train tiếp đến 50k–100k — model multitask cần thêm time để language steerability ổn định, dù loss đã phẳng.

Tăng tốc inference bằng DDIM

Khi training xong, đổi sampler sang DDIM với ít bước:

--policy.noise_scheduler_type=DDIM \
--policy.num_inference_steps=10

DDIM với 10 step ≈ DDPM với 100 step về chất lượng, nhanh hơn 10×.

6. Inference & deploy trên SO-101 thật

lerobot-eval \
  --robot.type=so101_follower \
  --robot.port=/dev/ttyACM0 \
  --policy.path=YOUR_USER/multitask-dit-so101 \
  --policy.device=cuda \
  --eval.task="pick up the red cube and place it in the bowl" \
  --eval.n_episodes=10

Tham số --eval.task chính là text prompt đưa vào CLIP text encoder. Bạn có thể đổi sang task khác trong cùng training set mà không cần reload model:

--eval.task="open the drawer slowly"

Đây là điểm khác biệt lớn so với ACT hay Diffusion Policy nguyên bản — bạn có 1 binary duy nhất, switch task bằng natural language.

7. Failure modes phổ biến & cách debug

Phần này là kinh nghiệm thực chiến, đáng giá nhất khi bạn tự train.

Failure 1: Robot đứng yên (idling / no motion)

Triệu chứng: action output gần như zero, robot không cử động hoặc chỉ rung nhẹ.

Nguyên nhân thường gặp:

• Dataset quá nhỏ (< 200 examples tổng).
• Nhiều task quá giống nhau (pick cube đỏ vs pick cube xanh) → model bám vào vision, bỏ qua text.
• Loss đã flatline nhưng train chưa đủ lâu.

Cách fix:

• Tăng dataset gấp đôi cho tới khi vượt 300 examples.
• Train đến 100k steps dù loss đã phẳng.
• Diversify text instruction: "grasp the crimson block" thay vì lặp "pick the cube".

Failure 2: Robot làm sai task

Triệu chứng: ra lệnh "pick cube đỏ" nhưng robot lại lau bàn.

Nguyên nhân: task description bị overlap quá nhiều, hoặc task lệch khỏi distribution có trong dataset.

Cách fix:

• Verify text instruction phải đặc thù.
• Re-weight task bị bỏ qua khi sampling.
• Fine-tune thêm vài hundred steps trên task khó.

Failure 3: Training instability

Triệu chứng: loss tăng giảm hỗn loạn, NaN sau vài k steps.

Cách fix:

• Hạ learning rate từ 2e-5 xuống 1e-5.
• Tăng batch size — dưới 64 sẽ rất bất ổn.
• Verify image normalization đúng range của CLIP (mean/std từ ImageNet).

8. Benchmark trên LIBERO

Đội LeRobot công bố Multitask DiT đạt 90.6% trung bình trên LIBERO với config 8 layers, hidden 768, horizon 48, train 100k steps batch 320:

So sánh: Diffusion Policy gốc ~78%, ACT ~73%, OpenVLA 7B (multi-billion params) ~76%. Multitask DiT với chỉ 450M params đã vượt nhiều VLA lớn hơn nhiều lần. Đây là tin cực vui cho nhóm chỉ có RTX 4090 / 5090.

9. Khi nào dùng Multitask DiT, khi nào không?

10. Kết luận

Multitask DiT trong LeRobot v0.5 là điểm sweet-spot rất khó tìm: đủ mạnh để học hàng chục task qua ngôn ngữ, đủ nhỏ để chạy được trên GPU consumer, và đã được package gọn trong lerobot-train. Đối với cộng đồng robotics Việt Nam đang xây sản phẩm thực tế với SO-100/SO-101 hoặc cánh tay chi phí thấp tương tự, đây là baseline mạnh nhất hiện có mà không cần thuê A100 cluster.

Lời khuyên cuối cùng: đừng cố train multitask ngay từ đầu. Hãy train 1 task trước với Multitask DiT để verify pipeline, rồi mới thêm task. Đa số "policy không học được" là do dataset bẩn, không phải do model.

Bài viết liên quan

• LeRobot v0.5: Tất tần tật về bản phát hành mới
• Diffusion Policy: nền tảng của Multitask DiT
• SO-101 + Isaac Lab Sim2Real với LeRobot