From 8414188db0b178b947985a7a9a91314708837315 Mon Sep 17 00:00:00 2001 From: Caroline Pascal Date: Tue, 30 Jun 2026 20:21:06 +0200 Subject: [PATCH] fix(datasets dependency): removing datasets dependency in pretrained.py (#3897) --- src/lerobot/policies/pretrained.py | 56 ++++++++++++++-------------- src/lerobot/scripts/lerobot_train.py | 4 +- 2 files changed, 29 insertions(+), 31 deletions(-) diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index aea5f1b08..702569b8c 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -11,6 +11,8 @@ # 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 __future__ import annotations + import abc import builtins import dataclasses @@ -19,7 +21,7 @@ import os from importlib.resources import files from pathlib import Path from tempfile import TemporaryDirectory -from typing import TypedDict, TypeVar, Unpack +from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack import packaging import safetensors @@ -38,10 +40,13 @@ from .utils import log_model_loading_keys T = TypeVar("T", bound="PreTrainedPolicy") +if TYPE_CHECKING: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + def _build_card_context( cfg: TrainPipelineConfig | None, - dataset_repo_id: str | None, + dataset_meta: LeRobotDatasetMetadata | None, input_features: dict | None, output_features: dict | None, ) -> dict: @@ -72,30 +77,16 @@ def _build_card_context( "lerobot_version": __version__, } - if dataset_repo_id: - dataset_cfg = getattr(cfg, "dataset", None) - try: - from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata - - meta = LeRobotDatasetMetadata( - dataset_repo_id, - root=getattr(dataset_cfg, "root", None), - revision=getattr(dataset_cfg, "revision", None), - ) - context["dataset"] = { - "repo_id": dataset_repo_id, - "episodes": meta.total_episodes, - "frames": meta.total_frames, - "fps": meta.fps, - "tasks": [str(task) for task in meta.tasks.index], - } - context["robot_type"] = meta.robot_type - context["cameras"] = [key.split(".")[-1] for key in meta.camera_keys] - except Exception as e: # noqa: BLE001 — dataset details are optional, never fail the push - logging.warning( - f"Could not load dataset metadata for '{dataset_repo_id}'; those sections will be " - f"omitted from the model card. ({e})" - ) + if dataset_meta is not None: + context["dataset"] = { + "repo_id": dataset_meta.repo_id, + "episodes": dataset_meta.total_episodes, + "frames": dataset_meta.total_frames, + "fps": dataset_meta.fps, + "tasks": [str(task) for task in dataset_meta.tasks.index], + } + context["robot_type"] = dataset_meta.robot_type + context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys] return context @@ -304,6 +295,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): cfg: TrainPipelineConfig, peft_model=None, state_dict: dict[str, Tensor] | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -325,7 +317,12 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( - cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg + cfg.dataset.repo_id, + self.config.type, + self.config.license, + self.config.tags, + cfg=cfg, + dataset_meta=dataset_meta, ) card.save(str(saved_path / "README.md")) @@ -352,6 +349,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): license: str | None, tags: list[str] | None, cfg: TrainPipelineConfig | None = None, + dataset_meta: LeRobotDatasetMetadata | None = None, ) -> ModelCard: base_model_mapping = { "smolvla": "lerobot/smolvla_base", @@ -372,7 +370,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): ) context = _build_card_context( - cfg, dataset_repo_id, self.config.input_features, self.config.output_features + cfg, dataset_meta, self.config.input_features, self.config.output_features ) # Used by the template to pre-fill commands and the "Fine-tuned from" line. context["policy_repo_id"] = getattr(self.config, "repo_id", None) @@ -389,7 +387,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, peft_config=None, peft_cli_overrides: dict | None = None, - ) -> "PreTrainedPolicy": + ) -> PreTrainedPolicy: """ Wrap this policy with PEFT adapters for parameter-efficient fine-tuning. diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index f2a152df9..44c94a1eb 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -736,9 +736,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): unwrapped_model = accelerator.unwrap_model(policy) # PEFT only applies when training a policy — reward models use the plain path. if not cfg.is_reward_model_training and cfg.policy.use_peft: - unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) + unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta) else: - unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id)