From f3723fe6e092c3ad9e45a1455c0d62a0647562cb Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Tue, 21 Apr 2026 14:44:26 +0200 Subject: [PATCH] refactor(pretrained_rm): add model card template --- src/lerobot/rewards/pretrained.py | 92 ++++++++++++++++--- .../lerobot_rewardmodel_modelcard_template.md | 55 +++++++++++ 2 files changed, 134 insertions(+), 13 deletions(-) create mode 100644 src/lerobot/templates/lerobot_rewardmodel_modelcard_template.md diff --git a/src/lerobot/rewards/pretrained.py b/src/lerobot/rewards/pretrained.py index 7cd112ba2..d44b31733 100644 --- a/src/lerobot/rewards/pretrained.py +++ b/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 diff --git a/src/lerobot/templates/lerobot_rewardmodel_modelcard_template.md b/src/lerobot/templates/lerobot_rewardmodel_modelcard_template.md new file mode 100644 index 000000000..933bf7586 --- /dev/null +++ b/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) }} + + + +{% 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}/ \ + --reward_model.type={{ model_name | default("reward_classifier", true) }} \ + --output_dir=outputs/train/ \ + --job_name=lerobot_reward_training \ + --reward_model.device=cuda \ + --reward_model.repo_id=${HF_USER}/ \ + --wandb.enable=true +``` + +_Writes checkpoints to `outputs/train//checkpoints/`._ + +### Load the reward model in Python + +```python +from lerobot.rewards import make_reward_model + +reward_model = make_reward_model(pretrained_path="/") +reward = reward_model.compute_reward(batch) +``` + +--- + +## Model Details + +- **License:** {{ license | default("\[More Information Needed]", true) }}