mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 02:06:15 +00:00
refactor(rl): make RLAlgorithmConfig an abstract base class for better extensibility
This commit is contained in:
@@ -14,6 +14,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import abc
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
@@ -46,7 +47,7 @@ class TrainingStats:
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class RLAlgorithmConfig(draccus.ChoiceRegistry):
|
class RLAlgorithmConfig(draccus.ChoiceRegistry, abc.ABC):
|
||||||
"""Registry for algorithm configs."""
|
"""Registry for algorithm configs."""
|
||||||
|
|
||||||
@property
|
@property
|
||||||
@@ -57,6 +58,7 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry):
|
|||||||
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
|
raise TypeError(f"Expected string from get_choice_name, got {type(choice_name)}")
|
||||||
return choice_name
|
return choice_name
|
||||||
|
|
||||||
|
@abc.abstractmethod
|
||||||
def build_algorithm(self, policy: torch.nn.Module) -> RLAlgorithm:
|
def build_algorithm(self, policy: torch.nn.Module) -> RLAlgorithm:
|
||||||
"""Construct the :class:`RLAlgorithm` for this config.
|
"""Construct the :class:`RLAlgorithm` for this config.
|
||||||
|
|
||||||
@@ -65,6 +67,7 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry):
|
|||||||
raise NotImplementedError(f"{type(self).__name__} must implement build_algorithm()")
|
raise NotImplementedError(f"{type(self).__name__} must implement build_algorithm()")
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
|
@abc.abstractmethod
|
||||||
def from_policy_config(cls, policy_cfg: Any) -> RLAlgorithmConfig:
|
def from_policy_config(cls, policy_cfg: Any) -> RLAlgorithmConfig:
|
||||||
"""Build an algorithm config from a policy config.
|
"""Build an algorithm config from a policy config.
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user