MA-VLA: Multi-Arm VLA on LeRobot

What Problem Does MA-VLA Solve?

MA-VLA, short for Multi-Arm Vision-Language-Action Model, is a new framework for collaborative robot manipulation introduced in the paper MA-VLA: Multi-Arm Vision-Language-Action Model for Collaboration and Compositional Generalization and the official repository zhangzaibin/future-robots. Its core idea is simple but important: instead of giving a multi-arm robot one global instruction such as "stack the cubes", MA-VLA decomposes the task into atomic actions and assigns those atomic prompts to individual arms.

For a single-arm robot, a high-level instruction may be enough. "Pick up the red cube" usually implies one actor, one gripper, one target, and one short manipulation sequence. Multi-arm manipulation is different. The same sentence hides several coordination decisions:

• Which arm should approach the object?
• Which arm should wait?
• Which arm should hand over or receive?
• What is the required order of collaboration?
• If the object layout or arm roles change, can the model still recombine known behaviors?

End-to-end VLAs such as Pi0 can perform well when training and test tasks share the same coordination pattern. The hard part is test-time collaboration that was absent from the training set. The paper calls this ability multi-arm compositional generalization: the robot must reuse known atomic behaviors in a new arm-role composition.

If you have already read Hands-on: Fine-tune OpenVLA with LeRobot, think of MA-VLA as the multi-arm extension of that workflow. The data still needs synchronized images, robot states, and actions in a LeRobot-compatible format. The difference is that each frame also carries arm-wise atomic prompts, for example left arm: grasp the red cube, right arm: wait, and third arm: place the blue cube. This intermediate language layer makes the policy easier to inspect and much less dependent on a single global instruction.

The Paper Idea in a Concrete Example

Consider a three-arm cube-stacking task. During training, the robot always sees the order blue -> green -> red. A model trained only with the instruction "stack the cubes" can overfit to arm identity and layout: arm 1 picks blue, arm 2 picks green, arm 3 picks red. At test time, if the required order becomes green -> blue -> red, the model may fail even though grasping, lifting, and placing are all familiar skills.

MA-VLA changes the representation:

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

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

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

The model no longer learns only "how to do the whole task". It learns which subgoal each arm should follow at each phase. When the order changes, a planner or parser can assign different atomic prompts while the executor reuses the same primitive behaviors: move, grasp, lift, place, hand over, receive, and wait.

In the paper, frame-level atomic action labels for simulation demonstrations are generated with a rule-based parser using task-specific predicates such as contact, grasp status, and object-pose thresholds. That detail matters for small labs. You do not need a perfect LLM planner to start. For a first dataset, atomic prompts can be generated from transparent rules:

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"

Start with four to eight atomic prompts. Add more only when the task truly needs more phases, such as pass-shoe, stack-bowls, object handover, or four-arm take-photo tasks.

MA-VLA Architecture

The paper describes MA-VLA as a two-level system: a VLM-based Planner and a VLA Executor. The planner decomposes a high-level instruction into temporally ordered atomic subgoals for each arm. The executor grounds those arm-wise linguistic subgoals into robot actions.

Conceptually, one training tuple for arm i looks like this:

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

For N arms, the training loss sums action prediction errors across all arms. It is still behavioral cloning: predict the demonstrator action from observations, proprioception, and language. But language is no longer a single sentence shared by the whole robot. It is a per-arm atomic prompt, which tells each action head what that arm is supposed to do.

The official future-robots repository follows the OpenPI/Pi0 and LeRobot ecosystem. Its README lists support for unified multi-agent training across 2-4 robot arms, unified and separate training modes, atomic action learning, agent order shuffling, image masking, Gaussian noise, and Mixture-of-Experts variants for 3+ arms. For a beginner, unified training is the best first path because it produces one checkpoint and one inference server:

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 per-arm models can still be useful for debugging. However, the paper shows that separate models do not solve compositional coordination by themselves. Each arm may follow its own prompt, but independently deployed policies do not share enough context to handle unseen collaboration patterns.

