MA-VLA: VLA đa cánh tay trên LeRobot

MA-VLA giải quyết vấn đề gì?

MA-VLA, viết tắt của Multi-Arm Vision-Language-Action Model, là một framework mới cho robot nhiều cánh tay, được giới thiệu trong paper MA-VLA: Multi-Arm Vision-Language-Action Model for Collaboration and Compositional Generalization và repo chính thức zhangzaibin/future-robots. Điểm quan trọng của paper không phải là "thêm nhiều camera rồi train model to hơn". Điểm quan trọng là cách biểu diễn nhiệm vụ: thay vì đưa cho robot một instruction chung kiểu "stack the cubes", MA-VLA tách nhiệm vụ thành các atomic actions ở mức trung gian, rồi gán từng atomic prompt cho từng arm.

Với robot single-arm, một instruction tổng quát đôi khi vẫn đủ. Ví dụ: "pick up the red cube" chỉ cần một tay tìm khối đỏ, kẹp, nâng lên. Nhưng với dual-arm, triple-arm hoặc four-arm, cùng một câu lệnh có thể ẩn nhiều quyết định khó:

• Arm nào giữ object?
• Arm nào chuyển object?
• Arm nào chờ?
• Thứ tự phối hợp là gì?
• Nếu đổi vị trí object hoặc đổi vai trò arm, policy có còn chạy được không?

Các VLA end-to-end như Pi0 có thể học behavior rất mạnh khi train/test cùng một kiểu phối hợp. Vấn đề xuất hiện khi test-time yêu cầu một collaboration pattern chưa từng xuất hiện trong training data. Paper gọi năng lực này là multi-arm compositional generalization: model phải tái tổ hợp những atomic actions đã biết thành một chiến lược phối hợp mới.

Nếu bạn đã quen với Hands-on: Fine-tune OpenVLA với LeRobot, MA-VLA có thể được hiểu như một bước tiếp theo cho multi-arm VLA. LeRobot-format data vẫn là phần xương sống để lưu observation, state và action. Khác biệt nằm ở chỗ mỗi frame còn có arm-wise atomic prompt, ví dụ left arm: grasp red cube, right arm: wait, third arm: place blue cube. Chính lớp prompt trung gian này giúp beginner nhìn được policy đang học vai trò gì, thay vì chỉ thấy một khối dữ liệu image-action khó giải thích.

Ý tưởng paper trong một ví dụ dễ hiểu

Hãy lấy task ba tay xếp ba khối màu. Trong training, robot luôn thấy thứ tự blue -> green -> red. Nếu model chỉ học instruction tổng quát "stack cubes", nó dễ gắn cứng hành vi với arm index và layout: arm 1 gắp blue, arm 2 gắp green, arm 3 gắp red. Khi test đổi thành green -> blue -> red, policy có thể không biết tái phân công, dù từng động tác grasp, lift, place đều đã học.

MA-VLA thay representation bằng chuỗi atomic prompts:

High-level instruction:
  Stack the cubes in order: green, blue, red.

Atomic assignment at time t:
  Arm 1: move to green cube
  Arm 2: wait near blue cube
  Arm 3: hold red cube stable

Atomic assignment at time t+k:
  Arm 1: place green cube on base
  Arm 2: grasp blue cube
  Arm 3: wait

Ở đây, model không chỉ học "làm task", mà học subgoal của từng arm tại từng phase. Khi thứ tự đổi, planner hoặc parser có thể gán lại atomic prompts. Executor vẫn dùng cùng skill primitives đã học: move, grasp, lift, place, handover, wait. Đây là lý do paper nhấn mạnh "atomic action assignment" thay vì chỉ "multi-arm model".

Trong paper, atomic labels cho simulation demonstrations được tạo từ rule-based parser dựa trên task-specific predicates như contact, grasp status và object pose threshold. Nói thực dụng: bạn không cần một LLM perfect để bắt đầu. Với dataset nhỏ, bạn có thể annotate phase bằng script hoặc rule đơn giản:

