genesis-forge-runtime 1.0.0__py3-none-any.whl

genesis_forge_runtime/__init__.py

@@ -0,0 +1,100 @@
+"""Simulation-free runtime for deploying Genesis Forge policies to real robots.
+
+This package reproduces the observation-assembly and action-processing pipelines
+from training, reading them out of a bundle exported by ``genesis_forge``. It
+depends on **numpy only** -- importing anything here must never pull in torch or
+the Genesis simulator, so it installs on a Raspberry Pi or Jetson.
+
+Typical robot-side use::
+
+    from genesis_forge_runtime import load_bundle
+
+    bundle = load_bundle("./go2_walk")
+    print(bundle.describe())          # what to wire up
+
+    observation_assembler = bundle.create_observation_assembler()
+    action_processor = bundle.create_action_processor()
+
+    while True:
+        obs = observation_assembler.assemble({
+            "robot_ang_vel": gyro,
+            "actions": action_processor.last_raw_actions_by_manager["action_manager"],
+        })
+        targets = action_processor.process(policy(obs))
+        send_to_motors(targets.by_joint)
+
+The modules behind it, roughly in dependency order:
+
+* :mod:`~genesis_forge_runtime.constants` -- the manifest's vocabulary
+* :mod:`~genesis_forge_runtime.errors` -- every exception raised here
+* :mod:`~genesis_forge_runtime.serialization` -- JSON to numpy and back
+* :mod:`~genesis_forge_runtime.observation_schema` /
+  :mod:`~genesis_forge_runtime.action_schema` -- the two halves of the contract
+* :mod:`~genesis_forge_runtime.manifest` -- the contract as a whole
+* :mod:`~genesis_forge_runtime.bundle` -- reading and writing a bundle directory
+* :mod:`~genesis_forge_runtime.observations` /
+  :mod:`~genesis_forge_runtime.processors` / :mod:`~genesis_forge_runtime.actions`
+  -- the runtime itself
+
+**Trust model:** a bundle is trusted input, equivalent to executable code --
+loading one may import processor classes it names. Only load bundles you produced.
+"""
+
+from importlib.metadata import version
+
+from .action_schema import ActionManagerSpec, ActuatorSpec
+from .actions import ActionProcessor, ProcessedActions
+from .bundle import Bundle, load_bundle, load_manifest, save_bundle
+from .constants import (
+    ARCHIVE_SUFFIX,
+    EXTRACT_MARKER,
+    GOLDEN_FILENAME,
+    MANIFEST_FILENAME,
+    MIN_SUPPORTED_SCHEMA_VERSION,
+    POLICY_DIRNAME,
+    SCHEMA_VERSION,
+)
+from .errors import (
+    ActionError,
+    BundleError,
+    MalformedBundleError,
+    ObservationError,
+    SchemaVersionError,
+)
+from .manifest import Manifest, Provenance
+from .observation_schema import ObservationEntry, ObservationLayout
+from .observations import ObservationAssembler
+from .processors import ActionManagerProcessor, AffineProcessor
+
+__version__ = version("genesis-forge-runtime")
+
+__all__ = [
+    "ARCHIVE_SUFFIX",
+    "ActionError",
+    "ActionManagerProcessor",
+    "ActionManagerSpec",
+    "ActionProcessor",
+    "ActuatorSpec",
+    "AffineProcessor",
+    "Bundle",
+    "BundleError",
+    "EXTRACT_MARKER",
+    "GOLDEN_FILENAME",
+    "MANIFEST_FILENAME",
+    "MIN_SUPPORTED_SCHEMA_VERSION",
+    "MalformedBundleError",
+    "Manifest",
+    "ObservationAssembler",
+    "ObservationEntry",
+    "ObservationError",
+    "ObservationLayout",
+    "POLICY_DIRNAME",
+    "ProcessedActions",
+    "Provenance",
+    "SCHEMA_VERSION",
+    "SchemaVersionError",
+    "__version__",
+    "load_bundle",
+    "load_manifest",
+    "save_bundle",
+]

genesis_forge_runtime/action_schema.py

