🌱 Future

Future is a research ecosystem working toward one goal: agentic, open-ended VLA policies — agents that decompose a goal into atomic skills, compose them on the fly, and coordinate multiple effectors to solve tasks they were never explicitly trained on.

Its first work, MA-VLA (Multi-Arm Vision-Language-Action Model for Collaboration and Compositional Generalization), introduces multi-arm compositional generalization and is accepted to ECCV 2026. More works are on the way. 🚀

📋 Table of Contents

• Features
• Installation
• Quick Start
• Data Processing
• Training Modes
• Deployment
• Data Transform Configuration
• Architecture Matrix
• Advanced Usage
• Roadmap
• Citation

✨ Features

• 🔄 Unified Multi-Agent Support: Single codebase for 2-4 robot arms
• 🎯 Flexible Training Modes:
  • Unified model (one model controls all arms)
  • Separate model (dedicated model per arm)
• 🎨 Rich Data Augmentation: Image masking, Gaussian noise, agent order shuffling
• 🧩 Mixture of Experts (MoE): Advanced architecture for 3+ arms
• 📊 Atomic Action Learning: Fine-grained action decomposition for compositional generalization
• 🧪 MACG Benchmark: Multi-Arm Compositional Generalization data-generation pipeline
• ⚡ Scalable Design: Easy to extend to 5+ arms

🚀 Installation

Prerequisites

• Python 3.10+
• CUDA-compatible GPU
• Git LFS

Setup

# Clone the repository
git clone https://github.com/yourusername/multivla-pi0.git
cd multivla-pi0

# Install dependencies
GIT_LFS_SKIP_SMUDGE=1 uv sync
GIT_LFS_SKIP_SMUDGE=1 uv pip install -e .

# Activate environment
source .venv/bin/activate

🎯 Quick Start

1. Prepare Your Dataset

# Generate atomic action dataset
python scripts/tasks/convert_h5_lerobot_stackcubes.py \
    --h5_path path/to/h5file/with/phase \
    --output_dir path/to/output \
    --task_name your_task_name \
    --use_phase

# Compute normalization statistics
uv run scripts/compute_norm_stats.py --config-name your_config_name

2. Train Your Model