if gripper_closed and object_height < lift_threshold:
    prompt = "grasp the target object"
elif object_height >= lift_threshold and distance_to_goal > eps:
    prompt = "move the object toward the goal"
elif distance_to_goal <= eps:
    prompt = "place the object at the target"
else:
    prompt = "wait"

Điểm này rất quan trọng cho lab nhỏ: atomic actions không nhất thiết là taxonomy phức tạp. Beginner nên bắt đầu với 4-8 atomic prompts rõ ràng, đủ bao phủ task. Chỉ khi task dài hơn như pass shoe, stack bowls hoặc take photo bằng bốn arm, bạn mới cần phase logic chi tiết hơn.

Kiến trúc MA-VLA

Paper mô tả MA-VLA gồm hai tầng chính: VLM-based Planner và VLA Executor. Planner nhận high-level instruction, tách thành chuỗi atomic subgoals theo thời gian và theo arm. Executor nhận observation, proprioception, atomic prompt của từng arm, rồi sinh action chunk.

Một tuple huấn luyện cho arm i có thể hình dung như sau:

arm_i_input = {
  "global_image": scene_camera,
  "wrist_image": wrist_camera_i,
  "state": joints_or_ee_state_i,
  "atomic_prompt": prompt_i
}

arm_i_target = action_i

Với N arms, model học tổng loss trên tất cả arms. Mục tiêu vẫn là behavioral cloning: dự đoán action của demonstrator từ observation + state + language. Nhưng language không còn là một câu chung cho toàn robot. Language được tách thành atomic prompt theo arm, giúp model biết chính xác arm đó đang chịu trách nhiệm gì.

Repo future-robots triển khai theo hướng tương thích hệ sinh thái OpenPI/Pi0 và LeRobot. README cho biết codebase hỗ trợ unified multi-agent training cho 2-4 robot arms, training mode unified hoặc separate, atomic action learning, agent order shuffling, image masking, Gaussian noise và MoE cho 3+ arms. Với beginner, bạn nên bắt đầu bằng unified model vì chỉ cần một server inference và một checkpoint:

multi-view images + per-arm states + per-arm atomic prompts
        |
        v
shared VLA backbone / Pi0-style action model
        |
        v
action chunk for all arms

Separate model vẫn hữu ích để debug, vì mỗi arm có một policy riêng. Nhưng paper cho thấy separate model không giải quyết được compositional coordination: từng arm có thể follow prompt riêng, nhưng các policy độc lập không có cùng context để phối hợp khi pattern mới xuất hiện.

Arm Shuffle hoạt động như thế nào?

Arm Shuffle là kỹ thuật training-time permutation. Ở mỗi iteration, model có xác suất p_shuffle để tráo thứ tự các bundle của từng arm: state, wrist-view observation, atomic prompt và action target. Nếu chỉ shuffle prompt mà không shuffle state/action tương ứng, dataset sẽ sai. Điểm đúng của MA-VLA là shuffle cả tuple của arm:

Before shuffle:
  Arm 1 tuple: (state_1, wrist_1, prompt_1, action_1)
  Arm 2 tuple: (state_2, wrist_2, prompt_2, action_2)
  Arm 3 tuple: (state_3, wrist_3, prompt_3, action_3)

After shuffle:
  Slot 1 receives tuple from Arm 3
  Slot 2 receives tuple from Arm 1
  Slot 3 receives tuple from Arm 2

Mục tiêu là phá thói quen "slot 1 luôn làm việc A". Model phải đọc atomic prompt và observation thật, thay vì nhớ identity của arm. Đây là một dạng regularization cho permutation invariance. Trong repo, tham số liên quan là shuffle_prob; README đưa ví dụ shuffle_prob=0.5 cho DroidMultiInputs_atom.

