| --- |
| library_name: lerobot |
| license: apache-2.0 |
| pipeline_tag: robotics |
| base_model: amd/CAE-FastWAM |
| datasets: |
| - yuanty/LIBERO-fastwam |
| tags: |
| - cae_fastwam |
| - lerobot |
| - libero |
| - robotics |
| - world-action-model |
| --- |
| |
| # CAE-FastWAM (LeRobot format) |
|
|
| [CAE-FastWAM](https://huggingface.co/amd/CAE-FastWAM) converted to the [LeRobot](https://github.com/huggingface/lerobot) checkpoint format. CAE-FastWAM is a [FastWAM](https://arxiv.org/abs/2603.16666) world action model with a Compact Action Expert: a small (32-wide) action expert denoises latent action tokens, and a LIBERO-suite-conditioned action decoder turns them into the action chunk, on top of FastWAM's Wan2.2-TI2V-5B video expert. |
|
|
| The weights are the same as [amd/CAE-FastWAM](https://huggingface.co/amd/CAE-FastWAM); this repo only changes the layout (`config.json`, `model.safetensors`, and the LeRobot pre/post-processors with the per-suite normalisation statistics). It also carries the frozen action-VAE encoder, so the policy can be fine-tuned with the original CAE-FastWAM objective. |
|
|
| > **Note:** loading this checkpoint needs the `cae_fastwam` policy type, which is not in a LeRobot release yet. This card will be updated with install instructions once it is available in LeRobot. |
| |
| ## Inputs and outputs |
| |
| | Feature | Type | Shape | |
| | --- | --- | --- | |
| | `observation.images.image` (agent view) | visual | (3, 224, 224) | |
| | `observation.images.image2` (wrist view) | visual | (3, 224, 224) | |
| | `observation.state` | state | (8,) | |
| | `action` | action | (7,) | |
| |
| Select the LIBERO suite with `--policy.task_suite` (`libero_spatial`, `libero_object`, `libero_goal` or `libero_10`); evaluate one suite per run. |
|
|
| ## Results |
|
|
| LIBERO, bfloat16, 10 denoising steps, 50 episodes per task: |
|
|
| | Suite | Success rate | Episodes | |
| | --- | ---: | ---: | |
| | libero_spatial | 97.0% | 500 | |
| | libero_object | 99.8% | 500 | |
| | libero_goal | 94.4% | 500 | |
| | libero_10 | 95.6% | 500 | |
| | **Average** | **96.7%** | 2000 | |
|
|
| Reproduce (per suite; `--env.episode_length=700` for `libero_10`), on one AMD Instinct MI355X: |
|
|
| ```bash |
| lerobot-eval --policy.path=amd/CAE-FastWAM-lerobot --policy.task_suite=libero_spatial \ |
| --env.type=libero --env.task=libero_spatial \ |
| --env.observation_height=256 --env.observation_width=256 --env.episode_length=400 \ |
| --eval.batch_size=5 --eval.n_episodes=50 --seed=42 |
| ``` |
|
|
| LIBERO stack: the one pinned by LeRobot's LIBERO benchmark image (hf-libero 0.1.4, robosuite 1.4.0, MuJoCo 3.3.1), CPU physics, OSMesa rendering. The MuJoCo version matters: with MuJoCo 3.8.1 the `libero_spatial` task 5 scene settles differently and the same evaluation scored 94.5% on average. Please report the MuJoCo version with LIBERO results. |
|
|
| For reference, the [amd/CAE-FastWAM](https://huggingface.co/amd/CAE-FastWAM) card reports 97.28 ± 0.20% with GEAK-optimized inference and a GPU-simulated LIBERO (mean of 11 runs of 2000 episodes). |
|
|
| ## Fine-tuning |
|
|
| The checkpoint holds the frozen action-VAE encoder and decoder, so `lerobot-train --policy.path=amd/CAE-FastWAM-lerobot --policy.task_suite=<suite> --accelerator.mixed_precision=bf16 ...` fine-tunes the video and action experts with the original CAE-FastWAM objective, one suite per run. |
|
|
| ## Provenance and license |
|
|
| Converted from [amd/CAE-FastWAM](https://huggingface.co/amd/CAE-FastWAM) (Apache-2.0, Advanced Micro Devices, Inc.), a fine-tune of [Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) with the [FastWAM](https://github.com/yuantianyuan01/FastWAM) recipe on LIBERO. See `LICENSE` and `NOTICE`. |
|
|
