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

genesis_forge_runtime/bundle.py

@@ -0,0 +1,220 @@
+"""Reading and writing a bundle directory.
+
+A bundle is what export writes and the robot reads::
+
+    go2_walk/
+      manifest.json   # the deployment contract, human readable
+      golden.npz      # recorded input/output pairs for the on-robot smoke test
+      policy.onnx     # optional: the exported policy
+
+Nothing in this package imports torch or genesis -- that is the whole point of it,
+and ``test_bundle.py`` asserts it in a clean subprocess.
+"""
+
+from __future__ import annotations
+
+import tempfile
+import zipfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from .actions import ActionProcessor
+from .archive import ensure_extracted, is_archive
+from .constants import GOLDEN_FILENAME, MANIFEST_FILENAME, POLICY_DIRNAME
+from .errors import MalformedBundleError
+from .manifest import Manifest
+from .observations import ObservationAssembler
+
+
+@dataclass(frozen=True)
+class Bundle:
+    """A deployment bundle: the manifest, and wherever its files are.
+
+    ``path`` is a directory when the bundle is unpacked, and an archive file when
+    it is not -- which is what :func:`~genesis_forge.deployment.export` hands back
+    after writing one. The manifest and the golden samples are held in memory
+    either way, so describing a bundle never unpacks it.
+    """
+
+    manifest: Manifest
+    path: Path
+    golden: dict[str, np.ndarray] | None = None
+
+    @property
+    def is_archive(self) -> bool:
+        """True when the files are still packed into a single file."""
+        return not self.path.is_dir()
+
+    @property
+    def policy_files(self) -> tuple[str, ...]:
+        """Every policy file in the bundle, in the order it was exported."""
+        return self.manifest.policy
+
+    @property
+    def policy_dir(self) -> Path | None:
+        """Where the policy files sit, relative to the bundle's own root.
+
+        Join it onto the directory :meth:`unpacked` yields, which works whether
+        the bundle arrived as a directory or an archive.
+        """
+        if not self.policy_files:
+            return None
+        return Path(POLICY_DIRNAME)
+
+    @property
+    def policy_path(self) -> Path | None:
+        """The first policy file, relative to the bundle's own root.
+
+        A convenience for the common single-file case; :attr:`policy_files` has
+        them all.
+
+        Example::
+
+            with bundle.unpacked() as directory:
+                session = onnxruntime.InferenceSession(
+                    str(directory / bundle.policy_path)
+                )
+        """
+        if not self.policy_files:
+            return None
+        return Path(POLICY_DIRNAME) / self.policy_files[0]
+
+    @contextmanager
+    def unpacked(self) -> Iterator[Path]:
+        """Yield a directory holding this bundle's files.
+
+        Already-unpacked bundles yield their own directory and nothing is copied.
+        An archive is unpacked to a temporary directory that is removed on the way
+        out, which is what you want on a training machine -- checking what you just
+        exported should not leave anything behind.
+        """
+        if not self.is_archive:
+            yield self.path
+            return
+        with tempfile.TemporaryDirectory(prefix="genesis-forge-bundle-") as scratch:
+            with zipfile.ZipFile(self.path) as archive:
+                archive.extractall(scratch)
+            yield Path(scratch)
+
+    def create_observation_assembler(self, **kwargs: Any) -> ObservationAssembler:
+        """Build a new :class:`ObservationAssembler` for this bundle.
+
+        Each call returns a fresh assembler with its own zero-filled history, so
+        create one and keep it for the life of the control loop -- a second one
+        starts with no history at all.
+        """
+        return ObservationAssembler(self.manifest.observations, **kwargs)
+
+    def create_action_processor(self, **kwargs: Any) -> ActionProcessor:
+        """Build a new :class:`ActionProcessor` for this bundle.
+
+        Each call returns a fresh processor with its own remembered actions, so
+        create one and keep it for the life of the control loop.
+        """
+        return ActionProcessor(self.manifest.actions, **kwargs)
+
+    def describe(self) -> str:
+        """Human-readable summary of what to wire up. Print this first."""
+        layout = self.manifest.observations
+        lines = [
+            f"Bundle: {self.path}",
+            f"  control rate: {self.manifest.control_hz:.1f} Hz (dt={self.manifest.dt})",
+            (
+                f"  observation vector: {layout.total_size} values "
+                f"({layout.single_size} per tick x {layout.history_length} history)"
+            ),
+            "  values you supply each tick:",
+        ]
+        lines.extend(f"    - {entry.describe()}" for entry in layout.entries)
+        joints = len(self.manifest.joint_names)
+        if joints == self.manifest.num_actions:
+            lines.append(f"  joint targets produced ({joints}):")
+        else:
+            lines.append(
+                f"  joint targets produced ({joints}, from "
+                f"{self.manifest.num_actions} policy outputs):"
+            )
+        for spec in sorted(self.manifest.actions, key=lambda item: item.slice_start):
+            joints = ", ".join(spec.joint_names)
+            lines.append(f"    - [{spec.deploy_type}] {joints}")
+        if self.policy_files:
+            files = ", ".join(f"{POLICY_DIRNAME}/{name}" for name in self.policy_files)
+            lines.append(f"  policy: {files}")
+        return "\n".join(lines)
+
+
+def load_manifest(path: str | Path) -> Manifest:
+    """Read and validate just the manifest, without touching the rest of the bundle."""
+    manifest_path = Path(path)
+    if manifest_path.is_dir():
+        manifest_path = manifest_path / MANIFEST_FILENAME
+    if not manifest_path.is_file():
+        raise MalformedBundleError(f"No manifest found at '{manifest_path}'.")
+    return Manifest.from_json(manifest_path.read_text())
+
+
+def load_bundle(path: str | Path, *, load_golden: bool = True) -> Bundle:
+    """Load a deployment bundle.
+
+    Args:
+        path: A bundle directory, or an archive holding one. An archive is
+            unpacked beside itself into ``.<name>/`` and reused on later loads,
+            so the files stay in one predictable place for as long as anything
+            needs them. Replacing the archive unpacks it again.
+        load_golden: Read ``golden.npz`` when present. Set False to skip the
+            recorded smoke-test samples.
+
+    Returns:
+        The loaded :class:`Bundle`. Its :attr:`~Bundle.path` is the directory the
+        files were read from, which for an archive is where they were unpacked.
+
+    Raises:
+        SchemaVersionError: The bundle was written by an incompatible version.
+        MalformedBundleError: The bundle is missing something required, or an
+            archive could not be unpacked.
+    """
+    bundle_path = Path(path)
+    if is_archive(bundle_path):
+        bundle_path = ensure_extracted(bundle_path)
+    if not bundle_path.is_dir():
+        raise MalformedBundleError(
+            f"'{bundle_path}' is not a bundle directory or archive. Expected a "
+            f"folder containing {MANIFEST_FILENAME}, or a single-file bundle."
+        )
+
+    manifest = load_manifest(bundle_path)
+
+    golden: dict[str, np.ndarray] | None = None
+    golden_path = bundle_path / GOLDEN_FILENAME
+    if load_golden and golden_path.is_file():
+        with np.load(golden_path, allow_pickle=False) as archive:
+            golden = {key: archive[key] for key in archive.files}
+
+    for name in manifest.policy:
+        if not (bundle_path / POLICY_DIRNAME / name).is_file():
+            raise MalformedBundleError(
+                f"Manifest lists policy file '{name}', but it is missing from "
+                f"'{bundle_path}'."
+            )
+
+    return Bundle(manifest=manifest, path=bundle_path, golden=golden)
+
+
+def save_bundle(
+    bundle_path: str | Path,
+    manifest: Manifest,
+    *,
+    golden: dict[str, np.ndarray] | None = None,
+) -> Path:
+    """Write a manifest (and optional golden samples) into a bundle directory."""
+    path = Path(bundle_path)
+    path.mkdir(parents=True, exist_ok=True)
+    (path / MANIFEST_FILENAME).write_text(manifest.to_json() + "\n")
+    if golden:
+        np.savez(path / GOLDEN_FILENAME, **golden)
+    return path