MA-VLA còn dùng View Dropout, tức mask một số image input trong training. Nếu quá phụ thuộc vào một wrist camera cụ thể, model dễ hỏng khi occlusion, lighting hoặc background thay đổi. View Dropout buộc model tận dụng global view, wrist views còn lại và proprioception một cách cân bằng hơn.

Với lab thực tế, bạn có thể bắt đầu như sau:

Đừng bật augmentation quá mạnh ngay từ đầu. Nếu dataset chỉ có vài trăm episodes, hãy train baseline sạch trước, kiểm tra action normalization và inference loop, sau đó mới bật Arm Shuffle.

Cài đặt repo và chuẩn bị môi trường

Repo chính thức yêu cầu Python 3.10+, CUDA-compatible GPU và Git LFS. README dùng uv, phù hợp cho project JAX/PyTorch lớn vì lock dependency rõ ràng.

# 1. Clone repo chính thức
git clone https://github.com/zhangzaibin/future-robots.git
cd future-robots

# 2. Cài dependencies, bỏ qua LFS weights nếu chỉ muốn đọc code hoặc chuẩn bị data
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

# 3. Activate environment
source .venv/bin/activate

Nếu bạn đã có stack LeRobot riêng, hãy giữ nguyên môi trường đó và xem repo MA-VLA như một training codebase có converter sang LeRobot-compatible format. Đừng trộn lung tung nhiều bản torch, jax, cuda trong cùng environment. Một cách an toàn là tạo environment riêng cho future-robots, export dataset ra thư mục chung, rồi mount hoặc symlink khi train.

Kiểm tra nhanh GPU:

python - <<'PY'
import jax
print(jax.devices())
PY

Với single GPU 24GB, bạn nên bắt đầu từ task 2-arm hoặc 3-arm nhỏ. Với 3+ arms, action dimension, token length và batch size đều tăng. README gợi ý action_dim=32 cho 3 arms vì 3 x 8 = 24 được pad lên 32, và max_token_len khoảng 200-220 khi có 3 wrist views cộng global view.

Chuẩn bị dataset LeRobot với atomic actions

Pipeline data cơ bản gồm ba bước:

1. Thu hoặc export expert demonstrations dạng H5 có observation, state, action và phase.
2. Convert sang LeRobot-compatible dataset bằng script trong scripts/tasks.
3. Compute normalization statistics trước khi train.

Ví dụ từ README cho task stack cubes:

python scripts/tasks/convert_h5_lerobot_stackcubes.py \
  --h5_path /path/to/stackcubes_with_phases.h5 \
  --output_dir /data/ma-vla/stackcubes_lerobot \
  --task_name stackcubes \
  --use_phase

Sau đó compute statistics:

uv run scripts/compute_norm_stats.py \
  --config-name pi0_base_3arms_stackcubes_mavla

Lưu ý nhỏ nhưng dễ gây lỗi: README nhắc rằng khi compute norm stats, nên đảm bảo shuffle_prob=0.0. Normalization phải phản ánh distribution action/state gốc, không phải distribution đã bị randomize bởi augmentation. Sau khi có stats, bạn mới bật shuffle_prob trong config training.

Dataset của một episode nên có cấu trúc concept như sau:

episode_0001/
  observation.images.global
  observation.images.wrist_0
  observation.images.wrist_1
  observation.images.wrist_2
  observation.state.arm_0
  observation.state.arm_1
  observation.state.arm_2
  action.arm_0
  action.arm_1
  action.arm_2
  language.atomic_prompt.arm_0
  language.atomic_prompt.arm_1
  language.atomic_prompt.arm_2

Tên field thực tế phụ thuộc converter, nhưng nguyên tắc không đổi: mỗi arm cần state, local view, target action và atomic prompt đồng bộ theo frame. Nếu camera timestamp lệch action timestamp, hãy sửa ở data layer trước khi đổ lỗi cho model.

Training unified MA-VLA

Training unified là đường đi nên thử đầu tiên. Một model nhận toàn bộ context và output action cho tất cả arms:

