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.
| 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 itsparams/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 andlabels.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_sizeandtraining.fsdp_devicesfor 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.mdfor 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 ownjax.sharding.Shardingto 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 themodel_fieldsfrommodel_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}
}