SereneC/AHA-WAM-Pretrained

AHA-WAM and Fast-WAM RoboCOIN Pretrained Checkpoints

This repository releases RoboCOIN-pretrained checkpoints for AHA-WAM and Fast-WAM. These checkpoints are intended as initialization weights for real-robot fine-tuning and deployment experiments with the public AHA-WAM codebase.

We pretrained both models for 45k steps on a selected subset of RoboCOIN. The goal of this release is to make real-world adaptation easier and to support comparisons requested by the community.

Repository Files

AHA-WAM-Pretrained/
├── README.md
├── AHA-WAM-pretrained.pt
└── Fast-WAM-pretrained.pt

• AHA-WAM-pretrained.pt: AHA-WAM checkpoint pretrained on a RoboCOIN subset.
• Fast-WAM-pretrained.pt: Fast-WAM checkpoint pretrained on the same RoboCOIN subset using our reproduction implementation.

Approximate file sizes:

File	Size
AHA-WAM-pretrained.pt	13.7 GiB
Fast-WAM-pretrained.pt	11.2 GiB

Model Details

• Model family: Wan2.2-based world-action models for robot policy learning
• Backbone: Wan2.2-TI2V-5B components with an ActionDiT action branch
• Released variants: AHA-WAM and Fast-WAM
• Pretraining data: a selected subset of RoboCOIN
• Pretraining length: 45k optimization steps
• Action dimension: 14
• Proprioception dimension: 14
• Observation cameras: front/head camera plus left and right wrist cameras
• Training frame setup: 65-frame clips with action_video_freq_ratio=8
• Primary use: initialization for downstream real-robot fine-tuning and deployment

These are PyTorch checkpoints for the AHA-WAM codebase. They are not transformers.from_pretrained() checkpoints.

Important Note on Fast-WAM

The released Fast-WAM-pretrained.pt is provided for reference and convenience. It is our reproduction version of Fast-WAM within this codebase. For a fair method comparison against the original Fast-WAM, please control all variables carefully and reproduce the original method under matched data, preprocessing, architecture, training schedule, inference settings, and evaluation protocol.

Installation

Install the public AHA-WAM codebase first:

git clone https://github.com/serene-sivy/AHA-WAM.git
cd AHA-WAM

conda create -n ahawam python=3.10 -y
conda activate ahawam
pip install -U pip setuptools wheel
pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 \
  --extra-index-url https://download.pytorch.org/whl/cu128
pip install -e .

AHA-WAM also needs the Wan base components. Follow the project README for the full setup. If the assets are already downloaded, point DiffSynth to them:

export DIFFSYNTH_MODEL_BASE_PATH=/path/to/wan_models

Download

Use huggingface_hub:

from huggingface_hub import hf_hub_download

repo_id = "SereneC/AHA-WAM-Pretrained"

aha_ckpt = hf_hub_download(
    repo_id=repo_id,
    filename="AHA-WAM-pretrained.pt",
    local_dir="checkpoints/AHA-WAM-Pretrained",
)
fastwam_ckpt = hf_hub_download(
    repo_id=repo_id,
    filename="Fast-WAM-pretrained.pt",
    local_dir="checkpoints/AHA-WAM-Pretrained",
)

print("AHA-WAM:", aha_ckpt)
print("Fast-WAM:", fastwam_ckpt)

Or use the CLI:

huggingface-cli download SereneC/AHA-WAM-Pretrained \
  AHA-WAM-pretrained.pt Fast-WAM-pretrained.pt \
  --local-dir checkpoints/AHA-WAM-Pretrained

Use for Fine-Tuning

The checkpoints can be used as weight initialization through the training config's init_checkpoint field. For example, to fine-tune AHA-WAM:

bash scripts/train_zero1.sh 8 \
  task=robotwin_ahawam \
  model=ahawam \
  init_checkpoint=checkpoints/AHA-WAM-Pretrained/AHA-WAM-pretrained.pt

For Fast-WAM:

bash scripts/train_zero1.sh 8 \
  task=robotwin_ahawam \
  model=fastwam \
  init_checkpoint=checkpoints/AHA-WAM-Pretrained/Fast-WAM-pretrained.pt

Replace task=robotwin_ahawam with your downstream real-robot task config and make sure the processor, action/state dimensions, camera layout, normalization statistics, and language instruction format match your data.

