Use FastWAM's processors; per-suite statistics in config.json; LIBERO 96.4%

The checkpoint now uses FastWAM's standard LeRobot processors: in-model image resize and LeRobot's min/max (un)normalizer, given the statistics of --policy.task_suite from config.json (task_suite_stats). The two cae_fastwam_suite_* processor state files are removed. model.safetensors is unchanged. README: LIBERO success rate with this version, 96.4% (1928/2000; lerobot-eval, bf16, 10 denoising steps, 50 episodes per task, MuJoCo 3.3.1). SHA256SUMS regenerated.

README.md

@@ -17,7 +17,7 @@ tags:
 
 [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.
+The weights are the same as [amd/CAE-FastWAM](https://huggingface.co/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.
 
 > **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.
 
@@ -25,8 +25,8 @@ The weights are the same as [amd/CAE-FastWAM](https://huggingface.co/amd/CAE-Fas
 
 | Feature | Type | Shape |
 | --- | --- | --- |
-| `observation.images.image` (agent view) | visual | (3, 224, 224) |
-| `observation.images.image2` (wrist view) | visual | (3, 224, 224) |
+| `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,) |
 
@@ -34,15 +34,17 @@ Select the LIBERO suite with `--policy.task_suite` (`libero_spatial`, `libero_ob
 
 ## Results
 
-LIBERO, bfloat16, 10 denoising steps, 50 episodes per task:
+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 | 97.0% | 500 |
-| libero_object | 99.8% | 500 |
-| libero_goal | 94.4% | 500 |
-| libero_10 | 95.6% | 500 |
-| **Average** | **96.7%** | 2000 |
+| 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:
 
@@ -53,7 +55,7 @@ lerobot-eval --policy.path=amd/CAE-FastWAM-lerobot --policy.task_suite=libero_sp
   --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.
+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](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).
 

SHA256SUMS

@@ -1,9 +1,7 @@
-d0eb115feb3471e0f1b98b4dfd49d559d98a0dab80f5e641615383b45b242443  config.json
+13c426f2b2b4b13e57f164c3cc85629955fb11e6b6a0daa0b471e96f8db2af9f  config.json
 cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30  LICENSE
-ce18712e4a1c6577dd7ec0970b4ebd61835763155149d8c9ec390529d7bf20f2  model.safetensors
 07b1f41737ff5020427958e08e62c07a176ae395e39d7e41e1f237defa38378a  NOTICE
-60c427e478f68f902bb019e50ffe9a0e55f71799faa7a24130cca31150f92161  policy_postprocessor.json
-d7eaf2818654f5f87424b23096096d437719d1a7e5f3dfa3ec5cbfd958c0829a  policy_postprocessor_step_0_cae_fastwam_suite_unnormalizer.safetensors
-8cdc442a212b6d414644dfdd130fdd291483282a9cad8b3db56dcda2ec5a8d57  policy_preprocessor.json
-21d67e2e1a64200163b7ae04c08aecde9f451e7ed4f386d35be3116dc3d8cdba  policy_preprocessor_step_4_cae_fastwam_suite_normalizer.safetensors
-9e28a2b08d21e0ecd28d160841120da0a41b953607e59465e02f2558156628c3  README.md
+8f8ce7bd200cf8aafa8480a255ded93a24353180185caffa68662dca53909f7b  README.md
+ce18712e4a1c6577dd7ec0970b4ebd61835763155149d8c9ec390529d7bf20f2  model.safetensors
+2e42c65ceb693e448f07687cf500a44cd4426ec1ed300af89fc0e168cae86d11  policy_postprocessor.json
+e97371868054240c3d98d6672dff18ab1154ed21f719e3bdfd01f149b104ecae  policy_preprocessor.json

config.json

@@ -85,7 +85,6 @@
     "sigma_shift": null,
     "tiled": false,
     "fp32_attention": false,
-    "fp32_upcast": false,
     "compile_action_infer": false,
     "use_gradient_checkpointing": false,
     "freeze_video_expert": false,
@@ -128,8 +127,7 @@
         "action_conditioned": false,
         "action_dim": 7,
         "action_group_causal_mask_mode": "group_diagonal",
-        "fp32_attention": false,
-        "fp32_upcast": false
+        "fp32_attention": false
     },
     "action_dit_config": {
         "action_dim": 16,
@@ -143,7 +141,6 @@
         "eps": 1e-06,
         "use_gradient_checkpointing": false,
         "fp32_attention": false,
-        "fp32_upcast": false,
         "use_text_cross_attention": true
     },
     "normalization_mapping": {
@@ -158,19 +155,189 @@
     "action_decoder_hidden_dim": 128,
     "action_encoder_hidden_dim": 128,
     "action_encoder_variational": true,
-    "delta_action_dims": [
-        0,
-        1,
-        2,
-        3,
-        4,
-        5
-    ],
     "task_suites": [
         "libero_spatial",
         "libero_object",
         "libero_goal",
         "libero_10"
     ],
-    "task_suite": null
+    "task_suite": null,
+    "task_suite_stats": {
+        "libero_spatial": {
+            "observation.state": {
+                "min": [
+                    -0.4828203022480011,
+                    -0.3255046010017395,
+                    0.008128181099891663,
+                    1.002794623374939,
+                    -3.641430377960205,
+                    -1.842738389968872,
+                    -0.0021671096328645945,
+                    -0.042015016078948975
+                ],
+                "max": [
+                    0.21031762659549713,
+                    0.39128610491752625,
+                    1.472778081893921,
+                    3.6714255809783936,
+                    3.560650587081909,
+                    1.386339545249939,
+                    0.04232141748070717,
+                    0.0013633022317662835
+                ]
+            },
+            "action": {
+                "min": [
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.18857142329216003,
+                    -0.3675000071525574,
+                    -0.36000001430511475,
+                    0.0
+                ],
+                "max": [
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.1971428543329239,
+                    0.33642858266830444,
+                    0.375,
+                    1.0
+                ]
+            }
+        },
+        "libero_object": {
+            "observation.state": {
+                "min": [
+                    -0.4828203022480011,
+                    -0.3255046010017395,
+                    0.008128181099891663,
+                    1.002794623374939,
+                    -3.641430377960205,
+                    -1.842738389968872,
+                    -0.0021671096328645945,
+                    -0.042015016078948975
+                ],
+                "max": [
+                    0.21031762659549713,
+                    0.39128610491752625,
+                    1.472778081893921,
+                    3.6714255809783936,
+                    3.560650587081909,
+                    1.386339545249939,
+                    0.04232141748070717,
+                    0.0013633022317662835
+                ]
+            },
+            "action": {
+                "min": [
+                    -0.8839285969734192,
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.15000000596046448,
+                    -0.29035714268684387,
+                    -0.32892856001853943,
+                    0.0
+                ],
+                "max": [
+                    0.9374999403953552,
+                    0.8919642567634583,
+                    0.9374999403953552,
+                    0.1767856925725937,
+                    0.35035714507102966,
+                    0.1810714304447174,
+                    1.0
+                ]
+            }
+        },
+        "libero_goal": {
+            "observation.state": {
+                "min": [
+                    -0.4828203022480011,
+                    -0.3255046010017395,
+                    0.008128181099891663,
+                    1.002794623374939,
+                    -3.641430377960205,
+                    -1.842738389968872,
+                    -0.0021671096328645945,
+                    -0.042015016078948975
+                ],
+                "max": [
+                    0.21031762659549713,
+                    0.39128610491752625,
+                    1.472778081893921,
+                    3.6714255809783936,
+                    3.560650587081909,
+                    1.386339545249939,
+                    0.04232141748070717,
+                    0.0013633022317662835
+                ]
+            },
+            "action": {
+                "min": [
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.24214285612106323,
+                    -0.375,
+                    -0.2871428430080414,
+                    0.0
+                ],
+                "max": [
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.35571423172950745,
+                    0.375,
+                    0.375,
+                    1.0
+                ]
+            }
+        },
+        "libero_10": {
+            "observation.state": {
+                "min": [
+                    -0.4828203022480011,
+                    -0.3255046010017395,
+                    0.008128181099891663,
+                    1.002794623374939,
+                    -3.641430377960205,
+                    -1.842738389968872,
+                    -0.0021671096328645945,
+                    -0.042015016078948975
+                ],
+                "max": [
+                    0.21031762659549713,
+                    0.39128610491752625,
+                    1.472778081893921,
+                    3.6714255809783936,
+                    3.560650587081909,
+                    1.386339545249939,
+                    0.04232141748070717,
+                    0.0013633022317662835
+                ]
+            },
+            "action": {
+                "min": [
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.9374999403953552,
+                    -0.23642855882644653,
+                    -0.3053571283817291,
+                    -0.3675000071525574,
+                    0.0
+                ],
+                "max": [
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.9374999403953552,
+                    0.30000001192092896,
+                    0.29357144236564636,
+                    0.375,
+                    1.0
+                ]
+            }
+        }
+    }
 }

policy_postprocessor.json

@@ -2,22 +2,29 @@
   "name": "policy_postprocessor",
   "steps": [
     {
-      "registry_name": "cae_fastwam_suite_unnormalizer",
+      "registry_name": "unnormalizer_processor",
       "config": {
-        "keys": [
-          "action"
-        ],
-        "task_suites": [
-          "libero_spatial",
-          "libero_object",
-          "libero_goal",
-          "libero_10"
-        ],
-        "task_suite": null,
-        "clamp": null,
-        "range_tol": 0.0001
-      },
-      "state_file": "policy_postprocessor_step_0_cae_fastwam_suite_unnormalizer.safetensors"
+        "eps": 1e-08,
+        "features": {
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              7
+            ]
+          }
+        },
+        "norm_map": {
+          "VISUAL": "IDENTITY",
+          "STATE": "MIN_MAX",
+          "ACTION": "MIN_MAX"
+        }
+      }
+    },
+    {
+      "registry_name": "absolute_actions_processor",
+      "config": {
+        "enabled": false
+      }
     },
     {
       "registry_name": "fastwam_action_toggle_processor",

policy_postprocessor_step_0_cae_fastwam_suite_unnormalizer.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d7eaf2818654f5f87424b23096096d437719d1a7e5f3dfa3ec5cbfd958c0829a
-size 856

policy_preprocessor.json

@@ -11,15 +11,6 @@
       "registry_name": "to_batch_processor",
       "config": {}
     },