How Arm Shuffle Works

Arm Shuffle is a training-time permutation strategy. At each iteration, with probability p_shuffle, the system shuffles the per-arm bundles. A bundle includes the arm state, wrist-view observation, atomic prompt, and target action. Shuffling only prompts would corrupt the dataset. MA-VLA shuffles complete arm tuples:

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

The goal is to stop the model from learning shortcuts such as "slot 1 always performs behavior A". The model must attend to the atomic prompt and observation rather than memorizing arm identity. This is a practical form of regularization toward permutation-invariant coordination. In the repository, the relevant parameter is shuffle_prob; the README example uses shuffle_prob=0.5 inside DroidMultiInputs_atom.

MA-VLA also uses View Dropout, which masks some image inputs during training. If a policy depends too heavily on one wrist camera, it can fail under occlusion, lighting changes, or clutter. View Dropout encourages the model to use global view, remaining wrist views, and proprioception more robustly.

A reasonable starting point:

Do not start with aggressive augmentation. First train a clean baseline, verify normalization and inference, then introduce Arm Shuffle.

Installation

The official repository requires Python 3.10+, a CUDA-compatible GPU, and Git LFS. It uses uv for dependency management.

# 1. Clone the official repository
git clone https://github.com/zhangzaibin/future-robots.git
cd future-robots

# 2. Install dependencies while skipping large LFS files if needed
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

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

If you already have a LeRobot environment, keep it separate. Treat future-robots as the MA-VLA training codebase and use its converters to export LeRobot-compatible data. Mixing several versions of torch, jax, CUDA, OpenPI, and LeRobot in one environment is a common source of hard-to-debug failures.

Check that the accelerator is visible:

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

On a 24GB GPU, start with a small 2-arm or 3-arm task. More arms increase the number of image tokens, proprioceptive streams, action dimensions, and memory pressure. The README suggests action_dim=32 for 3 arms because 3 x 8 = 24 is padded to 32, and max_token_len around 200-220 when using one global camera plus three wrist cameras.

Preparing LeRobot Data with Atomic Actions

The practical data pipeline has three steps:

1. Collect or export expert demonstrations as H5 files containing observations, states, actions, and phases.
2. Convert them into a LeRobot-compatible dataset with scripts under scripts/tasks.
3. Compute normalization statistics before training.

The README example for stack cubes is:

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

Then compute normalization statistics:

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

One easy mistake: keep shuffle_prob=0.0 when computing normalization statistics. Normalization should represent the original state/action distribution, not the augmented distribution after random arm permutations.

At a conceptual level, each episode should contain:

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

The exact field names depend on the converter, but the rule is fixed: each arm needs synchronized state, local observation, target action, and atomic prompt. If camera timestamps are misaligned with actions, fix the data layer first.

Training a Unified MA-VLA Policy

Unified training should be the first serious run. One model receives the full multi-arm context and outputs actions for all arms:

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

The key transform is DroidMultiInputs_atom paired with 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,
)

Track three metric groups:

• Total loss and per-arm loss, so one weak arm does not hide inside the average.
• In-domain success rate, using collaboration patterns seen in training.
• OOD success rate, using unseen role orders, object orders, or handover patterns.

If in-domain performance is poor, do not tune OOD generalization yet. Check data conversion, action normalization, camera order, gripper command scaling, and atomic labels. If in-domain is solid but OOD remains zero, improve atomic prompt quality, introduce Arm Shuffle, add View Dropout, or make the training set more diverse.

Inference and Deployment

The repository includes a Python policy server for unified deployment:

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

The inference loop looks like this:

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

For most labs, online atomic prompt generation is harder than calling the model. There are three common levels:

Start with a state machine. For a pass-shoe task, phases can be approach shoe, grasp shoe, handover shoe, receive shoe, place shoe, and wait. When arm A reaches the handover region, switch arm B from wait to receive shoe. This is crude, but it makes the executor testable.

Reported Results

