CAE-FastWAM (LeRobot format)

CAE-FastWAM converted to the LeRobot checkpoint format. CAE-FastWAM is a FastWAM 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; this repo only changes the layout: config.json (including the per-suite min/max normalisation statistics, task_suite_stats), model.safetensors, and FastWAM's standard LeRobot pre/post-processors, which use the statistics of the selected suite. It also carries the frozen action-VAE encoder, so the policy can be fine-tuned with the original CAE-FastWAM objective.

Installation

The cae_fastwam policy type is proposed for LeRobot in huggingface/lerobot#5317 and is not merged yet. Until then, install LeRobot from the commit used for this checkpoint:

pip install "lerobot[fastwam,libero] @ git+https://github.com/hiteshK03/lerobot.git@e0606546945b93f0cfe87c98918904279dba94ad"

Once the PR is merged, install from upstream instead:

pip install "lerobot[fastwam,libero] @ git+https://github.com/huggingface/lerobot.git@main"

Inclusion in a LeRobot release is decided by the LeRobot maintainers.

The script used to convert the original checkpoint is in converter/ (see its README).

Inputs and outputs

Feature	Type	Shape
`observation.images.image` (agent view)	visual	(3, 224, 224), resized by the model as in FastWAM
`observation.images.image2` (wrist view)	visual	(3, 224, 224), resized by the model as in FastWAM
`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 with LeRobot's lerobot-eval: bfloat16, 10 denoising steps, 256 x 256 camera frames, seed 42, 50 episodes per task (10 tasks per suite), at most 400 steps per episode (700 for libero_10):

Suite	Success rate	Episodes
libero_spatial	96.0% (480/500)	500
libero_object	99.4% (497/500)	500
libero_goal	94.4% (472/500)	500
libero_10	95.8% (479/500)	500
Average	96.4% (1928/2000)	2000

This version processes observations with FastWAM's own path (in-model image resize, LeRobot's standard min/max normaliser). The previous version of this repo, with its own image resize step and normaliser, scored 96.7% (1934/2000) in an earlier run of this protocol. Run side by side on the same 2000 episodes, the previous version scored 96.3% (1926/2000) and this version 96.4%; no task differed significantly.

Reproduce (per suite; --env.episode_length=700 for libero_10), on one AMD Instinct MI355X:

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 previous version of this repo scored 94.5% on average under it. Please report the MuJoCo version with LIBERO results.

For reference, the 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 (Apache-2.0, Advanced Micro Devices, Inc.), a fine-tune of Wan2.2-TI2V-5B with the FastWAM recipe on LIBERO. See LICENSE and NOTICE.