genesis_forge_runtime/constants.py

@@ -0,0 +1,28 @@
+"""The deployment bundle's vocabulary.
+
+These strings and numbers are the manifest's contract. They are deliberately
+separate from the code that uses them: a bundle may be read by tooling that never
+imports the rest of this package, and the values must not drift.
+"""
+
+#: Schema version this runtime writes and understands.
+SCHEMA_VERSION = 1
+
+#: Oldest bundle schema this runtime can still read.
+MIN_SUPPORTED_SCHEMA_VERSION = 1
+
+MANIFEST_FILENAME = "manifest.json"
+GOLDEN_FILENAME = "golden.npz"
+
+#: Suffix for a bundle written as one file: a zip holding what the directory
+#: would. Bundles are recognised by content, so a renamed one still loads.
+ARCHIVE_SUFFIX = ".gfb"
+
+#: Written inside an extracted archive, holding a fingerprint of the archive it
+#: came from. A later load re-extracts unless the fingerprint still matches, so
+#: replacing the archive can never leave a robot running the previous bundle.
+EXTRACT_MARKER = ".source"
+
+#: Directory inside the bundle holding the exported policy, under the filenames
+#: the exporter was given. The runtime never opens them.
+POLICY_DIRNAME = "policy"

genesis_forge_runtime/errors.py

