162 lines
4.7 KiB
Python
162 lines
4.7 KiB
Python
from __future__ import annotations
|
|
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Transition:
|
|
"""A minimal transition container for RL.
|
|
|
|
This is intentionally generic because the underlying env state type may be a
|
|
JAX pytree, a numpy struct, or something library-specific.
|
|
"""
|
|
|
|
obs: Any
|
|
action: Any
|
|
reward: float
|
|
next_obs: Any
|
|
terminated: bool
|
|
truncated: bool
|
|
info: dict[str, Any] | None = None
|
|
|
|
|
|
class RLAlgorithm(ABC):
|
|
"""Insertable RL algorithm interface."""
|
|
|
|
@abstractmethod
|
|
def select_action(self, *, obs: Any, rng: Any | None = None) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def observe(self, transition: Transition) -> None:
|
|
"""Optional hook to store transitions."""
|
|
|
|
def update(self, *, rng: Any | None = None) -> dict[str, float]:
|
|
"""Optional hook to run one training update."""
|
|
|
|
return {}
|
|
|
|
def save(self, path: str) -> None:
|
|
raise NotImplementedError("Save not implemented")
|
|
|
|
def load(self, path: str) -> None:
|
|
raise NotImplementedError("Load not implemented")
|
|
|
|
|
|
_RL_MODEL_REGISTRY: dict[str, type["RLModel"]] = {}
|
|
|
|
|
|
def registered_model_types() -> list[str]:
|
|
return sorted(_RL_MODEL_REGISTRY)
|
|
|
|
|
|
def create_model(type_name: str, *, payload: dict[str, Any]) -> "RLModel":
|
|
model_cls = _RL_MODEL_REGISTRY.get(type_name)
|
|
if model_cls is None:
|
|
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
|
|
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
|
|
return model_cls.from_payload(payload)
|
|
|
|
|
|
def get_rl_model_registry() -> dict[str, type["RLModel"]]:
|
|
"""Return a copy of the current RLModel registry.
|
|
|
|
The registry is populated by importing concrete model modules that use the
|
|
`@register_rl_model(...)` decorator.
|
|
"""
|
|
|
|
return dict(_RL_MODEL_REGISTRY)
|
|
|
|
|
|
def register_rl_model(*type_names: str):
|
|
"""Decorator to register an `RLModel` for generic loading.
|
|
|
|
Concrete model modules should apply this decorator, so `base.py` never needs
|
|
to import concrete models (avoids circular imports).
|
|
"""
|
|
|
|
if not type_names:
|
|
raise TypeError("register_rl_model() requires at least one type name")
|
|
|
|
primary = type_names[0]
|
|
|
|
def _decorator(cls: type[RLModel]):
|
|
for name in type_names:
|
|
_RL_MODEL_REGISTRY[name] = cls
|
|
cls.type_name = primary
|
|
return cls
|
|
|
|
return _decorator
|
|
|
|
|
|
class RLModel(ABC):
|
|
"""Serializable policy/model interface.
|
|
|
|
This is the artifact that `train.py` writes and `simulate.py` loads.
|
|
"""
|
|
|
|
# Overwritten by the `@register_rl_model(...)` decorator.
|
|
type_name: str = "RLModel"
|
|
|
|
def reset(self, seed: int | None = None) -> None:
|
|
"""Optional hook for RNG/stateful models."""
|
|
|
|
@abstractmethod
|
|
def act(self, *, obs: Any | None = None, t: float = 0.0) -> Any:
|
|
raise NotImplementedError
|
|
|
|
def train(self, *, env: Any, num_epochs: int = 1) -> None:
|
|
"""Optional training hook.
|
|
|
|
Many models won't learn; for those this can be a no-op.
|
|
"""
|
|
|
|
_ = (env, num_epochs)
|
|
|
|
def to_payload(self) -> dict[str, Any]:
|
|
"""Return JSON-serializable model parameters."""
|
|
|
|
return {}
|
|
|
|
@classmethod
|
|
def from_payload(cls, payload: dict[str, Any]) -> "RLModel":
|
|
"""Reconstruct a model from `to_payload()` output."""
|
|
|
|
return cls(**payload) # type: ignore[arg-type]
|
|
|
|
def save(self, path: str | Path) -> Path:
|
|
out = Path(path)
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
doc = {
|
|
"type": self.type_name,
|
|
"version": 1,
|
|
"payload": self.to_payload(),
|
|
}
|
|
out.write_text(json.dumps(doc, indent=2, sort_keys=True) + "\n")
|
|
return out
|
|
|
|
@classmethod
|
|
def load(cls, path: str | Path) -> "RLModel":
|
|
p = Path(path)
|
|
doc = json.loads(p.read_text())
|
|
|
|
type_name = doc.get("type")
|
|
if not isinstance(type_name, str):
|
|
raise ValueError("Model artifact missing string field 'type'")
|
|
|
|
model_cls = _RL_MODEL_REGISTRY.get(type_name)
|
|
if model_cls is None:
|
|
known = ", ".join(sorted(_RL_MODEL_REGISTRY)) or "<none>"
|
|
raise ValueError(f"Unknown RLModel type '{type_name}'. Known: {known}")
|
|
|
|
payload = doc.get("payload")
|
|
# Backward compatibility: older artifacts stored fields at top-level.
|
|
if payload is None:
|
|
payload = {k: v for k, v in doc.items() if k not in ("type", "version")}
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("Model artifact field 'payload' must be an object")
|
|
|
|
return model_cls.from_payload(payload)
|