Access A-0 on Hugging Face

This repository is publicly accessible, but you have to accept the conditions to access its files and content.

The A-0 weights are subject to the Gemma Terms of Use and the Gemma Prohibited Use Policy. To access A-0, log in to Hugging Face, review the terms and click below. Requests are processed immediately.

Log in or Sign Up to review the conditions and access this model content.

A-0

A-0 model weights: the pre-trained model, plus three post-trained checkpoints as examples of the post-training recipe. Weights only: no training state, data or inference server.

Project page arXiv Code License

Folder What it is Task instruction
pretrained/ A-0 pre-trained weights; the starting point for post-training โ€“
examples/orange_pyramid/ Post-trained on orange pyramid "Stack four oranges into a pyramid."
examples/towel_folding/ Post-trained on towel folding "Fold the towel."
examples/letter_blocks/ Post-trained on letter blocks "Spell YUBI by placing the U, B and I blocks to the right of the Y block already in place."

Post-trained folders contain EMA weights after 20k post-training steps.

Each folder contains:

File Content
params/ Weights in openpi Orbax/OCDBT format (about 12 GB, 3.36 B parameters)
model_config.json Model settings: model_fields are the Pi0Config arguments in the A-0 code. The post-trained folders also list the task instruction, action horizon, action label rate and number of stages.
norm_stats.json State/action normalization statistics (post-trained folders only)
LICENSE, LICENSE_GEMMA.txt Apache-2.0 and Gemma terms of use
SHA256SUMS Checksums; verify with sha256sum -c SHA256SUMS inside a folder

Download

# pre-trained weights only (about 12 GB)
hf download airoa-org/A-0 --include "pretrained/*" --local-dir A-0

# one post-trained example
hf download airoa-org/A-0 --include "examples/orange_pyramid/*" --local-dir A-0

# everything (about 50 GB)
hf download airoa-org/A-0 --local-dir A-0

cd A-0/pretrained && sha256sum -c SHA256SUMS

Fine-tune on your own task (from pretrained/)

Use the A-0 post-training code, github.com/airoa-org/AIRoA-0. Its default configuration expects the pretrained weights at checkpoints/airoa-0_pretrained, so download pretrained/ there:

git clone https://github.com/airoa-org/AIRoA-0.git && cd AIRoA-0
python3.11 -m venv .venv && . .venv/bin/activate
python -m pip install -e '.[train,test]'

hf download airoa-org/A-0 --include "pretrained/*" --local-dir checkpoints/A-0
mv checkpoints/A-0/pretrained checkpoints/airoa-0_pretrained
cp configs/post_training.toml configs/local.toml

Edit configs/local.toml:

  • Initial weights: init_params = "../checkpoints/airoa-0_pretrained" is the default; either this folder or its params/ works.
  • Dataset: data.repo_id, the state and action keys and dimensions, and the camera mapping. Your data must be in LeRobot v2.1 format; LeRobot v3.0 is not supported yet. The post-trained checkpoints here use an 11-dimensional state and a 20-dimensional action; the repository README shows the layout.
  • Labels: labels.task_instruction, the stage texts and labels.pa_to_stage. These map each subtask label in your data to a stage; 2, 3, 4 or 8 stages are supported.
  • Hardware: training.batch_size and training.fsdp_devices for your GPUs.

Then run:

./run.sh inspect --config configs/local.toml                # check the dataset and one example
./run.sh norm    --config configs/local.toml                # compute normalization statistics
./run.sh train   --config configs/local.toml --run trial01  # post-train
  • What is loaded: the pre-trained backbone, the action expert and the visual-budget gate. The task's stage head and the proprioceptive memory are created fresh.
  • Default optimization: batch 64, 20k steps, learning rate 5e-6 after 500 warm-up steps, EMA 0.99. The post-trained checkpoints here used the same optimization settings.
  • Requirements: Linux, an NVIDIA GPU with CUDA 12 and Python 3.11. The default uses 8 GPUs and needs at least 32 CPU cores for data loading.
  • Data conventions: see the repository README and docs/data.md for units, coordinate frames and the action convention.
  • Plain openpi: it can load the shared core weights, but not the visual-budget gate.

Load

With openpi, or the A-0 code (github.com/airoa-org/AIRoA-0), after downloading:

import jax.numpy as jnp
from openpi.models import model as _model

# jax.Array on the visible accelerators (GPU/TPU), replicated on each device
params = _model.restore_params("A-0/examples/orange_pyramid/params", dtype=jnp.bfloat16)
print(sorted(params))
# ['PaliGemma', 'action_in_proj', 'action_out_proj', 'proprio_memory_compressor', ..., 'stage_classifier', ...]

restore_params arguments:

  • sharding=: pass your own jax.sharding.Sharding to shard the weights instead of replicating them, for example for FSDP.
  • restore_type=np.ndarray: loads host-memory numpy arrays instead, for inspection without an accelerator.

With Orbax only:

import jax, jax.numpy as jnp, orbax.checkpoint as ocp

path = "/abs/path/to/A-0/examples/orange_pyramid/params"
mesh = jax.sharding.Mesh(jax.devices(), ("x",))
sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())  # replicated
with ocp.PyTreeCheckpointer() as ckptr:
    item = {"params": ckptr.metadata(path)["params"]}
    args = jax.tree.map(lambda _: ocp.ArrayRestoreArgs(sharding=sharding, dtype=jnp.bfloat16), item)
    params = ckptr.restore(path, ocp.args.PyTreeRestore(item=item, restore_args=args))["params"]
  • Memory: 3.36 B parameters, about 6.7 GB per device in bfloat16 when replicated, or 13 GB in float32 (leave out dtype).
  • Building the model: in the A-0 code, Pi0Config(**model_fields) with the model_fields from model_config.json (JSON lists converted to tuples) builds the matching architecture. Fields not listed keep their defaults.
  • Extra modules: the post-trained checkpoints add a stage head (stage_classifier) and a proprioceptive memory (proprio_memory_*). They are implemented in the A-0 code, not in upstream openpi.

License

The A-0 code is released under the Apache License 2.0 (LICENSE). The vision-language backbone is initialized from PaliGemma, so the weights are subject to the Gemma Terms of Use (LICENSE_GEMMA.txt) and the Gemma Prohibited Use Policy.

Gemma is provided under and subject to the Gemma Terms of Use found at ai.google.dev/gemma/terms (NOTICE).

Contact

Questions about A-0: ext-oss@airoa.org

Citation

@techreport{airoa0,
  title         = {{AIRoA-0}: Long-Horizon, Precise, Cross-Embodiment
                   Manipulation from Handheld Data},
  author        = {Yueh-Hua Wu and Kei Ota and Hanlong Li and
                   Hayata Yuasa and Jason Wang and Kengo Ikeuchi and
                   Kensuke Chiyoma and Makoto Sugiura and
                   Masafumi Takahashi and Masatoshi Tateno and
                   Naoaki Kanazawa and Nicola Scianca and
                   Petr Khrapchenkov and Satoshi Otsubo and
                   Takehiko Ohkawa and Takuya Okubo and
                   Tatsuya Matsushima and Yuki Noguchi and Yuma Shin and
                   Zixuan Zhong},
  institution   = {AIRoA},
  year          = {2026},
  month         = nov,
  eprint        = {2611.03815},
  archivePrefix = {arXiv}
}
Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading

Paper for airoa-org/A-0