MA-VLA is evaluated on RoboFactory, RoboTwin 2.0 Hard, and a real-world dual SO101 setup. The paper reports 150 expert demonstrations per simulation task in the in-domain setting. Frame-level atomic action labels are produced by a parser using state predicates.

On RoboFactory, MA-VLA reaches 83.5% average success on two-arm tasks and 83.3% on three/four-arm tasks. The gains become more visible as the number of arms and role ambiguity increase. For three/four-arm tasks, Pi0 reaches 76.5% average success while MA-VLA reaches 83.3%.

On RoboTwin 2.0 Hard, which includes distractors, background variation, and lighting changes, MA-VLA reaches 49.0% average success. Pi0 reaches 41.1%, while ACT, DP, DP3, and Pi0-FAST are lower. This suggests that atomic prompts help the policy stay grounded in the relevant subgoal even under visual disturbance.

The most important result is OOD compositional generalization. On unseen collaboration patterns, DP, Pi0-FAST, and Pi0 all report 0.0% average success. MA-VLA reaches 13.0%. That is not a solved problem, but it is a meaningful jump from total collapse to non-zero transfer.

The ablation tells the same story:

Atomic actions improve in-domain learning by clarifying each arm's responsibility. Arm Shuffle is the key step for OOD transfer because it discourages fixed arm identity shortcuts. View Dropout adds robustness to visual changes.

On real-world dual SO101 tasks, Pi0 achieves some in-domain success but records 0/20 OOD success on four tasks. MA-VLA is still imperfect, but it reaches non-zero OOD success: Stack Bowls 10/20, Place Cubes 8/20, Pass Toys 2/20, and Stack Cubes 2/20. For real robots, that matters because OOD changes affect role assignment, pose, timing, and contact dynamics, not just text labels.

Beginner Debugging Checklist

If you want to reproduce MA-VLA on your own dataset, debug in this order:

1. Replay the data before training. Visualize global camera, wrist cameras, states, actions, and prompts for every phase.
2. Train without shuffle first. Get in-domain success before optimizing OOD behavior.
3. Increase Arm Shuffle gradually. Start at 0.3, then try 0.5. If in-domain success drops sharply, prompts or camera slots may be wrong.
4. Separate ID and OOD evaluation. Use unseen object orders, role orders, or handover patterns for OOD.
5. Log inference prompts. Many failures come from the planner assigning the wrong atomic phase, not from the executor.
6. Keep a separate safety layer. VLA actions should always pass through joint, velocity, workspace, and collision limits.

When Should You Use MA-VLA?

Use MA-VLA when the task really requires multiple effectors: dual-arm handover, two-arm stacking, three-arm assembly, or four-arm workcells where arms hold, place, and inspect objects together. If your task is simple single-arm pick-and-place, MA-VLA is probably too heavy; ACoT-VLA with LeRobot data or a standard LeRobot policy is easier to start with.

For whole-body or humanoid manipulation, the MA-VLA idea is useful even if you do not use the repository directly. Atomic assignment can extend from "arm" to "end-effector": left hand, right hand, mobile base, torso, and head camera. That makes it a natural bridge to whole-body VLA, where the policy must decide which body part owns each subtask.

Conclusion

MA-VLA teaches a practical lesson: for multi-arm generalization, representation matters as much as model scale. Atomic actions make each arm's subgoal explicit. Arm Shuffle breaks fixed identity shortcuts. View Dropout improves robustness to missing or noisy views. Together, these choices produce non-zero success on collaboration patterns where Pi0 and diffusion baselines fail completely.

For a beginner, the best path is to reproduce a small LeRobot-style task, annotate four to eight atomic prompts, train a clean baseline, enable Arm Shuffle, and evaluate in-domain and OOD splits separately. Once that loop works, MA-VLA becomes a clear recipe for multi-arm robot learning rather than just another large VLA paper.

Related Posts

• Hands-on: Fine-tune OpenVLA with LeRobot
• ACoT-VLA with LeRobot data
• Run FineVLA for dual-arm robots