XLA_PYTHON_CLIENT_MEM_FRACTION=1 \
uv run scripts/train.py pi0_base_3arms_stackcubes_mavla \
  --exp-name=mavla_unified_threestackcubes \
  --overwrite

Trong config, phần data transform quan trọng nhất là DroidMultiInputs_atom và DroidMultiOutputs_atom:

DroidMultiInputs_atom(
    action_dim=32,
    num_agents=3,
    model_type=ModelType.PI0,
    shuffle_prob=0.5,
    image_mask_prob=0.3,
    image_noise_prob=0.2,
    mask_multiple_images=True,
    max_masked_images=4,
    noise_std=0.02,
    single_arm_id=None,
)
DroidMultiOutputs_atom(
    num_agents=3,
    single_arm_id=None,
)

Beginner nên log riêng ba nhóm metric:

• Loss tổng và loss theo arm, để phát hiện một arm học kém bất thường.
• Success rate in-domain, ví dụ đúng thứ tự stack đã thấy khi train.
• Success rate OOD, ví dụ đổi thứ tự stack hoặc đổi vai trò handover.

Nếu in-domain thấp, đừng tối ưu OOD vội. Hãy kiểm tra data conversion, action normalization, camera order và gripper action. Nếu in-domain tốt nhưng OOD bằng 0, lúc đó mới tăng chất lượng atomic prompts, bật Arm Shuffle, thêm View Dropout, hoặc tăng diversity của phase assignment.

Inference và deployment

Repo có ví dụ deployment bằng server Python:

python server.py \
  --config pi0_base_3arms_stackcubes_mavla \
  --checkpoint_dir /path/to/unified/checkpoint \
  --port 20019

Luồng inference concept:

robot runtime
  -> capture global camera + wrist cameras
  -> read per-arm joint / end-effector state
  -> build atomic prompts for current phase
  -> call MA-VLA policy server
  -> receive action chunk
  -> execute with safety limits

Trong lab, phần khó nhất thường không phải gọi model, mà là tạo atomic prompts online. Có ba mức triển khai:

Với beginner, hãy dùng state machine. Ví dụ task pass shoe có phase: approach shoe, grasp shoe, handover shoe, receive shoe, place shoe, wait. Khi gripper A đã kẹp và object tới vùng handover, đổi prompt của arm B từ wait sang receive shoe. Cách này thô nhưng minh bạch, đủ để kiểm chứng executor.

Kết quả paper

MA-VLA được đánh giá trên RoboFactory, RoboTwin 2.0 Hard và real-world dual SO101. Theo paper, mỗi simulation task dùng 150 expert demonstrations ở setting in-domain. Atomic action labels được sinh ở frame-level bằng parser dựa trên state predicates.

Trên RoboFactory, MA-VLA đạt trung bình 83.5% cho nhóm two-arm tasks và 83.3% cho nhóm three/four-arm tasks. So với Pi0, lợi thế lớn nhất xuất hiện khi số arm và độ mơ hồ vai trò tăng. Ví dụ trong nhóm three/four-arm tasks, Pi0 đạt 76.5% trung bình, MA-VLA đạt 83.3%.

Trên RoboTwin 2.0 Hard, nơi có distractors, background variation và lighting changes, MA-VLA đạt 49.0% trung bình, cao hơn Pi0 ở mức 41.1% và các baseline ACT, DP, DP3, Pi0-FAST. Đây là tín hiệu rằng atomic prompt không chỉ giúp coordination, mà còn giúp model bám vào subgoal khi visual scene nhiễu.

Quan trọng nhất là OOD compositional generalization. Trong các pattern chưa thấy, DP, Pi0-FAST và Pi0 đều đạt 0.0% trung bình, còn MA-VLA đạt 13.0%. Con số 13% không có nghĩa bài toán đã được giải quyết hoàn toàn. Nó có nghĩa là representation theo atomic assignment + Arm Shuffle tạo ra success khác 0 ở nơi baseline sụp hoàn toàn.

