From 6b518a3de7b50de0e54334f5bdf321c238303552 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Tue, 21 Apr 2026 14:58:56 +0200 Subject: [PATCH] refactor(tests): reward models --- tests/rewards/test_reward_model_base.py | 272 ++++++++++++++++++++++++ 1 file changed, 272 insertions(+) diff --git a/tests/rewards/test_reward_model_base.py b/tests/rewards/test_reward_model_base.py index 96557359f..c8755a0fa 100644 --- a/tests/rewards/test_reward_model_base.py +++ b/tests/rewards/test_reward_model_base.py @@ -14,13 +14,37 @@ """Tests for the reward model base classes and registry.""" +from dataclasses import dataclass +from pathlib import Path +from types import SimpleNamespace + import pytest import torch from lerobot.configs.rewards import RewardModelConfig +from lerobot.optim.optimizers import AdamWConfig from lerobot.rewards.pretrained import PreTrainedRewardModel +@RewardModelConfig.register_subclass(name="_dummy_hub_reward") +@dataclass +class _DummyHubRewardConfig(RewardModelConfig): + def get_optimizer_preset(self): + return AdamWConfig(lr=1e-4) + + +class _DummyHubReward(PreTrainedRewardModel): + config_class = _DummyHubRewardConfig + name = "_dummy_hub_reward" + + def __init__(self, config): + super().__init__(config) + self.bias = torch.nn.Parameter(torch.zeros(1)) + + def compute_reward(self, batch): + return self.bias.expand(1) + + def test_reward_model_config_registry(): """Verify that classifier and sarm are registered.""" known = RewardModelConfig.get_known_choices() @@ -99,3 +123,251 @@ def test_non_trainable_forward_raises(): with pytest.raises(NotImplementedError, match="not trainable"): model.forward({"x": torch.zeros(1)}) + + +# --------------------------------------------------------------------------- +# Trainable vs zero-shot (general-purpose) reward models. +# The proposal explicitly supports models like TOPReward that wrap a pretrained +# VLM and produce a reward signal without any training step. These tests pin +# the contract that lets such models coexist with trainable ones. +# --------------------------------------------------------------------------- + + +def test_is_trainable_false_when_forward_not_overridden(): + """A reward model that only implements ``compute_reward`` is zero-shot.""" + model, _ = _make_dummy_reward_model() + assert model.is_trainable is False + + +def test_is_trainable_true_when_forward_overridden(): + """Overriding ``forward`` flips ``is_trainable`` to True.""" + + class _TrainableReward(_DummyHubReward): + name = "_trainable_dummy_reward" + + def forward(self, batch): + loss = (self.bias**2).sum() + return loss, {} + + # Register a fresh config subclass so the subclass check passes. + @RewardModelConfig.register_subclass(name="_trainable_dummy_reward") + @dataclass + class _TrainableConfig(_DummyHubRewardConfig): + pass + + _TrainableReward.config_class = _TrainableConfig + model = _TrainableReward(_TrainableConfig()) + assert model.is_trainable is True + + +# --------------------------------------------------------------------------- +# RewardModelConfig.from_pretrained +# --------------------------------------------------------------------------- + + +def test_reward_model_config_from_pretrained_raises_when_config_missing(tmp_path): + """``from_pretrained`` must surface a clear ``FileNotFoundError`` when the + target directory exists but does not contain ``config.json``, instead of + crashing later inside ``draccus.parse``. + """ + # tmp_path exists but has no config.json + with pytest.raises(FileNotFoundError, match="config.json not found"): + RewardModelConfig.from_pretrained(tmp_path) + + +def test_reward_model_config_from_pretrained_roundtrip(tmp_path): + """Round-trip: save a RewardClassifierConfig, reload it, fields must match.""" + from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig + + original = RewardClassifierConfig( + num_classes=3, + hidden_dim=128, + latent_dim=64, + num_cameras=1, + learning_rate=5e-4, + ) + original._save_pretrained(tmp_path) + + loaded = RewardModelConfig.from_pretrained(tmp_path) + + assert isinstance(loaded, RewardClassifierConfig) + assert loaded.num_classes == 3 + assert loaded.hidden_dim == 128 + assert loaded.latent_dim == 64 + assert loaded.num_cameras == 1 + assert loaded.learning_rate == 5e-4 + + +# --------------------------------------------------------------------------- +# TrainPipelineConfig — reward model training path +# --------------------------------------------------------------------------- + + +def test_train_pipeline_config_path_fields_includes_reward_model(): + """``--reward_model.path=local/dir`` requires ``reward_model`` to be listed + as a draccus path-field on ``TrainPipelineConfig``.""" + from lerobot.configs.train import TrainPipelineConfig + + fields = TrainPipelineConfig.__get_path_fields__() + assert "policy" in fields + assert "reward_model" in fields + + +def test_train_pipeline_config_trainable_config_returns_reward_model_when_set(): + """When only ``reward_model`` is set, ``trainable_config`` (used by the + trainer for e.g. ``.device``) must return it — not ``None`` from ``policy``.""" + from lerobot.configs.default import DatasetConfig + from lerobot.configs.train import TrainPipelineConfig + from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig + + reward_cfg = RewardClassifierConfig(device="cpu") + cfg = TrainPipelineConfig( + dataset=DatasetConfig(repo_id="user/repo"), + reward_model=reward_cfg, + ) + + assert cfg.is_reward_model_training is True + assert cfg.trainable_config is reward_cfg + # This is what lerobot_train.py uses to decide force_cpu; ``cfg.policy.device`` + # would AttributeError here because policy is None. + assert cfg.trainable_config.device == "cpu" + + +def test_train_pipeline_config_trainable_config_returns_policy_when_set(): + """Mirror of the reward-model case: when only ``policy`` is set, + ``trainable_config`` must return it.""" + from lerobot.configs.default import DatasetConfig + from lerobot.configs.train import TrainPipelineConfig + from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig + + policy_cfg = DiffusionConfig(device="cpu") + cfg = TrainPipelineConfig( + dataset=DatasetConfig(repo_id="user/repo"), + policy=policy_cfg, + ) + + assert cfg.is_reward_model_training is False + assert cfg.trainable_config is policy_cfg + assert cfg.trainable_config.device == "cpu" + + +# --------------------------------------------------------------------------- +# PreTrainedRewardModel hub upload: push_model_to_hub + generate_model_card. +# We test the generation side (offline) fully, and the upload side with HfApi +# mocked so nothing actually hits the network. +# --------------------------------------------------------------------------- + + +def _make_dummy_reward_model(**config_kwargs): + return _DummyHubReward(_DummyHubRewardConfig(**config_kwargs)), _DummyHubRewardConfig + + +@pytest.fixture +def _offline_model_card(monkeypatch): + """``ModelCard.validate`` does a live ``POST`` to huggingface.co — bypass it + so tests can run offline.""" + from huggingface_hub import ModelCard + + monkeypatch.setattr(ModelCard, "validate", lambda self, *a, **kw: None) + + +def test_reward_model_generate_model_card_renders_expected_fields(_offline_model_card): + """``generate_model_card`` must produce a card with the right metadata and + body, using the dedicated reward-model template.""" + model, _ = _make_dummy_reward_model( + license="mit", + tags=["robot", "sim"], + ) + + card = model.generate_model_card( + dataset_repo_id="user/my_dataset", + model_type=model.config.type, + license=model.config.license, + tags=model.config.tags, + ) + + # Metadata (YAML header) — ModelCardData fields. + assert card.data.license == "mit" + assert card.data.library_name == "lerobot" + assert card.data.pipeline_tag == "robotics" + assert "reward-model" in card.data.tags + assert model.config.type in card.data.tags + assert card.data.model_name == model.config.type + assert card.data.datasets == "user/my_dataset" + + # Body — specific to the reward-model template, NOT the policy one. + body = str(card) + assert "Reward Model Card" in body + assert "This reward model has been trained" in body + assert "--reward_model.type=" in body # reward-model-specific usage block + + +def test_reward_model_generate_model_card_uses_default_license(_offline_model_card): + """When config.license is None the card falls back to apache-2.0.""" + model, _ = _make_dummy_reward_model() + + card = model.generate_model_card( + dataset_repo_id="user/my_dataset", + model_type=model.config.type, + license=model.config.license, + tags=None, + ) + + assert card.data.license == "apache-2.0" + + +def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _offline_model_card): + """``push_model_to_hub`` must: + 1. create the repo, + 2. assemble a temp folder with weights + config.json + train_config.json + README.md, + 3. call ``api.upload_folder`` on that folder. + All network calls are mocked. + """ + from huggingface_hub.constants import CONFIG_NAME + + from lerobot.configs.default import DatasetConfig + from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig + + model, _ = _make_dummy_reward_model( + repo_id="user/my_reward", + license="apache-2.0", + ) + # Point the reward model's train config at a dummy dataset repo. + train_cfg = TrainPipelineConfig( + dataset=DatasetConfig(repo_id="user/my_dataset"), + reward_model=model.config, + ) + + uploaded: dict = {} + fake_commit_info = SimpleNamespace(repo_url=SimpleNamespace(url="https://huggingface.co/user/my_reward")) + + class _FakeHfApi: + def create_repo(self, repo_id, private=None, exist_ok=False): + uploaded["create_repo_id"] = repo_id + uploaded["create_private"] = private + return SimpleNamespace(repo_id=repo_id) + + def upload_folder(self, *, repo_id, repo_type, folder_path, commit_message, **_kwargs): + uploaded["upload_repo_id"] = repo_id + uploaded["upload_repo_type"] = repo_type + uploaded["commit_message"] = commit_message + # Snapshot files assembled in the temp folder — this is the real + # contract we care about. + uploaded["files"] = sorted(p.name for p in Path(folder_path).iterdir()) + return fake_commit_info + + from lerobot.rewards import pretrained as reward_pretrained + + monkeypatch.setattr(reward_pretrained, "HfApi", lambda *a, **kw: _FakeHfApi()) + + model.push_model_to_hub(train_cfg) + + assert uploaded["create_repo_id"] == "user/my_reward" + assert uploaded["upload_repo_id"] == "user/my_reward" + assert uploaded["upload_repo_type"] == "model" + assert uploaded["commit_message"] == "Upload reward model weights, train config and readme" + # Minimum required files that must be uploaded with a reward model. + assert CONFIG_NAME in uploaded["files"] # config.json + assert TRAIN_CONFIG_NAME in uploaded["files"] # train_config.json + assert "README.md" in uploaded["files"] + assert any(name.endswith(".safetensors") for name in uploaded["files"])