lerobot/tests/policies/fastwam/test_fastwam_policy.py

#!/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 json

import pytest
import torch
from safetensors.torch import save_model
from torch import nn

from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
from lerobot.policies import FastWAMConfig, get_policy_class, make_policy_config, make_pre_post_processors
from lerobot.policies.fastwam import modeling_fastwam
from lerobot.policies.fastwam.modeling_fastwam import FastWAMPolicy, resolve_wan_component_paths
from lerobot.policies.fastwam.processor_fastwam import FastWAMActionToggleProcessorStep
from lerobot.policies.fastwam.wan_components import (
    WAN_DIT_PATTERN,
    WAN_T5_CHECKPOINT,
    WAN_T5_TOKENIZER,
    WAN_VAE_CHECKPOINT,
    resolve_wan_checkpoint_paths,
)
from lerobot.utils.constants import ACTION, OBS_STATE


class FakeFastWAMCore(nn.Module):
    def __init__(self):
        super().__init__()
        self.dit = nn.Linear(2, 2)

    def training_loss(self, sample):
        assert sample["video"].ndim == 5
        assert sample["context"].ndim == 3
        return sample[ACTION].sum() * 0.0 + torch.tensor(1.0), {"loss_action": 1.0}

    def infer_action(self, **kwargs):
        return {"action": torch.ones(1, kwargs["action_horizon"], 3)}


def test_fastwam_is_registered_and_publicly_exported():
    cfg = make_policy_config(
        "fastwam",
        action_dim=3,
        proprio_dim=2,
        action_horizon=4,
        n_action_steps=2,
        base_model_id=None,
    )

    assert isinstance(cfg, FastWAMConfig)
    assert cfg.type == "fastwam"
    assert get_policy_class("fastwam") is FastWAMPolicy


def test_config_validates_features_model_ids_and_saved_auto_route(tmp_path):
    cfg = FastWAMConfig()
    cfg.save_pretrained(tmp_path)
    saved = json.loads((tmp_path / "config.json").read_text())

    assert saved["pretrained_path"] is None
    assert cfg.image_features["observation.images.image"].type == FeatureType.VISUAL
    assert cfg.action_feature.shape == (7,)
    assert cfg.robot_state_feature.shape == (8,)
    with pytest.raises(ValueError, match="image feature"):
        FastWAMConfig(input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))})
    with pytest.raises(ValueError, match="tokenizer_model_id"):
        FastWAMConfig(tokenizer_model_id="somebody/other-tokenizer")


def test_preprocessor_normalizes_images_and_postprocessor_toggles_actions(tmp_path):
    cfg = FastWAMConfig(
        action_dim=3,
        proprio_dim=2,
        action_horizon=4,
        n_action_steps=2,
        image_size=(2, 2),
        device="cpu",
        toggle_action_dimensions=[-1],
        input_features={
            "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 2, 2)),
            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
        },
        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
        base_model_id=None,
    )
    dataset_stats = {
        "observation.images.image": {
            "mean": torch.full((3, 1, 1), 0.2),
            "std": torch.full((3, 1, 1), 0.1),
        },
        OBS_STATE: {
            "mean": torch.tensor([1.0, 3.0]),
            "std": torch.tensor([2.0, 4.0]),
        },
        ACTION: {
            "mean": torch.zeros(3),
            "std": torch.ones(3),
        },
    }

    preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_stats)
    processed = preprocessor(
        {
            "observation.images.image": torch.tensor(
                [
                    [[0.0, 0.5], [1.0, 0.5]],
                    [[0.0, 0.5], [1.0, 0.5]],
                    [[0.0, 0.5], [1.0, 0.5]],
                ]
            ),
            OBS_STATE: torch.tensor([3.0, 7.0]),
        }
    )
    preprocessor.save_pretrained(tmp_path, config_filename="policy_preprocessor.json")
    postprocessor.save_pretrained(tmp_path, config_filename="policy_postprocessor.json")
    _, loaded_postprocessor = make_pre_post_processors(cfg, pretrained_path=str(tmp_path))

    expected_image = torch.tensor(
        [[[[-1.0, 0.0], [1.0, 0.0]], [[-1.0, 0.0], [1.0, 0.0]], [[-1.0, 0.0], [1.0, 0.0]]]]
    )
    assert preprocessor.name == "policy_preprocessor"
    assert postprocessor.name == "policy_postprocessor"
    assert torch.allclose(processed["observation.images.image"], expected_image)
    assert torch.allclose(processed[OBS_STATE], torch.tensor([[1.0, 1.0]]))
    assert torch.equal(dataset_stats["observation.images.image"]["mean"], torch.full((3, 1, 1), 0.2))
    assert any(isinstance(step, FastWAMActionToggleProcessorStep) for step in loaded_postprocessor.steps)
    assert torch.equal(
        loaded_postprocessor(torch.tensor([[0.25, 0.5, 1.0]])), torch.tensor([[0.25, 0.5, -1.0]])
    )