-    {
-      "registry_name": "cae_fastwam_image_processor",
-      "config": {
-        "camera_size": [
-          224,
-          224
-        ]
-      }
-    },
     {
       "registry_name": "device_processor",
       "config": {
@@ -28,31 +19,55 @@
       }
     },
     {
-      "registry_name": "cae_fastwam_suite_normalizer",
+      "registry_name": "relative_actions_processor",
       "config": {
-        "keys": [
-          "observation.state",
-          "action"
-        ],
-        "task_suites": [
-          "libero_spatial",
-          "libero_object",
-          "libero_goal",
-          "libero_10"
+        "enabled": false,
+        "exclude_joints": [
+          "gripper"
         ],
-        "task_suite": null,
-        "clamp": 5.0,
-        "range_tol": 0.0001,
-        "delta_action_dims": [
-          0,
-          1,
-          2,
-          3,
-          4,
-          5
-        ]
-      },
-      "state_file": "policy_preprocessor_step_4_cae_fastwam_suite_normalizer.safetensors"
+        "action_names": null
+      }
+    },
+    {
+      "registry_name": "normalizer_processor",
+      "config": {
+        "eps": 1e-08,
+        "features": {
+          "observation.images.image": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              224,
+              224
+            ]
+          },
+          "observation.images.image2": {
+            "type": "VISUAL",
+            "shape": [
+              3,
+              224,
+              224
+            ]
+          },
+          "observation.state": {
+            "type": "STATE",
+            "shape": [
+              8
+            ]
+          },
+          "action": {
+            "type": "ACTION",
+            "shape": [
+              7
+            ]
+          }
+        },
+        "norm_map": {
+          "VISUAL": "IDENTITY",
+          "STATE": "MIN_MAX",
+          "ACTION": "MIN_MAX"
+        }
+      }
     }
   ]
 }

policy_preprocessor_step_4_cae_fastwam_suite_normalizer.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:21d67e2e1a64200163b7ae04c08aecde9f451e7ed4f386d35be3116dc3d8cdba
-size 1832