Unified Training (one model 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

Separate Training (one model per arm):

XLA_PYTHON_CLIENT_MEM_FRACTION=1 uv run scripts/train.py pi0_base_3arms_stackcubes_agent0 \
    --exp-name=mavla_separate_agent0 \
    --overwrite

3. Deploy the Model

python server_agent0.py \
    --config pi0_base_3arms_stackcubes_atom_shuffle \
    --checkpoint_dir /path/to/checkpoint/directory \
    --port 20019

📊 Data Processing

Phase Assignment

Assign task phases to trajectory segments for atomic action learning.

# Implementation coming soon
# This will enable automatic phase labeling for complex multi-step tasks

Language Instruction Assignment

Step 1: Generate Atomic Action Dataset

Convert H5 format data to LeRobot-compatible format with phase annotations:

python scripts/tasks/convert_h5_lerobot_stackcubes.py \
    --h5_path /path/to/your/h5_file_with_phases.h5 \
    --output_dir /path/to/empty/output/directory \
    --task_name stackcubes \
    --use_phase

Parameters:

• --h5_path: Path to H5 file containing trajectory data with phase labels
• --output_dir: Empty directory for output (will be created if not exists)
• --task_name: Task identifier (e.g., stackcubes, takephoto, passshoe)
• --use_phase: Enable phase-based atomic action generation

Step 2: Compute Normalization Statistics

Calculate mean and standard deviation for action normalization:

uv run scripts/compute_norm_stats.py --config-name <config_name>

Note: Ensure shuffle_prob=0.0 in your config when computing statistics to avoid data randomization.

MACG Data Generation

Data Generation for Multi-arm Compositional Generalization Benchmark

# set phase for stack macg task
python scripts/set_atom/macg/set_phases_to_h5_threestackcube.py

# generate video
python scripts/set_atom/macg/visualize_phases_threestackcube.py


# generate dataset
python scripts/set_atom/macg/placefactory/convert_h5_lerobot_placefactory.py --h5_path dataset/placefactory_order_123_20251216_with_phases.h5 --output_dir dataset/placefactory --task_name placefactory --use_phase

# compute the norm states
uv run scripts/compute_norm_stats.py --config-name pi0_base_3arms_placefactory_mavla_macg


# start run :)
XLA_PYTHON_CLIENT_MEM_FRACTION=1 uv run scripts/train.py pi0_base_3arms_placefactory_mavla_macg \
    --exp-name=macg_placefactory \
    --overwrite

🎓 Training Modes

Unified Arm Training

Description: Single model learns to control all robot arms simultaneously. The model receives observations from all arms and outputs actions for all arms in one forward pass.

Advantages:

• ✅ Shared representation learning across arms
• ✅ Better generalization to new configurations
• ✅ Single model deployment
• ✅ Implicit coordination learning

Example: 3-arm stack cubes task

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

Configuration:

data_transforms=lambda model_config: _transforms.Group(
    inputs=[droid_policy.DroidMultiInputs_atom(
        action_dim=model_config.action_dim, 
        num_agents=3,
        model_type=ModelType.PI0,
        shuffle_prob=0.5,
        image_mask_prob=0.3
    )],
    outputs=[droid_policy.DroidMultiOutputs_atom(num_agents=3)]
)

Separate Arm Training

Description: Train individual models for each robot arm. Each model specializes in controlling a single arm using the single_arm_id parameter to filter the data.

Advantages:

• ✅ Specialized control per arm
• ✅ Independent optimization
• ✅ Modular deployment (hot-swap arms)
• ✅ Easier debugging and analysis

Example: Train models for a 3-arm system

Agent 0:

XLA_PYTHON_CLIENT_MEM_FRACTION=1 uv run scripts/train.py pi0_base_3arms_stackcubes_agent0 \
    --exp-name=mavla_separate_agent0 \
    --overwrite

Agent 1:

XLA_PYTHON_CLIENT_MEM_FRACTION=1 uv run scripts/train.py pi0_base_3arms_stackcubes_agent1 \
    --exp-name=mavla_separate_agent1 \
    --overwrite

Agent 2:

XLA_PYTHON_CLIENT_MEM_FRACTION=1 uv run scripts/train.py pi0_base_3arms_stackcubes_agent2 \
    --exp-name=mavla_separate_agent2 \
    --overwrite

Configuration:

data_transforms=lambda model_config: _transforms.Group(
    inputs=[droid_policy.DroidMultiInputs_atom(
        action_dim=model_config.action_dim, 
        num_agents=3,
        model_type=ModelType.PI0,
        shuffle_prob=0.0,        # Disable shuffling for single-arm training
        image_mask_prob=0.0,     # Disable masking for single-arm training
        single_arm_id=0          # Target arm ID (0, 1, 2, ...)
    )],
    outputs=[droid_policy.DroidMultiOutputs_atom(
        num_agents=3, 
        single_arm_id=0          # Must match input single_arm_id
    )]
)

Key Parameter: single_arm_id

• Purpose: Filters multi-agent data to extract single-arm trajectories for specialized training
• Type: int or None
• Values:
  • 0, 1, 2, ...: Train on specific arm's data only
  • None: Train on all arms (unified mode)
• Usage: Set the same single_arm_id in both DroidMultiInputs_atom and DroidMultiOutputs_atom
• Data Augmentation: When using single_arm_id, set shuffle_prob=0.0 and image_mask_prob=0.0 to ensure consistent single-arm behavior

🌐 Deployment

Unified Model Deployment

Note: the following code now support Robofactory but not support Robotwin. We will update Robotwin later.
Deploy a single model serving all robot arms:

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

Separate Model Deployment

Deploy individual models for each arm with memory-efficient settings:

# Agent 0 (Port 20019)
XLA_PYTHON_CLIENT_PREALLOCATE=false \
XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 \
python server_agent0.py \
    --config pi0_base_3arms_stackcubes_agent0 \
    --checkpoint_dir /path/to/agent0/checkpoint \
    --port 20019

# Agent 1 (Port 20020)
XLA_PYTHON_CLIENT_PREALLOCATE=false \
XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 \
python server_agent1.py \
    --config pi0_base_3arms_stackcubes_agent1 \
    --checkpoint_dir /path/to/agent1/checkpoint \
    --port 20020

# Agent 2 (Port 20021)
XLA_PYTHON_CLIENT_PREALLOCATE=false \
XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 \
python server_agent2.py \
    --config pi0_base_3arms_stackcubes_agent2 \
    --checkpoint_dir /path/to/agent2/checkpoint \
    --port 20021

Memory Management:

• XLA_PYTHON_CLIENT_PREALLOCATE=false: Disable pre-allocation for flexible memory usage
• XLA_PYTHON_CLIENT_MEM_FRACTION=0.4: Limit each model to 40% GPU memory (adjust based on your GPU)

🔧 Data Transform Configuration

Unified Multi-Agent Transform Classes

The codebase uses unified data transform classes that support 2-4 robot arms with a single num_agents parameter, replacing 6+ redundant legacy classes.

Training with Data Augmentation (atom version)

For multi-arm unified training:

DroidMultiInputs_atom(
    action_dim=32,              # Model action dimension (padded)
    num_agents=3,               # Number of robot arms (2-4)
    model_type=ModelType.PI0,   # Model architecture
    shuffle_prob=0.5,           # Agent order shuffling probability
    image_mask_prob=0.3,        # Image masking probability
    image_noise_prob=0.2,       # Gaussian noise probability
    mask_multiple_images=True,  # Allow multiple image masking
    max_masked_images=4,        # Max images to mask simultaneously
    noise_std=0.02,             # Gaussian noise standard deviation
    single_arm_id=None          # None for unified training
)
DroidMultiOutputs_atom(
    num_agents=3,
    single_arm_id=None
)

For single-arm separate training:

DroidMultiInputs_atom(
    action_dim=32,
    num_agents=3,
    model_type=ModelType.PI0,
    shuffle_prob=0.0,           # Disable for single-arm
    image_mask_prob=0.0,        # Disable for single-arm
    image_noise_prob=0.0,       # Disable for single-arm
    single_arm_id=0             # Target arm (0, 1, 2, ...)
)
DroidMultiOutputs_atom(
    num_agents=3,
    single_arm_id=0             # Must match input
)

Inference (non-atom version)

For deployment without atomic action learning:

DroidMultiInputs(
    action_dim=32,
    num_agents=3,
    model_type=ModelType.PI0,
    task_instruction="Stack the cubes"  # Fixed task instruction
)
DroidMultiOutputs(num_agents=3)

📖 Data Augmentation Parameters

Smart Augmentation Logic:

• When all agents have identical instructions, shuffling and masking are automatically disabled to maintain consistency
• When single_arm_id is set, augmentation probabilities should typically be set to 0.0

🏆 Benefits of Unified Classes

⚠️ Legacy Classes (Deprecated)

The following classes remain for backward compatibility but are deprecated:

Migration: All training configs in config.py have been updated to use unified classes.

📐 Architecture Matrix

Legend:

• Baseline: Standard multi-agent VLA
• + Atom: Atomic action learning with phase decomposition
• + MoE: Mixture of Experts architecture (available for 3+ arms)

🔬 Advanced Usage

Creating a Custom Config

Define a new training configuration in config.py:

TrainConfig(
    name="my_3arms_custom_task",

    # Model configuration
    model=pi0.Pi0Config(
        action_dim=32,         # 3 agents × 8 dims = 24, padded to 32
        action_horizon=50,     # Action prediction horizon
        max_token_len=220,     # Maximum token length
    ),

    # Data configuration
    data=SimpleDataConfig(
        repo_id="path/to/your/dataset",
        assets=AssetsConfig(
            assets_dir="path/to/assets",
            asset_id="your_asset_id"
        ),

        # Data transforms with augmentation
        data_transforms=lambda model_config: _transforms.Group(
            inputs=[droid_policy.DroidMultiInputs_atom(
                action_dim=model_config.action_dim, 
                num_agents=3,
                model_type=ModelType.PI0,
                shuffle_prob=0.5,        # 50% shuffle
                image_mask_prob=0.3,     # 30% masking
                image_noise_prob=0.2     # 20% noise
            )],
            outputs=[droid_policy.DroidMultiOutputs_atom(num_agents=3)],
        ),

        base_config=DataConfig(prompt_from_task=True),
    ),

    # Weight initialization
    weight_loader=weight_loaders.CheckpointWeightLoader(
        "path/to/pretrained/params"
    ),

    # Training hyperparameters
    lr_schedule=_optimizer.CosineDecaySchedule(
        warmup_steps=300,
        peak_lr=3e-5,
        decay_steps=10000,
        decay_lr=1e-6,
    ),
    num_train_steps=15000,
    batch_size=32,
    log_interval=50,
    save_interval=1000,
)

Key Configuration Tips

1. Action Dimension: Set to nearest power of 2 above num_agents × action_dim_per_arm

   • 2 arms × 8 dims = 16 → action_dim=16 or 32
   • 3 arms × 8 dims = 24 → action_dim=32
   • 4 arms × 8 dims = 32 → action_dim=32

2. Token Length: Increase for more images

   • 2 arms + 1 global: max_token_len=150-180
   • 3 arms + 1 global: max_token_len=200-220
   • 4 arms + 1 global: max_token_len=250-280

3. Data Augmentation:

   • High diversity tasks: increase shuffle_prob and image_mask_prob
   • Low diversity tasks: use conservative augmentation
   • Single-arm training: set all augmentation to 0.0

4. Training Steps:

   • Small datasets (<1000 episodes): 5000-10000 steps
   • Medium datasets (1000-5000 episodes): 10000-20000 steps
   • Large datasets (>5000 episodes): 20000-50000 steps

📝 Development Notes

Configuration Management

All modifications correspond to specific config files. You can trace model and data processing files through the config system:

# Find config by name
grep -n "name=\"your_config_name\"" src/openpi/training/config.py

Version Control

Use git commit history to track changes:

# View file history
git log --follow -- src/openpi/policies/droid_policy.py

# Compare versions
git diff <commit1> <commit2> -- src/openpi/training/config.py

Model Variants

🗺️ Roadmap

This codebase evolves alongside our research program toward open-ended agentic embodied policies.

• [x] MA-VLA (ECCV 2026) — Unified multi-arm VLA, atomic action learning, MACG benchmark
• [x] Unified 2–4 arm training & deployment (unified / separate modes)
• [x] Mixture of Experts (MoE) for 3+ arms
• [ ] Single-arm baselines and broader task coverage
• [ ] Agentic skill composition — decompose goals into atomic skills and compose them at inference
• [ ] Open-ended task generalization — solve unseen task combinations without task-specific training
• [ ] Long-horizon planning and closed-loop self-correction

🚧 Each step builds on the last: from collaboration → compositional generalization → open-ended agentic policies.

📚 Citation

If you find MA-VLA or this codebase useful in your research, please cite:

comming soon

🤝 Contributing

We welcome contributions! Please:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Commit your changes (git commit -m 'Add amazing feature')
4. Push to the branch (git push origin feature/amazing-feature)
5. Open a Pull Request

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🙏 Acknowledgments

We would like to thank the following projects for their inspiration and contributions to the robotics community:

• OpenPI - Physical Intelligence foundation

📧 Contact

For questions and discussions:

• Open an Issue
• Start a Discussion