royalelearn 0.5.4__py3-none-any.whl

royalelearn/__init__.py

@@ -0,0 +1,60 @@
+"""royalelearn -- the self-play training harness for RoyaleGym environments.
+
+Layers (each user-facing behaviour is an ABC in ``api/`` with a shipped default):
+
+    RolloutSource     api/rollout.py     ProcessRolloutSource, InlineRolloutSource
+    ActorCritic       api/policy.py      SeparateActorCritic, SharedTrunkActorCritic
+    ObsCodec          api/buffer.py      SpatialObsCodec
+    ExperienceBuffer  api/buffer.py      RectBuffer
+    Matchmaker/Rater  api/ladder.py      MixMatchmaker, BradleyTerryDavidsonRater
+    MetricsSink       api/metrics.py     JsonlSink, ConsoleSink, WandbSink
+    CheckpointStore   api/checkpoint.py  DirCheckpointStore
+
+``LearningCoordinator`` (coordinator.py) is the only place the phases of an iteration are
+ordered. ``docs/harness-spec.md`` specifies all of it; ``docs/design.md`` says why.
+
+Importing this package pulls in nothing that needs torch: the ABCs use numpy and msgspec,
+and the concrete learner lives behind the lazy attributes below. That is what lets
+``royalelearn config``, ``royalelearn identity`` and ``royalelearn --help`` work in an
+environment where torch is not installed. Asking for a name that does need torch raises an
+ImportError naming the package that is missing.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+from .version import __version__, git_describe
+
+# Public name -> (module, attribute). A None attribute exports the module itself.
+# Resolution is deferred so that the import cost and the torch dependency of a name are
+# paid only by the caller that asks for it.
+_EXPORTS: dict[str, tuple[str, str | None]] = {
+    "Learner": (".learner", "Learner"),
+    "LearningCoordinator": (".coordinator", "LearningCoordinator"),
+    "RunConfig": (".config", "RunConfig"),
+    "load_config": (".config", "load_config"),
+    "api": (".api", None),
+}
+
+__all__ = ["__version__", "git_describe", *sorted(_EXPORTS)]
+
+
+def __getattr__(name: str) -> Any:
+    """Resolve a public name on first use.
+
+    An unknown name is an AttributeError. A known name whose module cannot be imported
+    raises that import's own ImportError unchanged, because its message already names what
+    is missing and rewording it would hide which package to install.
+    """
+    try:
+        module_name, attribute = _EXPORTS[name]
+    except KeyError:
+        raise AttributeError(f"module 'royalelearn' has no attribute {name!r}") from None
+    module = importlib.import_module(module_name, __name__)
+    return module if attribute is None else getattr(module, attribute)
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

royalelearn/__main__.py

@@ -0,0 +1,15 @@
+"""``python -m royalelearn``.
+
+The module body is one call. Everything the entry point has to do before torch is imported --
+the deterministic cuBLAS workspace and the BLAS thread counts -- ``cli`` does as it is imported,
+which is why this file imports it and nothing else.
+"""
+
+from __future__ import annotations
+
+import sys
+
+from .cli import main
+
+if __name__ == "__main__":
+    sys.exit(main())

royalelearn/api/__init__.py

@@ -0,0 +1,121 @@
+"""Every ABC and struct the harness is written against.
+
+This subpackage imports numpy and msgspec and nothing else that matters: torch appears in type
+annotations only, behind ``TYPE_CHECKING``. That is what lets a rollout worker, the CLI's
+config and identity commands, and ``import royalelearn`` itself run in an environment with no
+torch installed -- and what keeps a worker's resident memory three hundred megabytes smaller
+than the parent's.
+
+The concrete implementations live outside ``api/`` and may import whatever they need:
+
+    RolloutSource        rollout/farm.py, rollout/inline.py
+    ActorCritic          learn/actor_critic.py
+    ObsCodec             rollout/codec.py
+    ExperienceBuffer     learn/buffer.py
+    AdvantageEstimator   learn/gae.py
+    Update               learn/ppo.py
+    Schedule             learn/schedules.py
+    Matchmaker, Rater    ladder/matchmaker.py, ladder/rating.py
+    MetricsSink, Alarm   metrics/sinks.py, metrics/alarms.py
+    CheckpointStore      checkpoint.py
+"""
+
+from __future__ import annotations
+
+from .advantage import AdvantageEstimator, AdvantageStats
+from .buffer import CodecTable, ExperienceBuffer, ObsCodec
+from .checkpoint import Checkpointable, CheckpointStore, Manifest, RngState
+from .ladder import (
+    ConditionResult,
+    EvictionPolicy,
+    GateDecision,
+    Matchmaker,
+    PromotionGate,
+    Rater,
+    RatingTable,
+    SnapshotStore,
+)
+from .metrics import Alarm, AlarmResult, MetricRow, MetricsSink, MetricValue
+from .policy import (
+    ActionDistribution,
+    Actor,
+    ActorCritic,
+    ActResult,
+    BackpropResult,
+    Critic,
+    NetworkFactory,
+    ObsBatch,
+)
+from .rollout import (
+    Assignment,
+    Close,
+    Defer,
+    EnvSpec,
+    EpisodeRecord,
+    ObsKeySpec,
+    Plan,
+    RolloutRound,
+    RolloutSource,
+    SetState,
+    SlotPlan,
+    Spaces,
+    Step,
+    WorkerCommand,
+    WorkerFailure,
+)
+from .schedule import Schedule, ScheduleState
+from .update import ActorLossTerm, ActorTermInputs, Update, UpdateResult
+
+__all__ = [
+    "ActResult",
+    "ActionDistribution",
+    "Actor",
+    "ActorCritic",
+    "ActorLossTerm",
+    "ActorTermInputs",
+    "AdvantageEstimator",
+    "AdvantageStats",
+    "Alarm",
+    "AlarmResult",
+    "Assignment",
+    "BackpropResult",
+    "CheckpointStore",
+    "Checkpointable",
+    "Close",
+    "CodecTable",
+    "ConditionResult",
+    "Critic",
+    "Defer",
+    "EnvSpec",
+    "EpisodeRecord",
+    "EvictionPolicy",
+    "ExperienceBuffer",
+    "GateDecision",
+    "Manifest",
+    "Matchmaker",
+    "MetricRow",
+    "MetricValue",
+    "MetricsSink",
+    "NetworkFactory",
+    "ObsBatch",
+    "ObsCodec",
+    "ObsKeySpec",
+    "Plan",
+    "PromotionGate",
+    "Rater",
+    "RatingTable",
+    "RngState",
+    "RolloutRound",
+    "RolloutSource",
+    "Schedule",
+    "ScheduleState",
+    "SetState",
+    "SlotPlan",
+    "SnapshotStore",
+    "Spaces",
+    "Step",
+    "Update",
+    "UpdateResult",
+    "WorkerCommand",
+    "WorkerFailure",
+]

royalelearn/api/advantage.py

@@ -0,0 +1,61 @@
+"""How a return is turned into an advantage.
+
+One ABC, because the estimator is the piece most likely to be replaced -- V-trace against a
+frozen pool, a different lambda rule, a per-bucket normalisation -- and because the one thing it
+must never do is carry its recursion across an episode boundary.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+import msgspec
+
+from .checkpoint import Checkpointable
+
+if TYPE_CHECKING:  # pragma: no cover - annotations only
+    from torch import Tensor
+
+__all__ = ["AdvantageEstimator", "AdvantageStats"]
+
+
+class AdvantageStats(msgspec.Struct):
+    """What the estimator saw, for the metric row.
+
+    ``reward_scale`` is the divisor the return scaler applied and ``clipped_reward_frac`` is how
+    much of the batch the clip bound touched: both are how a scaled reward stays legible after
+    the scaling.
+    """
+
+    raw_return_mean: float
+    raw_return_std: float
+    reward_scale: float
+    clipped_reward_frac: float
+
+
+class AdvantageEstimator(ABC, Checkpointable):
+    """Rewards and values in, advantages and returns out."""
+
+    @abstractmethod
+    def compute(
+        self,
+        *,
+        rewards: Tensor,
+        values: Tensor,
+        final_values: Tensor,
+        terminated: Tensor,
+        truncated: Tensor,
+        trainable: Tensor,
+        gamma: float,
+        lam: float,
+    ) -> tuple[Tensor, Tensor, AdvantageStats]:
+        """``rewards``/``terminated``/``truncated``/``trainable`` are ``(T, R)``; ``values`` is
+        ``(T+1, R)``; ``final_values`` is ``(T, R)`` and is read only where truncated.
+        Returns ``(advantages (T, R), returns (T, R), stats)``.
+
+        Implementers MUST bootstrap a terminated cell from 0 and a truncated cell from
+        ``final_values``, and MUST NOT carry the recursion across an episode boundary. A
+        truncation is an episode that was cut, not one that was decided, and treating the two
+        alike throws away the value of every position a step limit ended.
+        """

