refactor(tests): reward models

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"])