Load the Checkpoints Directly

The checkpoints can also be loaded by instantiating the Hydra model config and calling load_checkpoint().

AHA-WAM

import sys
from pathlib import Path

import torch
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from hydra.utils import instantiate

PROJECT_ROOT = Path("/path/to/AHA-WAM").resolve()
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "src"))

ckpt_path = PROJECT_ROOT / "checkpoints/AHA-WAM-Pretrained/AHA-WAM-pretrained.pt"

if GlobalHydra.instance().is_initialized():
    GlobalHydra.instance().clear()

with initialize_config_dir(version_base="1.3", config_dir=str(PROJECT_ROOT / "configs")):
    cfg = compose(
        config_name="deploy",
        overrides=[
            "model=ahawam",
            "model.load_text_encoder=true",
            "model.skip_dit_load_from_pretrain=true",
            "model.action_dit_pretrained_path=null",
        ],
    )

model = instantiate(cfg.model, model_dtype=torch.bfloat16, device="cuda")
model.load_checkpoint(str(ckpt_path))
model = model.cuda().eval()

Fast-WAM

import sys
from pathlib import Path

import torch
from hydra import compose, initialize_config_dir
from hydra.core.global_hydra import GlobalHydra
from hydra.utils import instantiate

PROJECT_ROOT = Path("/path/to/AHA-WAM").resolve()
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "src"))

ckpt_path = PROJECT_ROOT / "checkpoints/AHA-WAM-Pretrained/Fast-WAM-pretrained.pt"

if GlobalHydra.instance().is_initialized():
    GlobalHydra.instance().clear()

with initialize_config_dir(version_base="1.3", config_dir=str(PROJECT_ROOT / "configs")):
    cfg = compose(
        config_name="deploy",
        overrides=[
            "model=fastwam",
            "model.load_text_encoder=true",
            "model.skip_dit_load_from_pretrain=true",
            "model.action_dit_pretrained_path=null",
        ],
    )

model = instantiate(cfg.model, model_dtype=torch.bfloat16, device="cuda")
model.load_checkpoint(str(ckpt_path))
model = model.cuda().eval()

Real-Robot Deployment

For real-robot deployment, use the deployment utilities in the AHA-WAM repository:

deploy/
├── deploy_example.yml
├── server/wam_policy_server.py
└── client/wam_remote_client_node.py

Start from the example config:

cp deploy/deploy_example.yml deploy/deploy.yml

Set at least:

checkpoint_path: checkpoints/AHA-WAM-Pretrained/AHA-WAM-pretrained.pt
dataset_stats_path: /path/to/your/downstream_dataset_stats.json
project_root: .
hydra_config_name: deploy
task: your_downstream_task

dataset_stats_path should match the downstream real-robot data used for fine-tuning or deployment. This repository only provides pretrained weights; it does not define a universal normalization file for every robot setup.

Limitations

• These checkpoints are pretrained on a subset of RoboCOIN and are intended as initialization weights, not as universal zero-shot robot policies.
• Downstream performance depends heavily on camera calibration, action conventions, state representation, normalization statistics, robot hardware, task instructions, and fine-tuning data quality.
• Real-robot deployment requires careful safety checks, action limits, and environment-specific validation.
• The Fast-WAM checkpoint is a reproduction checkpoint from this codebase and should not be treated as an official original Fast-WAM release.
• The checkpoints should be used with the AHA-WAM codebase; they are not standalone Hugging Face Transformers models.

License

The AHA-WAM code release is under the MIT License unless otherwise noted. Please also follow the licenses and usage terms of upstream dependencies, models, and datasets, including Wan2.2 and RoboCOIN.

Citation

If you use these checkpoints, please cite the AHA-WAM project and the upstream datasets or models used in your work.

@article{cai2026ahawam,
  title={AHA-WAM: Asynchronous Horizon-Adaptive World-Action Modeling with Observation-Guided Context Routing},
  author={Cai, Jisong and Ling, Long and Chu, Shiwei and Liu, Zhongshan and Kang, Jiayue and Liang, Zhixuan and Xu, Wenjie and Mao, Yinan and Zhang, Weinan and Yang, Xiaokang and Ying, Ru and Zheng, Ran and Mu, Yao},
  journal={arXiv preprint arXiv:2606.09811},
  year={2026}
}