def test_policy_forward_and_predict_action_adapt_lerobot_batches(monkeypatch):
    captured = []

    class CapturingCore(FakeFastWAMCore):
        def infer_action(self, **kwargs):
            captured.append(
                {
                    "image_shape": tuple(kwargs["input_image"].shape),
                    "proprio_shape": tuple(kwargs["proprio"].shape),
                    "prompt": kwargs["prompt"],
                }
            )
            return {"action": torch.full((1, kwargs["action_horizon"], 3), float(len(captured)))}

    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: CapturingCore())
    cfg = FastWAMConfig(
        action_dim=3,
        proprio_dim=2,
        action_horizon=4,
        n_action_steps=2,
        image_size=(16, 16),
        input_features={
            "observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(2,)),
        },
        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(3,))},
        base_model_id=None,
    )
    with pytest.warns(RuntimeWarning, match="does not load pretrained FastWAM weights"):
        policy = FastWAMPolicy(cfg)

    output = policy.forward(
        {
            "observation.images.image": torch.zeros(1, 3, 16, 16),
            OBS_STATE: torch.zeros(1, 2),
            ACTION: torch.zeros(1, 4, 3),
            "context": torch.zeros(1, 5, 4096),
            "context_mask": torch.ones(1, 5, dtype=torch.bool),
        }
    )
    action = policy.predict_action_chunk(
        {
            "observation.images.image": torch.stack(
                [
                    torch.zeros(3, 16, 16),
                    torch.ones(3, 16, 16),
                ]
            ),
            OBS_STATE: torch.tensor([[0.0, 1.0], [2.0, 3.0]]),
            "task": ["task 0", "task 1"],
        }
    )

    assert output["loss"].item() == 1.0
    assert output["loss_action"].item() == 1.0
    assert action.shape == (2, 4, 3)
    assert action[:, 0, 0].tolist() == [1.0, 2.0]
    assert [item["image_shape"] for item in captured] == [(1, 3, 16, 16), (1, 3, 16, 16)]
    assert [item["proprio_shape"] for item in captured] == [(1, 2), (1, 2)]
    assert [item["prompt"] for item in captured] == [
        cfg.prompt_template.format(task="task 0"),
        cfg.prompt_template.format(task="task 1"),
    ]


