refactor(pretrained_rm): add model card template

src/lerobot/rewards/pretrained.py

@@ -16,12 +16,14 @@ import abc
 import builtins
 import logging
 import os
+from importlib.resources import files
 from pathlib import Path
-from typing import Any, TypeVar
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any, TypeVar
 
 import packaging
 import safetensors
-from huggingface_hub import hf_hub_download
+from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
@@ -30,7 +32,8 @@ from torch import Tensor, nn
 from lerobot.configs.rewards import RewardModelConfig
 from lerobot.utils.hub import HubMixin
 
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from lerobot.configs.train import TrainPipelineConfig
 
 T = TypeVar("T", bound="PreTrainedRewardModel")
 
@@ -61,7 +64,7 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
     def _save_pretrained(self, save_directory: Path) -> None:
         self.config._save_pretrained(save_directory)
         model_to_save = self.module if hasattr(self, "module") else self
-        save_model_as_safetensor(model_to_save, str(Path(save_directory) / SAFETENSORS_SINGLE_FILE))
+        save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
 
     @classmethod
     def from_pretrained(
@@ -79,6 +82,10 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
         strict: bool = False,
         **kwargs,
     ) -> T:
+        """
+        The reward model is set in evaluation mode by default using `reward.eval()` (dropout modules are
+        deactivated). To train it, you should first set it back in training mode with `reward.train()`.
+        """
         if config is None:
             config = RewardModelConfig.from_pretrained(
                 pretrained_name_or_path=pretrained_name_or_path,
@@ -94,7 +101,7 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
         model_id = str(pretrained_name_or_path)
         instance = cls(config, **kwargs)
         if os.path.isdir(model_id):
-            logger.info("Loading reward model weights from local directory")
+            print("Loading weights from local directory")
             model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
             reward = cls._load_as_safetensor(instance, model_file, config.device or "cpu", strict)
         else:
@@ -123,7 +130,7 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
         # Create base kwargs
-        kwargs: dict[str, Any] = {"strict": strict}
+        kwargs = {"strict": strict}
 
         # Add device parameter for newer versions that support it
         if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
@@ -131,11 +138,10 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
 
         # Load the model with appropriate kwargs
         missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
-
         if missing_keys:
-            logger.warning(f"Missing keys when loading reward model: {missing_keys}")
+            logging.warning(f"Missing key(s) when loading model: {missing_keys}")
         if unexpected_keys:
-            logger.warning(f"Unexpected keys when loading reward model: {unexpected_keys}")
+            logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
 
         # For older versions, manually move to device if needed
         if "device" not in kwargs and map_location != "cpu":
@@ -148,16 +154,16 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
             model.to(map_location)
         return model
 
-    def reset(self) -> None:
-        """Reset any internal state."""
-        pass
-
     def get_optim_params(self):
         """
         Returns the reward-model-specific parameters dict to be passed on to the optimizer.
         """
         return self.parameters()
 
+    def reset(self) -> None:
+        """Reset any internal state."""
+        pass
+
     @abc.abstractmethod
     def compute_reward(self, batch: dict[str, Tensor]) -> Tensor:
         """Compute a scalar reward signal for a batch of observations.
@@ -176,3 +182,63 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
         raise NotImplementedError(
             f"{self.__class__.__name__} is not trainable. Only use compute_reward() for inference."
         )
+
+    @property
+    def is_trainable(self) -> bool:
+        """Whether this reward model can be trained via ``lerobot-train``.
+
+        Trainable reward models override :meth:`forward`; zero-shot models
+        inherit the base implementation that raises ``NotImplementedError``.
+        """
+        return type(self).forward is not PreTrainedRewardModel.forward
+
+    def push_model_to_hub(self, cfg: "TrainPipelineConfig"):
+        api = HfApi()
+        repo_id = api.create_repo(
+            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
+        ).repo_id
+
+        # Push the files to the repo in a single commit
+        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+            saved_path = Path(tmp) / repo_id
+
+            self.save_pretrained(saved_path)  # Calls _save_pretrained and stores model tensors
+
+            card = self.generate_model_card(
+                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
+            )
+            card.save(str(saved_path / "README.md"))
+
+            cfg.save_pretrained(saved_path)  # Calls _save_pretrained and stores train config
+
+            commit_info = api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message="Upload reward model weights, train config and readme",
+                allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
+                ignore_patterns=["*.tmp", "*.log"],
+            )
+
+            logging.info(f"Model pushed to {commit_info.repo_url.url}")
+
+    def generate_model_card(
+        self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None
+    ) -> ModelCard:
+        card_data = ModelCardData(
+            license=license or "apache-2.0",
+            library_name="lerobot",
+            pipeline_tag="robotics",
+            tags=list(set(tags or []).union({"robotics", "lerobot", "reward-model", model_type})),
+            model_name=model_type,
+            datasets=dataset_repo_id,
+        )
+
+        template_card = (
+            files("lerobot.templates")
+            .joinpath("lerobot_rewardmodel_modelcard_template.md")
+            .read_text(encoding="utf-8")
+        )
+        card = ModelCard.from_template(card_data, template_str=template_card)
+        card.validate()
+        return card

src/lerobot/templates/lerobot_rewardmodel_modelcard_template.md

@@ -0,0 +1,55 @@
+---
+# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+# Doc / guide: https://huggingface.co/docs/hub/model-cards
+# prettier-ignore
+{{card_data}}
+---
+
+# Reward Model Card for {{ model_name | default("Reward Model ID", true) }}
+
+<!-- Provide a quick summary of what the reward model is/does. -->
+
+{% if model_name == "reward_classifier" %}
+A reward classifier is a lightweight neural network that scores observations or trajectories for task success, providing a learned reward signal or offline evaluation when explicit rewards are unavailable.
+{% elif model_name == "sarm" %}
+A Success-Aware Reward Model (SARM) predicts a dense reward signal from observations, typically used downstream for reinforcement learning or human-in-the-loop fine-tuning when task success is not directly observable.
+{% else %}
+_Reward model type not recognized — please update this template._
+{% endif %}
+
+This reward model has been trained and pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot).
+See the full documentation at [LeRobot Docs](https://huggingface.co/docs/lerobot/index).
+
+---
+
+## How to Get Started with the Reward Model
+
+### Train from scratch
+
+```bash
+lerobot-train \
+  --dataset.repo_id=${HF_USER}/<dataset> \
+  --reward_model.type={{ model_name | default("reward_classifier", true) }} \
+  --output_dir=outputs/train/<desired_reward_model_repo_id> \
+  --job_name=lerobot_reward_training \
+  --reward_model.device=cuda \
+  --reward_model.repo_id=${HF_USER}/<desired_reward_model_repo_id> \
+  --wandb.enable=true
+```
+
+_Writes checkpoints to `outputs/train/<desired_reward_model_repo_id>/checkpoints/`._
+
+### Load the reward model in Python
+
+```python
+from lerobot.rewards import make_reward_model
+
+reward_model = make_reward_model(pretrained_path="<hf_user>/<reward_model_repo_id>")
+reward = reward_model.compute_reward(batch)
+```
+
+---
+
+## Model Details
+
+- **License:** {{ license | default("\[More Information Needed]", true) }}