From cbaf1acffdafd47525c3b13c0bb4b644da16be02 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Fri, 17 Apr 2026 15:07:01 +0200 Subject: [PATCH] feat(configs): add reward_model field to TrainPipelineConfig and Hub fields to RewardModelConfig --- src/lerobot/configs/rewards.py | 8 +++++ src/lerobot/configs/train.py | 48 +++++++++++++++++++-------- src/lerobot/policies/sarm/README.md | 1 - src/lerobot/policies/sarm/__init__.py | 18 ---------- 4 files changed, 43 insertions(+), 32 deletions(-) delete mode 120000 src/lerobot/policies/sarm/README.md delete mode 100644 src/lerobot/policies/sarm/__init__.py diff --git a/src/lerobot/configs/rewards.py b/src/lerobot/configs/rewards.py index 459d8787b..226db70b6 100644 --- a/src/lerobot/configs/rewards.py +++ b/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__) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 90547df18..f51cc211c 100644 --- a/src/lerobot/configs/train.py +++ b/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: - # Only load the policy config + reward_model_path = parser.get_path_arg("reward_model") + + 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.scheduler = self.policy.get_scheduler_preset() + self.optimizer = active_cfg.get_optimizer_preset() + self.scheduler = active_cfg.get_scheduler_preset() - if self.policy.push_to_hub and not self.policy.repo_id: - raise ValueError( - "'policy.repo_id' argument missing. Please specify it to push the model to the hub." - ) + if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id: + raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.") @classmethod def __get_path_fields__(cls) -> list[str]: diff --git a/src/lerobot/policies/sarm/README.md b/src/lerobot/policies/sarm/README.md deleted file mode 120000 index 18495860d..000000000 --- a/src/lerobot/policies/sarm/README.md +++ /dev/null @@ -1 +0,0 @@ -../../../../docs/source/policy_sarm_README.md \ No newline at end of file diff --git a/src/lerobot/policies/sarm/__init__.py b/src/lerobot/policies/sarm/__init__.py deleted file mode 100644 index b164c87ef..000000000 --- a/src/lerobot/policies/sarm/__init__.py +++ /dev/null @@ -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"]