feat(configs): add reward_model field to TrainPipelineConfig and Hub fields to RewardModelConfig

src/lerobot/configs/rewards.py

@@ -55,6 +55,14 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
 
     pretrained_path: str | None = None
 
+    push_to_hub: bool = False
+    repo_id: str | None = None
+
+    # Hub metadata
+    license: str | None = None
+    tags: list[str] | None = None
+    private: bool | None = None
+
     @property
     def type(self) -> str:
         choice_name = self.get_choice_name(self.__class__)

src/lerobot/configs/train.py

@@ -30,6 +30,7 @@ from lerobot.utils.sample_weighting import SampleWeightingConfig
 
 from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
+from .rewards import RewardModelConfig
 
 TRAIN_CONFIG_NAME = "train_config.json"
 
@@ -39,6 +40,7 @@ class TrainPipelineConfig(HubMixin):
     dataset: DatasetConfig
     env: envs.EnvConfig | None = None
     policy: PreTrainedConfig | None = None
+    reward_model: RewardModelConfig | None = None
     # Set `dir` to where you would like to save all of the run outputs. If you run another training session
     # with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
     output_dir: Path | None = None
@@ -78,16 +80,34 @@ class TrainPipelineConfig(HubMixin):
     rename_map: dict[str, str] = field(default_factory=dict)
     checkpoint_path: Path | None = field(init=False, default=None)
 
+    @property
+    def is_reward_model_training(self) -> bool:
+        """True when the config targets a reward model rather than a policy."""
+        return self.reward_model is not None
+
+    @property
+    def trainable_config(self) -> PreTrainedConfig | RewardModelConfig:
+        """Return whichever config (policy or reward_model) is active."""
+        if self.is_reward_model_training:
+            return self.reward_model  # type: ignore[return-value]
+        return self.policy  # type: ignore[return-value]
+
     def validate(self) -> None:
         # HACK: We parse again the cli args here to get the pretrained paths if there was some.
         policy_path = parser.get_path_arg("policy")
-        if policy_path:
+        reward_model_path = parser.get_path_arg("reward_model")
-            # Only load the policy config
+
+        if reward_model_path:
+            cli_overrides = parser.get_cli_overrides("reward_model")
+            self.reward_model = RewardModelConfig.from_pretrained(
+                reward_model_path, cli_overrides=cli_overrides
+            )
+            self.reward_model.pretrained_path = str(Path(reward_model_path))
+        elif policy_path:
             cli_overrides = parser.get_cli_overrides("policy")
             self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
             self.policy.pretrained_path = Path(policy_path)
         elif self.resume:
-            # The entire train config is already loaded, we just need to get the checkpoint dir
             config_path = parser.parse_arg("config_path")
             if not config_path:
                 raise ValueError(
@@ -103,18 +123,22 @@ class TrainPipelineConfig(HubMixin):
             policy_dir = Path(config_path).parent
             if self.policy is not None:
                 self.policy.pretrained_path = policy_dir
+            if self.reward_model is not None:
+                self.reward_model.pretrained_path = str(policy_dir)
             self.checkpoint_path = policy_dir.parent
 
-        if self.policy is None:
+        if self.policy is None and self.reward_model is None:
             raise ValueError(
-                "Policy is not configured. Please specify a pretrained policy with `--policy.path`."
+                "Neither policy nor reward_model is configured. "
+                "Please specify one with `--policy.path` or `--reward_model.path`."
             )
 
+        active_cfg = self.trainable_config
         if not self.job_name:
             if self.env is None:
-                self.job_name = f"{self.policy.type}"
+                self.job_name = f"{active_cfg.type}"
             else:
-                self.job_name = f"{self.env.type}_{self.policy.type}"
+                self.job_name = f"{self.env.type}_{active_cfg.type}"
 
         if not self.resume and isinstance(self.output_dir, Path) and self.output_dir.is_dir():
             raise FileExistsError(
@@ -132,13 +156,11 @@ class TrainPipelineConfig(HubMixin):
         if not self.use_policy_training_preset and (self.optimizer is None or self.scheduler is None):
             raise ValueError("Optimizer and Scheduler must be set when the policy presets are not used.")
         elif self.use_policy_training_preset and not self.resume:
-            self.optimizer = self.policy.get_optimizer_preset()
+            self.optimizer = active_cfg.get_optimizer_preset()
-            self.scheduler = self.policy.get_scheduler_preset()
+            self.scheduler = active_cfg.get_scheduler_preset()
 
-        if self.policy.push_to_hub and not self.policy.repo_id:
+        if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id:
-            raise ValueError(
+            raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.")
-                "'policy.repo_id' argument missing. Please specify it to push the model to the hub."
-            )
 
     @classmethod
     def __get_path_fields__(cls) -> list[str]:

src/lerobot/policies/sarm/README.md

@@ -1 +0,0 @@
-../../../../docs/source/policy_sarm_README.md

src/lerobot/policies/sarm/__init__.py

@@ -1,18 +0,0 @@
-# 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.
-
-from .configuration_sarm import SARMConfig
-from .modeling_sarm import SARMRewardModel
-
-__all__ = ["SARMConfig", "SARMRewardModel"]