CAE-FastWAM in LeRobot format (converted from amd/CAE-FastWAM; LIBERO 96.7%)

8ecb02e verified about 20 hours ago

3.52 kB

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 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, 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:

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 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.