royalelearn/api/buffer.py

@@ -0,0 +1,179 @@
+"""How an observation is stored, and the rectangle it is stored in.
+
+The two are one subject. The worker packs a row straight into its final resting place in the
+experience buffer, so there is no intermediate copy to agree about; the learner unpacks it on
+the device inside the same kernel that gathers the frame stack. What the codec may decide is
+how each key is stored, and that is decided from the observation space at preflight rather than
+from a list of plane indices written down here.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterator, Sequence
+from typing import TYPE_CHECKING
+
+import msgspec
+import numpy as np
+
+from .checkpoint import Checkpointable
+
+if TYPE_CHECKING:  # pragma: no cover - annotations only
+    from torch import Tensor
+
+    from ..learn.buffer import Batch
+    from ..rollout.layout import BufferHandle
+    from .policy import ObsBatch
+    from .rollout import EnvSpec, RolloutRound, SlotPlan
+
+__all__ = ["MIN_TABLE_STATES", "CodecTable", "ExperienceBuffer", "ObsCodec"]
+
+#: How many real observations a codec table may be decided from. Storage is decided per plane
+#: from what the sample contains, so the sample has to be large enough, and played rather than
+#: idle, to have reached the states in which a plane takes the values that decide it.
+MIN_TABLE_STATES = 1000
+
+
+class CodecTable(msgspec.Struct, frozen=True, omit_defaults=True):
+    """How each observation key is stored, decided from ``EnvSpec.obs_space`` at preflight
+    rather than from a list of plane indices. Logged, hashed and written into every snapshot.
+
+    ``plane`` is one entry per spatial plane: its name, its storage in
+    {"uint8", "float16", "static", "derived"}, and the divisor that takes the stored integer
+    back to the value the environment produced. "derived" is reserved: nothing rebuilds such a
+    plane yet, so the codec refuses it at bind. Two runs whose tables differ are not comparable,
+    and ``digest()`` is what says so.
+    """
+
+    plane: tuple[tuple[str, str, float], ...]
+    vector: str
+    mask: str
+    #: How ``card_ids`` is stored: "uint8", exact, or None when the observation has none. It is
+    #: omitted from the encoding when None (``omit_defaults``), so every table decided before
+    #: card identity existed hashes exactly as it did and no saved run's digest moves.
+    ids: str | None = None
+
+    def digest(self) -> str:
+        """sha256 of the canonical JSON of this table."""
+        from ..rollout.envspec import digest_of
+
+        return digest_of(self)
+
+
+class ObsCodec(ABC):
+    """Quantisation of one observation row. The worker packs; the learner unpacks on the GPU.
+
+    Implementers MUST be exact round-trips for the integer-valued channels and MUST declare
+    ``row_bytes`` as a constant given an ``EnvSpec`` and its ``CodecTable``.
+    """
+
+    @abstractmethod
+    def table(
+        self,
+        spec: EnvSpec,
+        sample: Sequence[dict[str, np.ndarray]],
+        *,
+        min_states: int = MIN_TABLE_STATES,
+    ) -> CodecTable:
+        """Decide storage per key from the declared bounds and a sample of real observations.
+
+        Storage is decided from a sample; EXISTENCE is decided from the declaration. A plane the
+        layout does not declare static is stored even when it is constant across the sample,
+        because the tower planes are constant in any sample in which no tower falls.
+
+        Implementers MUST refuse a sample of fewer than ``min_states`` observations. The row
+        size of the whole run follows from this one decision, and a sample too small or too
+        idle to have reached the states a plane varies in decides it wrongly and in silence.
+        """
+
+    @abstractmethod
+    def row_bytes(self, spec: EnvSpec) -> int: ...
+
+    @abstractmethod
+    def pack(self, obs: dict[str, np.ndarray], out: memoryview, row: int) -> None: ...
+
+    @abstractmethod
+    def static_planes(self, obs: dict[str, np.ndarray]) -> np.ndarray:
+        """The planes ``EnvSpec.spatial_layout`` declares static; stored once per seat, never
+        per row."""
+
+    @abstractmethod
+    def unpack_to_device(self, raw: Tensor, statics: Tensor, out: ObsBatch) -> None:
+        """Dequantise, scatter the static planes in, reshape the stored mask into the mask
+        planes, and gather the frame-stack history."""
+
+    @property
+    @abstractmethod
+    def codec_version(self) -> int:
+        """The RULE's version. The table it produces is data and travels separately."""
+
+
+class ExperienceBuffer(ABC, Checkpointable):
+    """A rectangle of ``(T + frame_stack)`` cycles x ``R`` slots: ``T`` collected cycles, one
+    bootstrap row, and ``frame_stack - 1`` history rows carried over from the previous
+    iteration. Owns the shared-memory block the workers write into.
+
+    Implementers may assume each ``(cycle, slot)`` cell is written exactly once, by the worker
+    that owns that slot; the learner only reads.
+    """
+
+    @abstractmethod
+    def shared_handle(self) -> BufferHandle:
+        """Name, size and the numbers the offsets follow from; picklable, and sent to workers."""
+
+    @abstractmethod
+    def begin_iteration(self, plan: SlotPlan, cycles: int) -> None: ...
+
+    @abstractmethod
+    def record_round(
+        self, r: RolloutRound, actions: np.ndarray, log_probs: np.ndarray
+    ) -> None:
+        """Scalars only: observations are already in place. O(n), no observation copy."""
+
+    @abstractmethod
+    def set_values(self, values: Tensor) -> None:
+        """``(T+1, R)`` float32, from the whole-iteration critic pass."""
+
+    @abstractmethod
+    def set_n_legal(self, counts: Tensor) -> None:
+        """``(T+1, R)`` integer: how many actions each cell's mask left, from the critic's pass.
+
+        Implementers keep the collected cycles and may drop the bootstrap row, in which no
+        action was taken. The count is what tells the update which rows had a choice at all
+        without unpacking an observation to find out.
+        """
+
+    @abstractmethod
+    def set_final_values(self, cells: np.ndarray, values: Tensor) -> None:
+        """V(final_obs) for truncated cells; ``cells`` is int64[(k, 2)] of (cycle, slot)."""
+
+    @abstractmethod
+    def set_advantages(self, adv: Tensor, ret: Tensor) -> None:
+        """``(T, R)`` float32 each."""
+
+    @abstractmethod
+    def trainable_mask(self) -> Tensor:
+        """``(T, R)`` bool: which cells reach the update. What decides it is the seat's group
+        and the cell's validity, never where the row was written."""
+
+    @abstractmethod
+    def batches(
+        self,
+        batch_size: int,
+        minibatch_size: int,
+        epochs: int,
+        rng_for_epoch: Callable[[int], np.random.Generator],
+        *,
+        choice_first: bool = False,
+    ) -> Iterator[Batch]:
+        """Yield batches; each Batch knows its true sample count and iterates device-resident
+        minibatches. A batch never straddles an epoch boundary, and an epoch holds as many whole
+        batches as it can fill with the rows over spread one each across them, so every batch is
+        at least ``batch_size`` and an epoch is exactly ``n // batch_size`` optimizer steps. An
+        epoch with nothing trainable in it yields no batches at all. Each minibatch is weighted by
+        its share of its own batch. Gathers per MINIBATCH, never per batch.
+
+        ``choice_first`` reorders each batch's cells so that the ones with more than one legal
+        action come first, keeping the permutation's order inside each class. Implementers MUST
+        move no cell between batches: it is a reordering, so that a caller skipping the forced
+        rows skips whole minibatches of them, and every batch-level denominator is unchanged."""