Ablation càng làm rõ vai trò từng thành phần:

Nói cách khác, atomic actions giúp in-domain vì làm rõ trách nhiệm của từng arm. Arm Shuffle mới là cú hích cho OOD vì buộc model tách behavior khỏi arm identity. View Dropout bổ sung robustness thị giác.

Trên real-world dual SO101, Pi0 có in-domain success nhưng OOD bằng 0/20 ở bốn task. MA-VLA vẫn chưa hoàn hảo, nhưng đạt OOD khác 0: Stack Bowls 10/20, Place Cubes 8/20, Pass Toys 2/20, Stack Cubes 2/20. Với robot thật, kết quả này đáng chú ý vì OOD không chỉ đổi label; nó đổi cách phối hợp, pose, timing và contact.

Checklist debug cho beginner

Nếu bạn muốn tái hiện MA-VLA trên dataset riêng, hãy debug theo thứ tự này:

1. Replay data trước khi train. Mở từng episode, visualize global camera, wrist camera, state và action. Nếu replay đã lệch, model sẽ không cứu được.
2. Train without shuffle. Đạt in-domain success trước. Atomic prompts sai sẽ thể hiện ngay ở loss hoặc rollout.
3. Bật Arm Shuffle từ thấp lên cao. Bắt đầu 0.3, sau đó 0.5. Nếu in-domain rơi mạnh, prompt có thể chưa đủ rõ hoặc camera slot bị xử lý sai.
4. Tách evaluation ID và OOD. Không dùng cùng order, cùng layout, cùng vai trò để tự đánh lừa mình.
5. Log prompt thực tế ở inference. Nhiều lỗi đến từ planner/state machine gán prompt sai phase, không phải executor yếu.
6. Giữ safety layer độc lập. VLA output phải đi qua joint limits, velocity limits, workspace limits và collision checks.

Khi nào nên dùng MA-VLA?

Dùng MA-VLA khi task của bạn thật sự có nhiều effector cần phối hợp: dual-arm handover, hai tay xếp vật, ba robot cùng lắp ráp, bốn arm giữ/chụp/đặt object trong dây chuyền nhỏ. Nếu task chỉ là single-arm pick-and-place, MA-VLA có thể quá nặng; ACoT-VLA với LeRobot data hoặc các policy LeRobot đơn giản sẽ dễ bắt đầu hơn.

Ngược lại, nếu bạn đang làm whole-body hoặc humanoid manipulation, MA-VLA là một ý tưởng nên học dù không dùng nguyên repo. Atomic assignment có thể mở rộng từ "arm" sang "end-effector": left hand, right hand, base, torso, head camera. Đây là cầu nối tự nhiên với các hướng whole-body VLA, nơi policy cần biết bộ phận nào làm subtask nào.

Kết luận

MA-VLA đưa ra một bài học rất thực tế: muốn robot nhiều cánh tay generalize, đừng chỉ tăng model size và trộn thêm data. Hãy biểu diễn nhiệm vụ theo cách có cấu trúc hơn. Atomic actions làm rõ subgoal theo arm. Arm Shuffle phá bias vào arm identity. View Dropout giúp model chịu được camera noise và occlusion. Kết hợp lại, hệ thống đạt kết quả khác 0 trên các collaboration patterns mà Pi0 và diffusion baselines thất bại hoàn toàn.

Với người mới, con đường hợp lý là: tái hiện một task LeRobot nhỏ, annotate 4-8 atomic prompts, train baseline sạch, sau đó bật Arm Shuffle và tạo split OOD thật sự. Khi bạn đo được ID/OOD riêng, MA-VLA không còn là một paper khó đọc, mà trở thành một recipe khá rõ cho multi-arm robot learning.

Bài viết liên quan

• Hands-on: Fine-tune OpenVLA với LeRobot
• ACoT-VLA với LeRobot data
• Chạy FineVLA cho robot dual-arm