Add FastWAM policy

2026-07-25 18:56:09 +00:00 · 2026-05-24 16:10:04 +00:00
parent 49755a3d9e
commit b3a28a49f6
31 changed files with 9957 additions and 1 deletions
@@ -0,0 +1,254 @@
+#!/usr/bin/env python
+
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import typing
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from lerobot.configs import FeatureType, PolicyFeature
+from lerobot.policies.fastwam.configuration_fastwam import FastWAMConfig
+from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy
+from lerobot.policies.fastwam.processor_fastwam import make_fastwam_pre_post_processors
+from lerobot.utils.constants import OBS_STATE
+
+ROOT = Path(__file__).resolve().parents[3]
+
+
+def test_package_init_exports_required_symbols():
+    init_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "__init__.py").read_text()
+
+    assert "FastWAMConfig" in init_source
+    assert "make_fastwam_pre_post_processors" in init_source
+
+
+def test_policy_config_is_exported_from_public_policies_package():
+    import lerobot.policies as policies
+
+    assert policies.FastWAMConfig is FastWAMConfig
+    assert "FastWAMConfig" in policies.__all__
+
+
+def test_fastwam_policy_docs_are_registered():
+    readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "README.md"
+    wan_readme_path = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan" / "README.md"
+    policy_readme_path = ROOT / "docs" / "source" / "policy_fastwam_README.md"
+    guide_path = ROOT / "docs" / "source" / "fastwam.mdx"
+    toctree_path = ROOT / "docs" / "source" / "_toctree.yml"
+
+    assert readme_path.is_symlink()
+    assert readme_path.resolve() == policy_readme_path.resolve()
+    assert wan_readme_path.exists()
+    wan_readme = wan_readme_path.read_text()
+    assert "Wan-Video/Wan2.2" in wan_readme
+    assert "42bf4cfaa384bc21833865abc2f9e6c0e67233dc" in wan_readme
+    assert policy_readme_path.exists()
+    assert guide_path.exists()
+    assert "local: fastwam" in toctree_path.read_text()
+
+
+def test_wan_backbone_code_is_isolated_from_lerobot_adapter():
+    wan_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan"
+
+    assert (wan_dir / "modules" / "attention.py").exists()
+    assert (wan_dir / "modules" / "model.py").exists()
+    assert (wan_dir / "modules" / "t5.py").exists()
+    assert (wan_dir / "modules" / "tokenizers.py").exists()
+    assert (wan_dir / "modules" / "vae2_1.py").exists()
+    assert (wan_dir / "modules" / "vae2_2.py").exists()
+    assert (wan_dir / "utils" / "fm_solvers.py").exists()
+    assert (wan_dir / "utils" / "fm_solvers_unipc.py").exists()
+
+    assert (wan_dir.parent / "wan_video_dit.py").exists()
+    assert (wan_dir.parent / "wan_adapters.py").exists()
+    assert (wan_dir.parent / "wan_components.py").exists()
+    assert not (wan_dir / "wan_video_dit.py").exists()
+    assert not (wan_dir / "wan_adapters.py").exists()
+    assert not (wan_dir / "wan_components.py").exists()
+
+
+def test_fastwam_text_encoder_uses_upstream_wan_modules_directly():
+    fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
+    modular_source = (fastwam_dir / "modular_fastwam.py").read_text()
+    components_source = (fastwam_dir / "wan_components.py").read_text()
+
+    assert not (fastwam_dir / "wan_video_text_encoder.py").exists()
+    assert "from .wan.modules.t5 import umt5_xxl" in components_source
+    assert "from .wan.modules.tokenizers import HuggingfaceTokenizer" in components_source
+    assert "WAN_T5_ENCODER_KWARGS" not in components_source
+    assert "wan_video_text_encoder" not in modular_source
+
+
+def test_fastwam_vae_reuses_upstream_wan_modules():
+    fastwam_dir = ROOT / "src" / "lerobot" / "policies" / "fastwam"
+    vae_source = (fastwam_dir / "wan_adapters.py").read_text()
+
+    assert not (fastwam_dir / "wan_video_vae.py").exists()
+    assert "from .wan.modules.vae2_2 import Wan2_2_VAE" in vae_source
+    assert "mean = [" not in vae_source
+    assert "std = [" not in vae_source
+    assert "class Encoder3d_38" not in vae_source
+    assert "class Decoder3d_38" not in vae_source
+    assert "class VideoVAE38_" not in vae_source
+
+
+def test_fastwam_component_loading_uses_fixed_wan_checkpoint_layout():
+    modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
+    modeling_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modeling_fastwam.py").read_text()
+    components_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_components.py").read_text()
+
+    assert "class ModelConfig" not in modular_source
+    assert "def load_state_dict" not in modular_source
+    assert "WAN22_MODEL_REGISTRY" not in modular_source
+    assert "class ModelConfig" not in components_source
+    assert "class WanComponentSource" not in components_source
+    assert "def load_state_dict" not in components_source
+    assert "WAN22_MODEL_REGISTRY" not in components_source
+    assert "hash_model_file" not in components_source
+    assert "_resolve_component_sources" not in components_source
+    assert "origin_file_pattern" not in components_source
+    assert "inspect.signature" not in components_source
+    assert "class FastWAMWanComponentPaths" not in modeling_source
+    assert "def _first_existing" not in modeling_source
+    assert "def _missing_wan_component_names" not in modeling_source
+    assert "WAN_T5_CHECKPOINT" in components_source
+    assert "WAN_VAE_CHECKPOINT" in components_source
+    assert "WAN_DIT_PATTERN" in components_source
+
+
+def test_fastwam_dit_reuses_upstream_wan_primitives():
+    dit_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "wan_video_dit.py").read_text()
+
+    assert "from .wan.modules.model import" in dit_source
+    assert "WanModel" in dit_source
+    for duplicated_symbol in [
+        "def flash_attention(",
+        "def sinusoidal_embedding_1d(",
+        "def rope_apply(",
+        "def unpatchify(",
+        "def _dense_video_freqs(",
+        "class RMSNorm(",
+        "class SelfAttention(",
+        "class CrossAttention(",
+        "class Head(",
+    ]:
+        assert duplicated_symbol not in dit_source
+
+
+def test_fastwam_inference_schedule_reuses_upstream_wan_sigmas():
+    modular_source = (ROOT / "src" / "lerobot" / "policies" / "fastwam" / "modular_fastwam.py").read_text()
+
+    assert "def _get_wan_sampling_sigmas" in modular_source
+    assert "from .wan.utils.fm_solvers import get_sampling_sigmas" in modular_source
+    assert "_get_wan_sampling_sigmas(num_inference_steps, shift)" in modular_source
+
+
+def test_policy_config_rejects_missing_required_image_and_action_features():
+    with pytest.raises(ValueError, match="image feature"):
+        FastWAMConfig(
+            input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
+        )
+
+    with pytest.raises(ValueError, match="action"):
+        FastWAMConfig(
+            output_features={"not_action": PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
+        )
+
+
+def test_policy_init_calls_validate_features_even_for_prebuilt_configs(monkeypatch):
+    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2)
+    calls = []
+
+    def record_validate_features():
+        calls.append("called")
+
+    monkeypatch.setattr(cfg, "validate_features", record_validate_features)
+    monkeypatch.setattr(
+        FastWAMPolicy,
+        "_build_core_model",
+        lambda self, config: nn.Linear(1, 1),
+    )
+    FastWAMPolicy(cfg)
+
+    assert calls == ["called"]
+
+
+def test_required_policy_entrypoints_exist_with_discoverable_names():
+    assert FastWAMPolicy.config_class is FastWAMConfig
+    assert FastWAMPolicy.name == "fastwam"
+    assert callable(FastWAMPolicy.reset)
+    assert callable(FastWAMPolicy.get_optim_params)
+    assert callable(FastWAMPolicy.predict_action_chunk)
+    assert callable(FastWAMPolicy.select_action)
+    assert callable(FastWAMPolicy.forward)
+    assert callable(make_fastwam_pre_post_processors)
+    assert make_fastwam_pre_post_processors.__name__ == "make_fastwam_pre_post_processors"
+
+
+def test_policy_constructor_and_forward_match_byo_template_contract():
+    init_signature = inspect.signature(FastWAMPolicy.__init__)
+
+    assert "dataset_stats" in init_signature.parameters
+    assert "core_model" not in init_signature.parameters
+    assert typing.get_type_hints(FastWAMPolicy.forward)["return"] == dict[str, torch.Tensor]
+
+
+def test_saved_config_round_trips_policy_features(tmp_path):
+    cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448))
+    cfg.save_pretrained(tmp_path)
+
+    loaded = FastWAMConfig.from_pretrained(tmp_path)
+
+    assert loaded.type == "fastwam"
+    assert loaded.image_features["observation.images.image"].type == FeatureType.VISUAL
+    assert loaded.action_feature.shape == (7,)
+    assert loaded.robot_state_feature.shape == (8,)
+
+
+def test_config_from_pretrained_ignores_unknown_fields(tmp_path):
+    cfg = FastWAMConfig()
+    cfg.save_pretrained(tmp_path)
+    config_path = tmp_path / "config.json"
+    payload = config_path.read_text()
+    payload = payload.replace(
+        '"torch_dtype": "bfloat16"',
+        '"torch_dtype": "bfloat16",\n  "unknown_fastwam_field": true',
+    )
+    config_path.write_text(payload)
+
+    loaded = FastWAMConfig.from_pretrained(tmp_path)
+
+    assert loaded.type == "fastwam"
+    assert not hasattr(loaded, "unknown_fastwam_field")
+
+
+def test_config_from_pretrained_does_not_use_non_wan22_tokenizer_repo_id(tmp_path):
+    cfg = FastWAMConfig()
+    cfg.save_pretrained(tmp_path)
+    config_path = tmp_path / "config.json"
+    payload = config_path.read_text()
+    payload = payload.replace(
+        '"tokenizer_model_id": "Wan-AI/Wan2.2-TI2V-5B"',
+        '"tokenizer_model_id": "somebody/old-tokenizer"',
+    )
+    config_path.write_text(payload)
+
+    loaded = FastWAMConfig.from_pretrained(tmp_path)
+
+    assert loaded.tokenizer_model_id == "Wan-AI/Wan2.2-TI2V-5B"