@@ -0,0 +1,32 @@
+"""Every error this package raises.
+
+Collected in one module so a robot-side caller can see the whole hierarchy at a
+glance and catch at whatever granularity it wants::
+
+    try:
+        bundle = load_bundle(path)
+    except SchemaVersionError:
+        ...   # re-export from a matching Genesis Forge
+    except BundleError:
+        ...   # anything else wrong with the bundle
+"""
+
+
+class BundleError(Exception):
+    """Base class for every error raised while reading a bundle."""
+
+
+class SchemaVersionError(BundleError):
+    """The bundle was written by an incompatible version of Genesis Forge."""
+
+
+class MalformedBundleError(BundleError):
+    """The bundle is missing a required section or holds an unusable value."""
+
+
+class ObservationError(Exception):
+    """A value handed to the assembler was missing, mis-sized, or unexpected."""
+
+
+class ActionError(Exception):
+    """A processor could not be resolved, or was handed unusable policy output."""

genesis_forge_runtime/manifest.py

@@ -0,0 +1,200 @@
+"""The deployment contract as a whole.
+
+Ties the observation and action halves together with the control rate, the policy
+description, and provenance, and validates that the pieces are mutually consistent
+before any of it reaches a robot.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+from .action_schema import ActionManagerSpec, ActuatorSpec
+from .constants import MIN_SUPPORTED_SCHEMA_VERSION, SCHEMA_VERSION
+from .errors import MalformedBundleError, SchemaVersionError
+from .observation_schema import ObservationLayout
+from .serialization import require
+
+
+@dataclass(frozen=True)
+class Provenance:
+    """Where this bundle came from -- the first thing to check when debugging.
+
+    The top-level fields are measured by the exporter itself. Anything the
+    developer chose to record sits under :attr:`additional`, so a reader can tell
+    a fact the tooling observed from a claim a human typed.
+    """
+
+    exported_at: str | None = None
+    genesis_forge_version: str | None = None
+    torch_version: str | None = None
+    additional: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        # Ensure the additional dict is JSON-serializable
+        json.dumps(self.additional)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Provenance:
+        additional = data.get("additional") or {}
+        if not isinstance(additional, dict):
+            raise MalformedBundleError(
+                f"'provenance.additional' must be an object, got "
+                f"{type(additional).__name__}."
+            )
+        return cls(
+            exported_at=data.get("exported_at"),
+            genesis_forge_version=data.get("genesis_forge_version"),
+            torch_version=data.get("torch_version"),
+            additional=dict(additional),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        data = {
+            key: value
+            for key, value in {
+                "exported_at": self.exported_at,
+                "genesis_forge_version": self.genesis_forge_version,
+                "torch_version": self.torch_version,
+            }.items()
+            if value is not None
+        }
+        if self.additional:
+            data["additional"] = dict(self.additional)
+        return data
+
+
+@dataclass(frozen=True)
+class Manifest:
+    """The full deployment contract, as read from ``manifest.json``."""
+
+    dt: float
+    observations: ObservationLayout
+    actions: tuple[ActionManagerSpec, ...]
+    actuators: tuple[ActuatorSpec, ...] = ()
+    #: The policy's files, in the order they were exported. Empty when the bundle
+    #: carries no policy. They live in the bundle's ``policy/`` directory.
+    policy: tuple[str, ...] = ()
+    provenance: Provenance = field(default_factory=Provenance)
+    schema_version: int = SCHEMA_VERSION
+
+    @property
+    def control_hz(self) -> float:
+        """Rate the policy was trained at; run the robot loop at this rate."""
+        return 1.0 / self.dt
+
+    @property
+    def num_actions(self) -> int:
+        return max((spec.slice_end for spec in self.actions), default=0)
+
+    @property
+    def joint_names(self) -> tuple[str, ...]:
+        """Every controlled joint, in policy-output order."""
+        names: list[str] = []
+        for spec in sorted(self.actions, key=lambda item: item.slice_start):
+            names.extend(spec.joint_names)
+        return tuple(names)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Manifest:
+        version = require(data, "schema_version", where="manifest")
+        _check_schema_version(version)
+
+        control = require(data, "control", where="manifest")
+        dt = float(require(control, "dt", where="control"))
+        if dt <= 0:
+            raise MalformedBundleError(f"'control.dt' must be positive, got {dt}.")
+
+        raw_actions = require(data, "actions", where="manifest")
+        actions = tuple(
+            ActionManagerSpec.from_dict(item, where="actions")
+            for item in require(raw_actions, "managers", where="actions")
+        )
+        _check_action_slices(actions)
+
+        return cls(
+            schema_version=int(version),
+            dt=dt,
+            observations=ObservationLayout.from_dict(
+                require(data, "observations", where="manifest")
+            ),
+            actions=actions,
+            actuators=tuple(
+                ActuatorSpec.from_dict(item, where="actuators")
+                for item in data.get("actuators", [])
+            ),
+            policy=tuple(data.get("policy", ())),
+            provenance=Provenance.from_dict(data.get("provenance", {})),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        data: dict[str, Any] = {
+            "schema_version": self.schema_version,
+            "provenance": self.provenance.to_dict(),
+            "control": {"dt": self.dt, "control_hz": self.control_hz},
+            "observations": self.observations.to_dict(),
+            "actions": {
+                "total_size": self.num_actions,
+                "managers": [spec.to_dict() for spec in self.actions],
+            },
+        }
+        if self.actuators:
+            data["actuators"] = [spec.to_dict() for spec in self.actuators]
+        if self.policy:
+            data["policy"] = list(self.policy)
+        return data
+
+    def to_json(self, *, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent, sort_keys=False)
+
+    @classmethod
+    def from_json(cls, text: str) -> Manifest:
+        try:
+            data = json.loads(text)
+        except json.JSONDecodeError as error:
+            raise MalformedBundleError(
+                f"manifest.json is not valid JSON: {error}"
+            ) from error
+        return cls.from_dict(data)
+
+
+def _check_schema_version(version: Any) -> None:
+    """Fail loudly, and specifically, on a bundle this runtime cannot read."""
+    if not isinstance(version, int) or isinstance(version, bool):
+        raise MalformedBundleError(
+            f"'schema_version' must be an integer, got {version!r}."
+        )
+    if version > SCHEMA_VERSION:
+        raise SchemaVersionError(
+            f"This bundle uses schema version {version}, but this runtime only "
+            f"understands up to version {SCHEMA_VERSION}. Upgrade genesis-forge-runtime."
+        )
+    if version < MIN_SUPPORTED_SCHEMA_VERSION:
+        raise SchemaVersionError(
+            f"This bundle uses schema version {version}, which is older than the "
+            f"oldest version this runtime supports ({MIN_SUPPORTED_SCHEMA_VERSION}). "
+            f"Re-export it with a current version of Genesis Forge."
+        )
+
+
+def _check_action_slices(actions: tuple[ActionManagerSpec, ...]) -> None:
+    """Action manager slices must tile the policy output with no gaps or overlaps."""
+    if not actions:
+        return
+    ordered = sorted(actions, key=lambda spec: spec.slice_start)
+    cursor = 0
+    for spec in ordered:
+        if spec.slice_start != cursor:
+            raise MalformedBundleError(
+                f"Action manager '{spec.name}' starts at index {spec.slice_start} but the "
+                f"previous manager ended at {cursor}; slices must tile the policy output "
+                f"without gaps or overlaps."
+            )
+        if spec.slice_end <= spec.slice_start:
+            raise MalformedBundleError(
+                f"Action manager '{spec.name}' has an empty slice "
+                f"[{spec.slice_start}, {spec.slice_end}]."
+            )
+        cursor = spec.slice_end

genesis_forge_runtime/observation_schema.py

@@ -0,0 +1,131 @@
+"""How the policy's observation vector is laid out.
+
+The half of the manifest that :mod:`genesis_forge_runtime.observations` consumes:
+what each slot holds, how wide it is, what it is clipped to, and what it is scaled by.
+
+Every slot is an input you supply each tick. Most come from sensors; some echo the
+policy's own previous output, which you read off the processor. The bundle does not
+distinguish them -- from the assembler's side they are passed in the same way, and
+the deployment guide covers which of your entries is which.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .errors import MalformedBundleError
+from .serialization import require
+
+
+@dataclass(frozen=True)
+class ObservationEntry:
+    """One named slot in the policy's observation vector."""
+
+    name: str
+    size: int
+    scale: float = 1.0
+    clip: tuple[float, float] | None = None
+    description: str | None = None
+    units: str | None = None
+
+    def describe(self) -> str:
+        """One-line summary of what to supply for this entry."""
+        parts = [f"{self.name} ({self.size} value{'s' if self.size != 1 else ''})"]
+        if self.units:
+            parts.append(f"in {self.units}")
+        # No scale or clip here: this listing is what the caller supplies, and the
+        # caller supplies raw readings. The assembler clips and scales them afterwards.
+        summary = ", ".join(parts)
+        if self.description:
+            summary = f"{summary} -- {self.description}"
+        return summary
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any], *, where: str) -> ObservationEntry:
+        name = require(data, "name", where=where)
+        clip = data.get("clip")
+        entry = cls(
+            name=name,
+            size=int(require(data, "size", where=f"{where}.{name}")),
+            scale=float(data.get("scale", 1.0)),
+            clip=None if clip is None else (float(clip[0]), float(clip[1])),
+            description=data.get("description"),
+            units=data.get("units"),
+        )
+        return entry
+
+    def to_dict(self) -> dict[str, Any]:
+        data: dict[str, Any] = {
+            "name": self.name,
+            "size": self.size,
+            "scale": self.scale,
+        }
+        if self.clip is not None:
+            data["clip"] = [float(self.clip[0]), float(self.clip[1])]
+        if self.description is not None:
+            data["description"] = self.description
+        if self.units is not None:
+            data["units"] = self.units
+        return data
+
+
+@dataclass(frozen=True)
+class ObservationLayout:
+    """Ordering, scaling, and history configuration of the policy's input vector."""
+
+    entries: tuple[ObservationEntry, ...]
+    history_length: int = 1
+
+    @property
+    def single_size(self) -> int:
+        """Width of one observation tick, before history stacking."""
+        return sum(entry.size for entry in self.entries)
+
+    @property
+    def total_size(self) -> int:
+        """Width of the full vector handed to the policy."""
+        return self.single_size * self.history_length
+
+    def entry(self, name: str) -> ObservationEntry:
+        for entry in self.entries:
+            if entry.name == name:
+                return entry
+        known = ", ".join(item.name for item in self.entries)
+        raise KeyError(f"No observation entry named '{name}'. Known entries: {known}.")
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ObservationLayout:
+        where = "observations"
+        raw_entries = require(data, "entries", where=where)
+        if not isinstance(raw_entries, list) or not raw_entries:
+            raise MalformedBundleError(
+                "'observations.entries' must be a non-empty list of observation entries."
+            )
+        entries = tuple(
+            ObservationEntry.from_dict(item, where=f"{where}.entries")
+            for item in raw_entries
+        )
+        names = [entry.name for entry in entries]
+        duplicates = {name for name in names if names.count(name) > 1}
+        if duplicates:
+            raise MalformedBundleError(
+                f"Duplicate observation entry names: {', '.join(sorted(duplicates))}."
+            )
+        history_length = int(data.get("history_length", 1))
+        if history_length < 1:
+            raise MalformedBundleError(
+                f"'observations.history_length' must be at least 1, got {history_length}."
+            )
+        return cls(
+            entries=entries,
+            history_length=history_length,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "history_length": self.history_length,
+            "single_size": self.single_size,
+            "total_size": self.total_size,
+            "entries": [entry.to_dict() for entry in self.entries],
+        }