def test_from_pretrained_loads_weights_without_initializing_wan_backbone(monkeypatch, tmp_path):
    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2, base_model_id=None)
    cfg.save_pretrained(tmp_path)
    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: FakeFastWAMCore())
    reference_policy = FastWAMPolicy(cfg, _suppress_base_init_warning=True)
    save_model(reference_policy, str(tmp_path / "model.safetensors"))

    def fail_if_wan_pretrained_is_loaded(*args, **kwargs):
        raise AssertionError("from_pretrained must not initialize or download Wan2.2 backbone components")

    monkeypatch.setattr(
        "lerobot.policies.fastwam.modular_fastwam.FastWAM.from_wan22_pretrained",
        fail_if_wan_pretrained_is_loaded,
    )
    monkeypatch.setattr(
        modeling_fastwam,
        "_build_core_model_from_architecture",
        lambda config: FakeFastWAMCore(),
        raising=False,
    )
    loaded_components_from = []
    monkeypatch.setattr(
        FastWAMPolicy,
        "load_wan_components_from_pretrained",
        lambda self, path: loaded_components_from.append(path),
    )

    policy = FastWAMPolicy.from_pretrained(tmp_path, strict=False)

    assert isinstance(policy.model, FakeFastWAMCore)
    assert loaded_components_from == [tmp_path]


def test_save_pretrained_copies_required_wan_sidecars(monkeypatch, tmp_path):
    cfg = FastWAMConfig(action_dim=3, proprio_dim=2, action_horizon=4, n_action_steps=2, base_model_id=None)
    source = tmp_path / "source"
    tokenizer = source / WAN_T5_TOKENIZER
    tokenizer.mkdir(parents=True)
    vae = source / WAN_VAE_CHECKPOINT
    text_encoder = source / WAN_T5_CHECKPOINT
    tokenizer_file = tokenizer / "tokenizer.json"
    vae.write_bytes(b"vae")
    text_encoder.write_bytes(b"text")
    tokenizer_file.write_text("{}")
    core = FakeFastWAMCore()
    core.model_paths = {
        "vae": str(vae),
        "text_encoder": str(text_encoder),
        "tokenizer": str(tokenizer),
    }
    monkeypatch.setattr(FastWAMPolicy, "_build_core_model", lambda self, config: core)
    policy = FastWAMPolicy(cfg, _suppress_base_init_warning=True)

    save_dir = tmp_path / "saved"
    policy.save_pretrained(save_dir)

    assert (save_dir / "model.safetensors").is_file()
    assert (save_dir / WAN_VAE_CHECKPOINT).read_bytes() == b"vae"
    assert (save_dir / WAN_T5_CHECKPOINT).read_bytes() == b"text"
    assert (save_dir / WAN_T5_TOKENIZER / "tokenizer.json").read_text() == "{}"


def test_wan_component_resolution_uses_fixed_safetensors_layout(tmp_path):
    tokenizer = tmp_path / WAN_T5_TOKENIZER
    tokenizer.mkdir(parents=True)
    (tmp_path / WAN_VAE_CHECKPOINT).touch()
    (tmp_path / WAN_T5_CHECKPOINT).touch()
    (tmp_path / "diffusion_pytorch_model-00001-of-00001.safetensors").touch()
    (tokenizer / "tokenizer.json").touch()

    paths = resolve_wan_checkpoint_paths(tmp_path)
    sidecar_paths = resolve_wan_component_paths(tmp_path)

    assert paths.dit == [tmp_path / "diffusion_pytorch_model-00001-of-00001.safetensors"]
    assert paths.vae == tmp_path / WAN_VAE_CHECKPOINT
    assert paths.text_encoder == tmp_path / WAN_T5_CHECKPOINT
    assert paths.tokenizer == tmp_path / WAN_T5_TOKENIZER
    assert sidecar_paths.dit == []
    assert WAN_DIT_PATTERN == "diffusion_pytorch_model*.safetensors"

    (tmp_path / WAN_T5_CHECKPOINT).unlink()
    with pytest.raises(FileNotFoundError, match="text encoder"):
        resolve_wan_checkpoint_paths(tmp_path)


def test_pretrained_config_round_trips_fastwam_features(tmp_path):
    cfg = FastWAMConfig(action_dim=7, proprio_dim=8, image_size=(224, 448), base_model_id=None)
    cfg.save_pretrained(tmp_path)

    loaded = PreTrainedConfig.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,)