royalelearn/api/checkpoint.py

@@ -0,0 +1,132 @@
+"""What a checkpoint is: components that save themselves, a store that writes them atomically,
+and a manifest that says what was written and proves the bytes are still what was written.
+
+Every component owns its folder and its own pair of methods. Adding a component to a checkpoint
+is adding a folder name and a dict entry; there is no central serialiser to edit, and no place
+where a component's state is described twice.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, Protocol, runtime_checkable
+
+import msgspec
+
+# Imported at run time, not behind TYPE_CHECKING: Manifest carries a RunIdentity, and msgspec
+# resolves a Struct's annotations when a decoder for it is built, so a name that exists only
+# for a type checker would make a manifest unreadable.
+from ..identity import RunIdentity
+
+__all__ = ["CheckpointStore", "Checkpointable", "Manifest", "RngState"]
+
+
+@runtime_checkable
+class Checkpointable(Protocol):
+    """Anything that goes into a checkpoint folder of its own.
+
+    ``FORMAT_VERSION`` is the component's own, not the checkpoint's: a component may change its
+    files without the checkpoint format moving, and the manifest records each one so that a
+    load can say which component it could not read.
+    """
+
+    FORMAT_VERSION: int
+
+    def save_checkpoint(self, folder: Path) -> None: ...
+
+    def load_checkpoint(self, folder: Path, *, strict: bool) -> None:
+        """With ``strict=False``, a missing file prints the exact path it wanted and continues
+        with a default. With ``strict=True`` it raises. ``strict`` defaults to True on resume:
+        tolerate-everything is right for a research tool and wrong for a harness that promises
+        the curve continues."""
+
+
+class Manifest(msgspec.Struct):
+    """What one checkpoint is, beside the folders that hold it.
+
+    ``config`` is the resolved config verbatim, so a checkpoint is self-describing without the
+    run directory around it, and ``files`` is every file in the checkpoint with its sha256, so a
+    truncated write from a crash six weeks ago is an error on load rather than a silently wrong
+    resume.
+    """
+
+    format_version: int
+    run_id: str
+    run_name: str
+    identity: RunIdentity
+    config: dict[str, Any]
+    config_hash: str
+    iteration: int
+    cumulative_env_steps: int
+    cumulative_timesteps: int
+    cumulative_model_updates: int
+    wall_seconds: float
+    created_unix_ns: int
+    state_digest: str
+    component_versions: dict[str, int]
+    files: dict[str, str]
+    #: Seconds spent gating so far, and the last gate's decision: run state, like the wall clock.
+    #: None on a manifest written before either was recorded.
+    gate_seconds: float | None = None
+    #: Stored as plain data and converted by the coordinator, so this module need not import
+    #: the ladder's types: ``api.ladder`` already imports this one.
+    last_decision: dict[str, Any] | None = None
+    #: The env step each cadence last fired at: the periodic checkpoint, the gate's candidate
+    #: and the floor's admission. A resume puts them back, so each fires where it would have in
+    #: a run that never stopped. None on a manifest written before they were recorded; a resume
+    #: then starts all three at this checkpoint's env step, as every resume did before.
+    last_checkpoint_step: int | None = None
+    last_candidate_step: int | None = None
+    last_floor_step: int | None = None
+
+
+class RngState(msgspec.Struct):
+    """Every random stream's position at the moment a checkpoint was written.
+
+    One reference learner saves none of this; the other re-seeds from the run's INITIAL seed on
+    load, so a run resumed at ten million steps draws the same action noise it drew at step
+    zero. Because every stream here is name-addressed, restoring the iteration counter and the
+    shard positions restores the stream, not merely the parameters.
+    """
+
+    master_seed: int
+    torch_cpu: str
+    torch_cuda: list[str]
+    python_random: list[Any]
+    numpy_minibatch: dict[str, Any]
+    iteration: int
+    shard_streams: list[dict[str, Any]]
+    eval_seed_set_sha: str
+
+
+class CheckpointStore(ABC):
+    """Where checkpoints live, and the only thing that writes or reads them."""
+
+    @abstractmethod
+    def write(self, components: Mapping[str, Checkpointable], manifest: Manifest) -> Path:
+        """Write atomically: into ``<name>.partial/``, fsync each file, then ``os.replace``.
+
+        The directory itself is not fsynced. ``os.fsync`` on a directory handle is a POSIX
+        guarantee and raises on Windows, where this harness's default profile runs, so the
+        durability step is per file and the atomic step is the rename -- which is atomic on both
+        platforms, and is the property recovery actually needs.
+        """
+
+    @abstractmethod
+    def read(
+        self, path: Path, components: Mapping[str, Checkpointable], *, strict: bool
+    ) -> Manifest:
+        """Verify every hash in the manifest, then load each component. A mismatch raises
+        ``CheckpointFormatError`` naming the first file that failed."""
+
+    @abstractmethod
+    def latest(self, run_dir: Path) -> Path | None:
+        """The newest checkpoint, from the run index rather than from parsing directory names:
+        ``int(x) for x in os.listdir(...)`` crashes on a stray file, and the save is called from
+        the crash handler, which is where the user most needs it not to."""
+
+    @abstractmethod
+    def prune(self, run_dir: Path, keep: int) -> list[Path]:
+        """Remove all but the newest ``keep``, and return what was removed."""