| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| from __future__ import annotations |
|
|
| import abc |
| import builtins |
| import os |
| from collections.abc import Iterator |
| from pathlib import Path |
| from typing import TYPE_CHECKING, Any, TypeVar |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
| from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE |
| from huggingface_hub.errors import HfHubHTTPError |
| from safetensors.torch import load_file as load_safetensors, save_file as save_safetensors |
| from torch.optim import Optimizer |
|
|
| from lerobot.types import BatchType |
| from lerobot.utils.hub import HubMixin |
|
|
| from .configs import RLAlgorithmConfig, TrainingStats |
|
|
| if TYPE_CHECKING: |
| from torch import nn |
|
|
| from ..data_sources.data_mixer import DataMixer |
|
|
| T = TypeVar("T", bound="RLAlgorithm") |
|
|
|
|
| class RLAlgorithm(HubMixin, abc.ABC): |
| """Base for all RL algorithms.""" |
|
|
| config_class: type[RLAlgorithmConfig] |
| name: str |
| config: RLAlgorithmConfig |
|
|
| @abc.abstractmethod |
| def update(self, batch_iterator: Iterator[BatchType]) -> TrainingStats: |
| """One complete training step. |
| |
| The algorithm calls ``next(batch_iterator)`` as many times as it |
| needs (e.g. ``utd_ratio`` times for SAC) to obtain fresh batches. |
| The iterator is owned by the trainer; the algorithm just consumes |
| from it. |
| """ |
| raise NotImplementedError |
|
|
| def configure_data_iterator( |
| self, |
| data_mixer: DataMixer, |
| batch_size: int, |
| *, |
| async_prefetch: bool = True, |
| queue_size: int = 2, |
| ) -> Iterator[BatchType]: |
| """Create the data iterator this algorithm needs. |
| |
| The default implementation uses the standard ``data_mixer.get_iterator()``. |
| Algorithms that need specialised sampling should override this method. |
| """ |
| return data_mixer.get_iterator( |
| batch_size=batch_size, |
| async_prefetch=async_prefetch, |
| queue_size=queue_size, |
| ) |
|
|
| @abc.abstractmethod |
| def make_optimizers_and_scheduler(self) -> dict[str, Optimizer]: |
| """Build and return the optimizers used during training. |
| |
| Called once on the learner side after construction. |
| """ |
| raise NotImplementedError |
|
|
| def get_optimizers(self) -> dict[str, Optimizer]: |
| """Return optimizers for checkpointing / external scheduling.""" |
| return {} |
|
|
| @property |
| def optimization_step(self) -> int: |
| """Current learner optimization step. |
| |
| Part of the stable contract for checkpoint/resume. Algorithms can |
| either use this default storage or override for custom behavior. |
| """ |
| return getattr(self, "_optimization_step", 0) |
|
|
| @optimization_step.setter |
| def optimization_step(self, value: int) -> None: |
| self._optimization_step = int(value) |
|
|
| def get_weights(self) -> dict[str, Any]: |
| """Policy state-dict to push to actors.""" |
| return {} |
|
|
| @abc.abstractmethod |
| def load_weights(self, weights: dict[str, Any], device: str | torch.device = "cpu") -> None: |
| """Load policy state-dict received from the learner.""" |
| raise NotImplementedError |
|
|
| @abc.abstractmethod |
| def state_dict(self) -> dict[str, torch.Tensor]: |
| """Algorithm-owned trainable tensors. |
| |
| Must return a flat tensor mapping for everything the algorithm owns |
| that is not part of the policy (e.g. critic ensembles, target networks, |
| temperature parameters). Algorithms with no training-only tensors |
| should explicitly return an empty dict. |
| """ |
| raise NotImplementedError |
|
|
| @abc.abstractmethod |
| def load_state_dict( |
| self, |
| state_dict: dict[str, torch.Tensor], |
| device: str | torch.device = "cpu", |
| ) -> None: |
| """In-place load of algorithm-owned tensors. |
| |
| Implementations MUST keep the identity of any ``nn.Parameter`` that an |
| optimizer references (e.g. SAC's ``log_alpha``) by using ``.copy_()`` |
| rather than rebinding the attribute. |
| """ |
| raise NotImplementedError |
|
|
| def _save_pretrained(self, save_directory: Path) -> None: |
| """Persist the algorithm's tensors and config to ``save_directory``. |
| |
| Writes ``model.safetensors`` (algorithm tensors via :meth:`state_dict`) |
| and ``config.json`` (via :meth:`RLAlgorithmConfig.save_pretrained`). |
| """ |
| tensors = {k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()} |
| save_safetensors(tensors, str(save_directory / SAFETENSORS_SINGLE_FILE)) |
| self.config._save_pretrained(save_directory) |
|
|
| @classmethod |
| def from_pretrained( |
| cls: builtins.type[T], |
| pretrained_name_or_path: str | Path, |
| *, |
| policy: nn.Module, |
| config: RLAlgorithmConfig | None = None, |
| force_download: bool = False, |
| resume_download: bool | None = None, |
| proxies: dict | None = None, |
| token: str | bool | None = None, |
| cache_dir: str | Path | None = None, |
| local_files_only: bool = False, |
| revision: str | None = None, |
| device: str | torch.device = "cpu", |
| **algo_kwargs: Any, |
| ) -> T: |
| """Build an algorithm and load its weights from ``pretrained_name_or_path``.""" |
| if config is None: |
| config = cls.config_class.from_pretrained( |
| pretrained_name_or_path, |
| force_download=force_download, |
| resume_download=resume_download, |
| proxies=proxies, |
| token=token, |
| cache_dir=cache_dir, |
| local_files_only=local_files_only, |
| revision=revision, |
| ) |
| if hasattr(config, "policy_config"): |
| config.policy_config = policy.config |
|
|
| instance = cls(policy=policy, config=config, **algo_kwargs) |
|
|
| model_id = str(pretrained_name_or_path) |
| if os.path.isdir(model_id): |
| model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE) |
| else: |
| try: |
| model_file = hf_hub_download( |
| repo_id=model_id, |
| filename=SAFETENSORS_SINGLE_FILE, |
| revision=revision, |
| cache_dir=cache_dir, |
| force_download=force_download, |
| proxies=proxies, |
| resume_download=resume_download, |
| token=token, |
| local_files_only=local_files_only, |
| ) |
| except HfHubHTTPError as e: |
| raise FileNotFoundError( |
| f"{SAFETENSORS_SINGLE_FILE} not found on the HuggingFace Hub in {model_id}" |
| ) from e |
|
|
| tensors = load_safetensors(model_file) |
| instance.load_state_dict(tensors, device=device) |
| return instance |
|
|