@@ -0,0 +1,155 @@
+"""How the policy's output maps onto real joints.
+
+The half of the manifest that :mod:`genesis_forge_runtime.processors` consumes: which
+slice of the policy vector belongs to each action manager, how to process it, and
+the actuator settings the robot should match.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .errors import MalformedBundleError
+from .serialization import require
+
+
+def _read_action_index(value: Any, name: str) -> tuple[int, ...] | None:
+    """Read the joint-to-action mapping as integers."""
+    if value is None:
+        return None
+    if not isinstance(value, (list, tuple)):
+        raise MalformedBundleError(
+            f"Action manager '{name}' has a 'joint_action_index' that is not a list."
+        )
+    try:
+        return tuple(int(index) for index in value)
+    except (TypeError, ValueError) as error:
+        raise MalformedBundleError(
+            f"Action manager '{name}' has a non-integer entry in 'joint_action_index'."
+        ) from error
+
+
+@dataclass(frozen=True)
+class ActionManagerSpec:
+    """How one action manager's slice of the policy output is processed."""
+
+    name: str
+    deploy_type: str
+    joint_names: tuple[str, ...]
+    slice_start: int
+    slice_end: int
+    config: dict[str, Any]
+    processor_import_path: str | None = None
+    #: One action index per joint, positionally matched to :attr:`joint_names`.
+    #: None when every joint has its own action.
+    joint_action_index: tuple[int, ...] | None = None
+
+    @property
+    def num_actions(self) -> int:
+        """How many policy outputs this manager consumes."""
+        return self.slice_end - self.slice_start
+
+    @property
+    def num_joints(self) -> int:
+        """How many joint targets it produces.
+
+        Larger than :attr:`num_actions` when joints share an action.
+        """
+        return len(self.joint_names)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], *, where: str) -> ActionManagerSpec:
+        name = require(data, "name", where=where)
+        scope = f"{where}.{name}"
+        bounds = require(data, "slice", where=scope)
+        if not isinstance(bounds, list) or len(bounds) != 2:
+            raise MalformedBundleError(
+                f"Action manager '{name}' has a malformed 'slice': expected [start, end], "
+                f"got {bounds!r}."
+            )
+        joint_names = tuple(require(data, "joint_names", where=scope))
+        spec = cls(
+            name=name,
+            deploy_type=require(data, "deploy_type", where=scope),
+            joint_names=joint_names,
+            slice_start=int(bounds[0]),
+            slice_end=int(bounds[1]),
+            config=data.get("config", {}),
+            processor_import_path=data.get("processor_import_path"),
+            joint_action_index=_read_action_index(data.get("joint_action_index"), name),
+        )
+        grouped = spec.joint_action_index is not None
+        if grouped and len(spec.joint_action_index) != len(joint_names):
+            raise MalformedBundleError(
+                f"Action manager '{name}' maps "
+                f"{len(spec.joint_action_index)} joint(s) to actions but names "
+                f"{len(joint_names)}."
+            )
+        if grouped and any(
+            index < 0 or index >= spec.num_actions for index in spec.joint_action_index
+        ):
+            raise MalformedBundleError(
+                f"Action manager '{name}' maps a joint to an action outside its "
+                f"slice of {spec.num_actions} action(s)."
+            )
+        if not grouped and spec.num_actions != len(joint_names):
+            raise MalformedBundleError(
+                f"Action manager '{name}' covers {spec.num_actions} actions but names "
+                f"{len(joint_names)} joints, and records no mapping between them; "
+                f"the bundle is inconsistent."
+            )
+        if grouped and len(joint_names) < spec.num_actions:
+            raise MalformedBundleError(
+                f"Action manager '{name}' groups {len(joint_names)} joints across "
+                f"{spec.num_actions} actions, so at least one action drives nothing; "
+                f"the bundle is inconsistent."
+            )
+        return spec
+
+    def to_dict(self) -> dict[str, Any]:
+        data: dict[str, Any] = {
+            "name": self.name,
+            "deploy_type": self.deploy_type,
+            "slice": [self.slice_start, self.slice_end],
+            "joint_names": list(self.joint_names),
+            "config": self.config,
+        }
+        if self.processor_import_path is not None:
+            data["processor_import_path"] = self.processor_import_path
+        if self.joint_action_index is not None:
+            data["joint_action_index"] = [
+                int(index) for index in self.joint_action_index
+            ]
+        return data
+
+
+@dataclass(frozen=True)
+class ActuatorSpec:
+    """Nominal actuator gains and defaults, recorded so the robot can match training."""
+
+    name: str
+    joint_names: tuple[str, ...]
+    values: dict[str, list[float]]
+    randomized: tuple[str, ...] = ()
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], *, where: str) -> ActuatorSpec:
+        name = require(data, "name", where=where)
+        scope = f"{where}.{name}"
+        return cls(
+            name=name,
+            joint_names=tuple(require(data, "joint_names", where=scope)),
+            values=data.get("values", {}),
+            randomized=tuple(data.get("randomized", ())),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        data: dict[str, Any] = {
+            "name": self.name,
+            "joint_names": list(self.joint_names),
+            "values": self.values,
+        }
+        if self.randomized:
+            data["randomized"] = list(self.randomized)
+        return data

genesis_forge_runtime/actions.py

@@ -0,0 +1,246 @@
+"""Turning a full policy output vector into named joint targets.
+
+Slices the vector across every action manager, hands each slice to its processor,
+and remembers the result so the caller can feed it back into the next observation.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .action_schema import ActionManagerSpec
+from .errors import ActionError
+from .processors import ActionManagerProcessor, resolve_processor_class
+
+
+class ProcessedActions:
+    """The result of one process, viewable per joint, per manager, or as a vector."""
+
+    __slots__ = ("_joint_names", "by_manager", "targets")
+
+    def __init__(
+        self,
+        targets: np.ndarray,
+        by_manager: dict[str, np.ndarray],
+        joint_names: tuple[str, ...],
+    ) -> None:
+        self.targets = targets
+        self.by_manager = by_manager
+        self._joint_names = joint_names
+
+    @property
+    def joint_names(self) -> tuple[str, ...]:
+        return self._joint_names
+
+    @property
+    def by_joint(self) -> dict[str, float]:
+        """Joint name to target value -- what you hand to the motor controllers."""
+        return {
+            name: float(value)
+            for name, value in zip(self._joint_names, self.targets, strict=True)
+        }
+
+    def __len__(self) -> int:
+        return int(self.targets.size)
+
+    def __repr__(self) -> str:  # pragma: no cover - debugging aid
+        return f"ProcessedActions({self.by_joint})"
+
+
+class ActionProcessor:
+    """Processes a full policy output vector across every action manager.
+
+    Args:
+        specs: The action manager specs from a loaded bundle's manifest.
+        check_finite: Raise when the policy emits NaN or infinity. On by default:
+            motors must never be handed those values.
+        dtype: Output dtype, float32 by default.
+    """
+
+    def __init__(
+        self,
+        specs: tuple[ActionManagerSpec, ...],
+        *,
+        check_finite: bool = True,
+        dtype: Any = np.float32,
+    ) -> None:
+        self._specs = tuple(sorted(specs, key=lambda spec: spec.slice_start))
+        self._check_finite = check_finite
+        self._dtype = np.dtype(dtype)
+
+        self._processors: list[ActionManagerProcessor] = [
+            resolve_processor_class(spec)(spec, dtype=self._dtype)
+            for spec in self._specs
+        ]
+        self._joint_names = tuple(
+            name for spec in self._specs for name in spec.joint_names
+        )
+        self.reset()
+
+    """
+    Properties
+    """
+
+    @property
+    def num_actions(self) -> int:
+        """Width of the policy output vector this processor consumes."""
+        return self._specs[-1].slice_end if self._specs else 0
+
+    @property
+    def num_joints(self) -> int:
+        """How many joint targets one call produces.
+
+        Larger than :attr:`num_actions` when joints share an action.
+        """
+        return len(self._joint_names)
+
+    @property
+    def joint_names(self) -> tuple[str, ...]:
+        """Every joint this processor produces a target for, in output order."""
+        return self._joint_names
+
+    @property
+    def clip_range_by_joint(self) -> dict[str, tuple[float, float]]:
+        """Per joint, the clip its processor applies to the target it produces.
+
+        The range a joint's targets can span, which is what a robot needs to map
+        them onto whatever its motors take. Joints whose targets are unbounded
+        are absent; see :meth:`ActionManagerProcessor.clip_range_by_joint`.
+        """
+        clips: dict[str, tuple[float, float]] = {}
+        for processor in self._processors:
+            clips.update(processor.clip_range_by_joint)
+        return clips
+
+    @property
+    def processors(self) -> tuple[ActionManagerProcessor, ...]:
+        return tuple(self._processors)
+
+    @property
+    def last_raw_actions(self) -> np.ndarray:
+        """The policy's raw output from the previous process.
+
+        Zeros before the first process, matching how training starts an episode.
+        Pass this to :meth:`ObservationAssembler.assemble` for an observation entry
+        that echoes raw policy output.
+        """
+        return self._last_raw_actions.copy()
+
+    @property
+    def last_target_actions(self) -> np.ndarray:
+        """The processed joint targets from the previous process.
+
+        Zeros before the first process.
+
+        Note this is *not* what the built-in ``current_actions`` observation feeds
+        back -- that one is raw policy output at both of its call shapes, so use
+        :attr:`last_raw_actions` or :attr:`last_raw_actions_by_manager`. This
+        property is for an observation that genuinely echoes processed targets,
+        which you would have written yourself.
+        """
+        return self._last_target_actions.copy()
+
+    @property
+    def last_target_actions_by_manager(self) -> dict[str, np.ndarray]:
+        """Per-manager view of :attr:`last_target_actions`, for multi-manager robots.
+
+        As with :attr:`last_target_actions`, this is processed output -- not what
+        ``current_actions`` feeds back.
+        """
+        return {name: values.copy() for name, values in self._last_by_manager.items()}
+
+    @property
+    def last_raw_actions_by_manager(self) -> dict[str, np.ndarray]:
+        """What the policy last emitted for each manager, before it was processed.
+
+        This is what ``current_actions(action_manager=...)`` feeds back during
+        training, and it is the whole policy vector only when a single manager is
+        registered.
+        """
+        return {name: value.copy() for name, value in self._last_raw_by_manager.items()}
+
+    """
+    Public methods
+    """
+
+    def reset(self) -> None:
+        """Clear processor state and remembered outputs."""
+        for processor in self._processors:
+            processor.reset()
+        self._last_raw_actions = np.zeros(self.num_actions, dtype=self._dtype)
+        self._last_target_actions = np.zeros(self.num_joints, dtype=self._dtype)
+        self._last_by_manager = {
+            spec.name: np.zeros(spec.num_joints, dtype=self._dtype)
+            for spec in self._specs
+        }
+        self._last_raw_by_manager = {
+            spec.name: np.zeros(spec.num_actions, dtype=self._dtype)
+            for spec in self._specs
+        }
+
+    def process(self, actions: Any) -> ProcessedActions:
+        """Process one policy output vector into joint targets.
+
+        Args:
+            actions: The policy's raw output, length :attr:`num_actions`. A
+                leading batch dimension of 1 is accepted and squeezed.
+
+        Returns:
+            A :class:`ProcessedActions` view of the result.
+
+        Raises:
+            ActionError: Wrong length, or non-finite values when
+                ``check_finite`` is on.
+        """
+        values = np.asarray(actions, dtype=self._dtype).ravel()
+        if values.size != self.num_actions:
+            raise ActionError(
+                f"Expected {self.num_actions} action(s) from the policy, got "
+                f"{values.size}."
+            )
+        if self._check_finite and not np.all(np.isfinite(values)):
+            bad = np.flatnonzero(~np.isfinite(values))
+            raise ActionError(
+                f"Policy produced non-finite action(s) at index/indices "
+                f"{bad.tolist()}: {values[bad].tolist()}. Refusing to send these to "
+                f"the actuators. Pass check_finite=False to override for debugging."
+            )
+
+        by_manager: dict[str, np.ndarray] = {}
+        raw_by_manager: dict[str, np.ndarray] = {}
+        pieces: list[np.ndarray] = []
+        for spec, processor in zip(self._specs, self._processors, strict=True):
+            chunk = values[spec.slice_start : spec.slice_end]
+            raw_by_manager[spec.name] = chunk.copy()
+            processed = processor.process(chunk)
+            by_manager[spec.name] = processed
+            pieces.append(processed)
+
+        # copy() so `targets` does not alias `by_manager[name]`; concatenate
+        # already returns a fresh array.
+        targets = (
+            pieces[0].copy() if len(pieces) == 1 else np.concatenate(pieces)
+        ).astype(self._dtype, copy=False)
+
+        self._last_raw_actions = values.copy()
+        self._last_target_actions = targets.copy()
+        self._last_raw_by_manager = raw_by_manager
+        self._last_by_manager = {
+            name: chunk.copy() for name, chunk in by_manager.items()
+        }
+
+        return ProcessedActions(targets, by_manager, self._joint_names)
+
+    def describe_outputs(self) -> str:
+        """Human-readable listing of the joint targets this processor produces."""
+        lines = [f"Joint targets produced ({self.num_joints}):"]
+        for spec in self._specs:
+            joints = ", ".join(spec.joint_names)
+            lines.append(f"  - [{spec.deploy_type}] {joints}")
+        return "\n".join(lines)
+
+    """
+    Internal methods
+    """

genesis_forge_runtime/archive.py

@@ -0,0 +1,122 @@
+"""Bundles stored as a single file.
+
+A bundle is a directory. An *archive* is that same directory written as one zip,
+which is what you copy to a robot -- one artifact instead of a folder, and a
+transfer that either arrives whole or does not arrive.
+
+Reading one extracts it beside the archive, into a dot-directory, and reuses that
+on later loads. The extracted copy records a fingerprint of the archive it came
+from, so replacing the archive re-extracts rather than leaving a robot quietly
+running the bundle it had before.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import shutil
+import zipfile
+from pathlib import Path
+
+from .constants import EXTRACT_MARKER, MANIFEST_FILENAME
+from .errors import MalformedBundleError
+
+#: Every zip begins with this. Read by content rather than by extension, so a
+#: bundle someone renamed still loads.
+_ZIP_MAGIC = b"PK\x03\x04"
+
+
+def is_archive(path: Path) -> bool:
+    """True when this file is a zip, whatever it happens to be called."""
+    if not path.is_file():
+        return False
+    with path.open("rb") as handle:
+        return handle.read(4) == _ZIP_MAGIC
+
+
+def fingerprint(archive: Path) -> str:
+    """Identify an archive's contents without reading them.
+
+    A zip's central directory already holds a CRC and a size for every entry, so
+    this costs the same whether the bundle is a kilobyte or a gigabyte -- no
+    payload is read. It is enough to notice that an archive has been replaced,
+    which is what it is for; it is not a tamper check, and a bundle is trusted
+    input either way.
+    """
+    with zipfile.ZipFile(archive) as zipped:
+        entries = sorted(
+            (item.filename, item.CRC, item.file_size) for item in zipped.infolist()
+        )
+    digest = hashlib.sha256()
+    for name, crc, size in entries:
+        digest.update(f"{name}\0{crc}\0{size}\0".encode())
+    return digest.hexdigest()
+
+
+def extract_dir_for(archive: Path) -> Path:
+    """Where an archive's contents live once unpacked: ``.<name>`` beside it."""
+    return archive.parent / f".{archive.stem}"
+
+
+def ensure_extracted(archive: Path) -> Path:
+    """Unpack an archive beside itself, reusing a previous extraction when valid.
+
+    Returns:
+        The directory holding the bundle's files.
+
+    Raises:
+        MalformedBundleError: The archive holds no manifest, the directory it
+            would extract into belongs to something else, or there is nowhere
+            writable to put it.
+    """
+    with zipfile.ZipFile(archive) as zipped:
+        if MANIFEST_FILENAME not in zipped.namelist():
+            raise MalformedBundleError(
+                f"'{archive.name}' is a zip archive but holds no "
+                f"{MANIFEST_FILENAME}, so it is not a deployment bundle."
+            )
+
+    destination = extract_dir_for(archive)
+    marker = destination / EXTRACT_MARKER
+    expected = fingerprint(archive)
+
+    if destination.exists():
+        if not destination.is_dir():
+            raise MalformedBundleError(
+                f"'{destination}' already exists and is not a directory, so "
+                f"'{archive.name}' cannot be unpacked beside itself."
+            )
+        if not marker.is_file():
+            raise MalformedBundleError(
+                f"'{destination}' already exists but was not unpacked from an "
+                f"archive, so it will not be overwritten -- it may be someone's "
+                f"work. If you know it came from '{archive.name}', deleting it is "
+                f"safe and the next load will unpack it again."
+            )
+        if marker.read_text().strip() == expected:
+            return destination
+
+    # Unpack elsewhere and swap it in, so an interrupted unpack leaves either the
+    # previous extraction or nothing -- never a half-filled directory, which the
+    # next load would mistake for someone else's and refuse.
+    staging = destination.with_name(f"{destination.name}.unpacking")
+    try:
+        if staging.exists():
+            shutil.rmtree(staging)
+        staging.mkdir(parents=True)
+        with zipfile.ZipFile(archive) as zipped:
+            # extractall keeps every member inside the target directory, so a
+            # crafted path cannot escape it.
+            zipped.extractall(staging)
+        (staging / EXTRACT_MARKER).write_text(expected)
+
+        if destination.exists():
+            shutil.rmtree(destination)
+        staging.rename(destination)
+    except OSError as error:
+        shutil.rmtree(staging, ignore_errors=True)
+        raise MalformedBundleError(
+            f"Could not unpack '{archive.name}' into '{destination}' ({error}). "
+            f"Unpack it yourself somewhere writable and load that directory."
+        ) from error
+
+    return destination