--- license: mit base_model: - Wan-AI/Wan2.2-TI2V-5B tags: - pytorch - robotics - robot-learning - imitation-learning - world-model - WAM - robotwin - ahawam model-index: - name: AHA-WAM-RoboTwin2.0 results: [] --- # AHA-WAM RoboTwin2.0 Checkpoint This repository provides a RoboTwin2.0 checkpoint for **AHA-WAM**, together with the dataset normalization statistics required for reproducible evaluation. AHA-WAM is a Wan2.2-based world-action model for robot policy learning. This checkpoint was trained with the `robotwin_ahawam` task configuration from the AHA-WAM codebase: ```text configs/task/robotwin_ahawam.yaml ``` The checkpoint is intended for direct loading, reproduction, and RoboTwin evaluation with the public AHA-WAM code release. ## Repository Files ```text AHA-WAM-RoboTwin2.0/ ├── README.md ├── config.json ├── dataset_stats.json ├── robotwin_ahawam.pt └── robotwin_ahawam-flash.pt ``` - `config.json`: checkpoint manifest and Hugging Face download-count query file. - `robotwin_ahawam.pt`: AHA-WAM checkpoint trained on RoboTwin2.0-format data. - `robotwin_ahawam-flash.pt`: AHA-WAM-Flash checkpoint trained on RoboTwin2.0-format data, after ODE distillation from AHA-WAM. - `dataset_stats.json`: normalization statistics used by the AHA-WAM processor for action/state normalization and denormalization. - `README.md`: this model card. ## Model Details - **Model type:** AHA-WAM robot world-action model - **Backbone:** Wan2.2-TI2V-5B components with an ActionDiT action branch - **Training task config:** `configs/task/robotwin_ahawam.yaml` - **RoboTwin action dimension:** 14 - **RoboTwin proprioception dimension:** 14 - **Observation cameras:** `cam_high`, `cam_left_wrist`, `cam_right_wrist` - **Training/evaluation frame setup:** 65-frame trajectories with `action_video_freq_ratio=8` - **Default evaluation task config:** RoboTwin `demo_randomized` - **Default evaluation instruction split:** `unseen` ## Installation First install the public AHA-WAM codebase: ```bash # Clone the AHA-WAM GitHub repository, then enter the codebase. cd AHA-WAM-release conda create -n ahawam python=3.10 -y conda activate ahawam pip install -U pip pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 \ --extra-index-url https://download.pytorch.org/whl/cu128 pip install -e . pip install huggingface_hub ``` AHA-WAM also uses Wan model assets. Put the Wan assets in your local model directory and point DiffSynth to them: ```bash export DIFFSYNTH_MODEL_BASE_PATH=/path/to/wan_models ``` Please follow the AHA-WAM repository README for the full environment and RoboTwin setup. ## Download Download `config.json` together with the checkpoint and normalization stats so the Hub can count downloads server-side using the repository's query file: ```python from huggingface_hub import hf_hub_download repo_id = "SereneC/AHA-WAM-RoboTwin2.0" manifest_path = hf_hub_download( repo_id=repo_id, filename="config.json", local_dir="checkpoints/AHA-WAM-RoboTwin2.0", ) ckpt_path = hf_hub_download( repo_id=repo_id, filename="robotwin_ahawam.pt", local_dir="checkpoints/AHA-WAM-RoboTwin2.0", ) stats_path = hf_hub_download( repo_id=repo_id, filename="dataset_stats.json", local_dir="checkpoints/AHA-WAM-RoboTwin2.0", ) print("manifest:", manifest_path) print("checkpoint:", ckpt_path) print("dataset stats:", stats_path) ``` Or use the CLI: ```bash hf download SereneC/AHA-WAM-RoboTwin2.0 \ config.json robotwin_ahawam.pt dataset_stats.json \ --local-dir checkpoints/AHA-WAM-RoboTwin2.0 ``` ## Load the Model This checkpoint is not a `transformers.from_pretrained()` checkpoint. Load it with the AHA-WAM codebase by instantiating the Hydra model config and calling `load_checkpoint()`. ```python 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-release").resolve() sys.path.insert(0, str(PROJECT_ROOT)) sys.path.insert(0, str(PROJECT_ROOT / "src")) ckpt_path = PROJECT_ROOT / "checkpoints/AHA-WAM-RoboTwin2.0/robotwin_ahawam.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="sim_robotwin.yaml", overrides=[ "task=robotwin_ahawam", "model.load_text_encoder=true", "model.skip_dit_load_from_pretrain=true", "model.action_dit_pretrained_path=null", ], ) device = "cuda" model = instantiate(cfg.model, model_dtype=torch.bfloat16, device=device) model.load_checkpoint(str(ckpt_path)) model = model.to(device).eval() print("Loaded AHA-WAM checkpoint:", ckpt_path) ``` For RoboTwin evaluation, the policy adapter loads the same checkpoint internally, so most users can use the evaluation command below directly. ## RoboTwin Evaluation Install and prepare RoboTwin following the upstream RoboTwin instructions. Then point AHA-WAM to your RoboTwin checkout: ```bash ROBOTWIN_ROOT=/path/to/RoboTwin ln -sfn "$(pwd)/experiments/robotwin/ahawam_policy" \ "$ROBOTWIN_ROOT/policy/ahawam_policy" ``` Run the multi-GPU RoboTwin evaluation manager: ```bash python experiments/robotwin/run_robotwin_manager.py \ task=robotwin_ahawam \ ckpt=checkpoints/AHA-WAM-RoboTwin2.0/robotwin_ahawam.pt \ EVALUATION.dataset_stats_path=checkpoints/AHA-WAM-RoboTwin2.0/dataset_stats.json \ EVALUATION.robotwin_root=/path/to/RoboTwin \ MULTIRUN.num_gpus=8 ``` Useful overrides: ```bash # Change the number of GPUs used by the evaluation manager. MULTIRUN.num_gpus= # Select a specific RoboTwin task. EVALUATION.task_name= # Change evaluation episodes. EVALUATION.eval_num_episodes= ``` ## Training Configuration The released checkpoint corresponds to: ```yaml task: robotwin_ahawam batch_size: 6 num_epochs: 5 learning_rate: 5e-5 weight_decay: 1e-2 gradient_accumulation_steps: 1 model: action_horizon: 64 action_chunk_size: 16 num_history_frames: 6 action_video_read_mode: current_only chunk_kv_editor_num_queries: 32 loss: lambda_action_prior: 1.0 data: train: pretrained_norm_stats: dataset_stats.json val: pretrained_norm_stats: dataset_stats.json ``` For the complete configuration, see `configs/task/robotwin_ahawam.yaml` in the AHA-WAM codebase. ## Limitations - The checkpoint is trained for the RoboTwin2.0 setup used by the AHA-WAM release and expects matching observation/action conventions. - `dataset_stats.json` must be used with the checkpoint for correct action and state normalization. - Evaluation results can vary with RoboTwin version, simulator assets, task split, random seed, GPU precision, and local environment setup. - The checkpoint should be used together with the AHA-WAM codebase; it is not a standalone Hugging Face Transformers model. ## 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 and assets, including Wan2.2 and RoboTwin. ## Citation If you use this checkpoint, please cite the AHA-WAM project and any upstream datasets or models used in your work. ```bibtex @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} } ```