simforge-gym 0.2.0__py3-none-any.whl

simforge_gym/__init__.py

@@ -0,0 +1,58 @@
+"""Gymnasium environments and the runtime-protocol client for the SimForge runtime."""
+
+import importlib.metadata
+
+from .env import SimForgeEnv, action_space_for, sensor_space, state_space
+from .rewards import EventPenalty, WeightedTerms, native
+from .batch import SimForgeBatchVectorEnv
+from .runtime import BatchResults, BatchRow, CommandResult, RuntimeBatch, RuntimeClient, RuntimeProtocolError
+from .datasets import DatasetError, dataset_scenarios, pull_dataset, read_dataset
+from .scenarios import MapPin, RuntimeSessions, Scenario, ScenarioServers, ScenarioSet, ScenarioSetEnv, bundled_sets
+from .server import EnvServer, template_workspace
+from .world import Subscription, TruthDelivery, WorldCheckpoint, WorldSession
+from .transport import EnvTransport, RuntimeEnvTransport, SocketTransport
+from .vector import make_vector_env
+from .wire import PROTOCOL, EnvServeError, ProtocolError
+
+__version__ = importlib.metadata.version("simforge-gym")
+
+__all__ = [
+    "PROTOCOL",
+    "BatchResults",
+    "BatchRow",
+    "CommandResult",
+    "RuntimeBatch",
+    "SimForgeBatchVectorEnv",
+    "Subscription",
+    "TruthDelivery",
+    "WorldCheckpoint",
+    "WorldSession",
+    "EnvServeError",
+    "EnvServer",
+    "EnvTransport",
+    "DatasetError",
+    "EventPenalty",
+    "MapPin",
+    "ProtocolError",
+    "RuntimeClient",
+    "RuntimeEnvTransport",
+    "RuntimeProtocolError",
+    "RuntimeSessions",
+    "Scenario",
+    "ScenarioServers",
+    "ScenarioSet",
+    "ScenarioSetEnv",
+    "SimForgeEnv",
+    "SocketTransport",
+    "WeightedTerms",
+    "action_space_for",
+    "bundled_sets",
+    "dataset_scenarios",
+    "make_vector_env",
+    "native",
+    "pull_dataset",
+    "read_dataset",
+    "sensor_space",
+    "state_space",
+    "template_workspace",
+]

simforge_gym/batch.py

@@ -0,0 +1,138 @@
+"""Gymnasium VectorEnv over one runtime ``env.batch`` (state observations, in-runtime CPU parallelism)."""
+
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+import gymnasium as gym
+import numpy as np
+from gymnasium.vector import AutoresetMode
+from gymnasium.vector.utils import batch_space
+
+from .env import ActionMode, action_space_for, encode_row, state_space, unpack_state
+from .rewards import RewardFn, native
+from .runtime import BatchResults, RuntimeBatch, RuntimeClient
+from .scenarios import SEED_SPACE
+from .wire import PROTOCOL, EnvServeError, ProtocolError
+
+
+class SimForgeBatchVectorEnv(gym.vector.VectorEnv[Any, Any, Any]):
+    """``num_envs`` episodes of one workspace; an ended env is reset (masked) on the next step."""
+
+    metadata: dict[str, Any] = {"autoreset_mode": AutoresetMode.NEXT_STEP}
+
+    def __init__(
+        self,
+        runtime: RuntimeClient,
+        argv: Sequence[str],
+        num_envs: int,
+        *,
+        shards: int | None = None,
+        cwd: str | None = None,
+        action_mode: ActionMode = "setpoint",
+        reward_fn: RewardFn = native,
+        info_channel: bool = True,
+    ) -> None:
+        self.runtime = runtime
+        self.batch: RuntimeBatch | None = runtime.batch_create(argv, num_envs, shards=shards, cwd=cwd)
+        try:
+            self.hello = self.batch.hello
+            if self.hello.get("protocol") != PROTOCOL:
+                raise ProtocolError(f"env.batch hello speaks {self.hello.get('protocol')!r}, not {PROTOCOL}")
+            self.num_envs = num_envs
+            self.action_mode = action_mode
+            self.reward_fn = reward_fn
+            self.info_channel = info_channel
+            self._slot = {name: index for index, name in enumerate(self.hello["actionFields"])}
+            self._width = int(self.hello["actionWidth"])
+            self._state = dict(state_space(self.hello).spaces)
+            self.single_observation_space = gym.spaces.Dict(self._state)
+            self.single_action_space = action_space_for(action_mode)
+            self.observation_space = batch_space(self.single_observation_space, num_envs)
+            self.action_space = batch_space(self.single_action_space, num_envs)
+        except BaseException:
+            self.close()
+            raise
+        self._ended = np.zeros(num_envs, dtype=bool)
+        self._rng = np.random.default_rng()
+
+    def _rows(self, op: str, results: BatchResults) -> None:
+        if results.failed:
+            codes = {k: results.rows[k].header.get("code") for k in results.failed}  # type: ignore[union-attr]
+            error = EnvServeError(op, "batch_rows_failed", f"envs {list(results.failed)} failed: {codes}", codes)
+            error.results = results  # type: ignore[attr-defined]
+            raise error
+
+    def _seeds(self, mask: np.ndarray) -> list[int | None]:
+        return [int(self._rng.integers(SEED_SPACE)) if m else None for m in mask]
+
+    def reset(self, *, seed: int | list[int] | None = None, options: dict[str, Any] | None = None) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
+        """Reset every env; ``seed`` reseeds the episode-seed stream (a list gives each env's seed)."""
+        if self.batch is None:
+            raise RuntimeError("the batch is closed")
+        if options:
+            raise ValueError(f"SimForgeBatchVectorEnv.reset takes no options, got {sorted(options)}")
+        if isinstance(seed, list):
+            if len(seed) != self.num_envs:
+                raise ValueError(f"{len(seed)} seeds for {self.num_envs} envs")
+            seeds: list[int | None] = [int(s) for s in seed]
+            self._rng = np.random.default_rng(seeds)
+        else:
+            if seed is not None:
+                self._rng = np.random.default_rng(seed)
+            seeds = self._seeds(np.ones(self.num_envs, dtype=bool))
+        results = self.runtime.batch_reset(self.batch, seeds)
+        self._rows("reset", results)
+        self._ended[:] = False
+        obs, infos = self._collect(results.rows)
+        infos["episode_seed"] = np.array(seeds)
+        return obs, infos
+
+    def step(self, actions: Any) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
+        if self.batch is None:
+            raise RuntimeError("the batch is closed")
+        actions = np.asarray(actions, dtype=np.float64)
+        stepping = ~self._ended
+        rows = [list(encode_row(self.action_mode, self._slot, actions[k], np.empty(self._width))) if stepping[k] else None for k in range(self.num_envs)]
+        merged: list[Any] = [None] * self.num_envs
+        if stepping.any():
+            stepped = self.runtime.batch_step(self.batch, rows, mask=list(stepping))
+            self._rows("step", stepped)
+            merged = list(stepped.rows)
+        seeds: list[int | None] = [None] * self.num_envs
+        if self._ended.any():
+            seeds = self._seeds(self._ended)
+            restarted = self.runtime.batch_reset(self.batch, seeds, mask=list(self._ended))
+            self._rows("reset", restarted)
+            merged = [restarted.rows[k] if self._ended[k] else merged[k] for k in range(self.num_envs)]
+        obs, infos = self._collect(merged)
+        rewards = np.zeros(self.num_envs)
+        terminated = np.zeros(self.num_envs, dtype=bool)
+        truncated = np.zeros(self.num_envs, dtype=bool)
+        for k in np.flatnonzero(stepping):
+            header = merged[k].header
+            info_k = {key: value[k] for key, value in infos.items() if key != "episode_seed"}
+            rewards[k] = float(self.reward_fn(float(header["reward"]), info_k))
+            terminated[k], truncated[k] = bool(header["terminated"]), bool(header["truncated"])
+        infos["episode_seed"] = np.array(seeds, dtype=object)
+        infos["native_reward"] = np.array([float(merged[k].header["reward"]) if stepping[k] else 0.0 for k in range(self.num_envs)])
+        self._ended = terminated | truncated
+        return obs, rewards, terminated, truncated, infos
+
+    def _collect(self, rows: Sequence[Any]) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
+        observations, infos = [], []
+        for row in rows:
+            if row is None:
+                raise ProtocolError("a batch row is missing")
+            state, info = unpack_state(row.header, row.tail, self._state, self.info_channel)
+            observations.append(state)
+            infos.append(info)
+        obs = {name: np.stack([o[name] for o in observations]) for name in self._state}
+        keys = set().union(*(i.keys() for i in infos))
+        info = {key: np.array([i.get(key) for i in infos], dtype=object) for key in keys}
+        return obs, info
+
+    def close_extras(self, **kwargs: Any) -> None:
+        batch, self.batch = self.batch, None
+        if batch is not None and self.runtime.alive:
+            self.runtime.batch_close(batch)

simforge_gym/datasets.py

@@ -0,0 +1,238 @@
+"""SimForge datasets as scenario-set tasks.
+
+Two folder formats, both with ``manifest.json``, ``scenarios.jsonl`` (one row per scenario, with
+its split, map and files) and ``SHA256SUMS``:
+
+- ``simforge.dataset-files/v1``, a hosted dataset pulled as files (``simforge datasets pull``,
+  which checks every file): each scenario's document and the record of the simulation that
+  produced its trace become a workspace whose episode is exactly the input that simulation ran.
+- ``simforge.public-dataset/v1``, the public SimForge dataset: each scenario is a
+  ``document.json`` and the recordings it pins, laid out as a workspace the runtime compiles.
+
+Either way an episode runs on the installed map whose OpenDRIVE digest the scenario names.
+"""
+
+from __future__ import annotations
+
+import gzip
+import hashlib
+import json
+import os
+import re
+import shutil
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from .server import find_binary
+
+FILES_SCHEMA = "simforge.dataset-files/v1"
+PUBLIC_SCHEMA = "simforge.public-dataset/v1"
+DATASET_ID = re.compile(r"^usds_[0-9A-Za-z_]+$")
+MARKER = ".simforge-gym-package"
+
+
+class DatasetError(ValueError):
+    """A named refusal: ``code`` is stable, ``reason`` says what and where."""
+
+    def __init__(self, code: str, reason: str) -> None:
+        super().__init__(f"{code}: {reason}")
+        self.code = code
+        self.reason = reason
+
+
+def cache_root() -> Path:
+    """``$XDG_CACHE_HOME/simforge-gym``, else ``~/.cache/simforge-gym``."""
+    base = os.environ.get("XDG_CACHE_HOME")
+    return (Path(base) if base and os.path.isabs(base) else Path.home() / ".cache") / "simforge-gym"
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1 << 20), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _relative(folder: Path, path: str) -> Path:
+    """A path a row or SHA256SUMS names: relative, ``/``-separated, plain names."""
+    parts = path.split("/")
+    if path.startswith("/") or any(part in ("", ".", "..") or "\\" in part for part in parts):
+        raise DatasetError("dataset_path_invalid", f"{path!r} is not a relative path")
+    return folder.joinpath(*parts)
+
+
+def pull_dataset(dataset_id: str, out: str | os.PathLike[str], *, binary: str | os.PathLike[str] | None = None) -> Path:
+    """Download hosted dataset ``dataset_id``'s scenarios as files with ``simforge datasets pull``
+    (signed in; it resumes and checks every file)."""
+    if not DATASET_ID.match(dataset_id):
+        raise DatasetError("dataset_id_invalid", f"{dataset_id!r} is not a dataset id (usds_...)")
+    folder = Path(out)
+    completed = subprocess.run(
+        [find_binary(binary), "datasets", "pull", dataset_id, "--include", "scenarios", "--out", os.fspath(folder), "--json"],
+        capture_output=True,
+        text=True,
+    )
+    if completed.returncode != 0:
+        lines = completed.stderr.strip().splitlines()
+        try:
+            error: Any = json.loads(lines[-1]) if lines else {}
+        except json.JSONDecodeError:
+            error = {}
+        if isinstance(error, dict) and "code" in error:
+            raise DatasetError(str(error["code"]), str(error.get("reason", "")))
+        raise DatasetError("dataset_pull_failed", f"simforge datasets pull exited {completed.returncode}: {completed.stderr.strip()}")
+    return folder
+
+
+def _listed(folder: Path, manifest: dict[str, Any]) -> dict[str, str]:
+    """SHA256SUMS by path, after checking it is the file the manifest names."""
+    sums = folder / "SHA256SUMS"
+    if not sums.is_file() or _sha256(sums) != manifest["files"]["SHA256SUMS"]["sha256"]:
+        raise DatasetError("dataset_digest_mismatch", f"{sums} is missing or not the SHA256SUMS the manifest names")
+    listed = {}
+    for line in sums.read_text().splitlines():
+        if line:
+            digest, _, path = line.partition("  ")
+            listed[path] = digest
+    return listed
+
+
+def read_dataset(root: str | os.PathLike[str]) -> tuple[dict[str, Any], list[dict[str, Any]]]:
+    """The manifest and rows of a dataset folder; ``scenarios.jsonl`` must be the file the manifest names."""
+    folder = Path(root)
+    manifest_path, rows_path = folder / "manifest.json", folder / "scenarios.jsonl"
+    for path in (manifest_path, rows_path):
+        if not path.is_file():
+            raise DatasetError("dataset_file_missing", f"{path} is missing: {folder} is not a pulled dataset folder")
+    manifest = json.loads(manifest_path.read_bytes())
+    schema = manifest.get("schema") if isinstance(manifest, dict) else None
+    if schema not in (FILES_SCHEMA, PUBLIC_SCHEMA):
+        raise DatasetError("dataset_schema_unsupported", f"{manifest_path} is neither {FILES_SCHEMA} nor {PUBLIC_SCHEMA}")
+    listed = _listed(folder, manifest)
+    if "scenarios.jsonl" not in listed:
+        raise DatasetError("dataset_file_missing", "SHA256SUMS does not list scenarios.jsonl")
+    want = listed["scenarios.jsonl"]
+    if schema == FILES_SCHEMA and manifest["files"]["scenarios.jsonl"]["sha256"] != want:
+        raise DatasetError("dataset_digest_mismatch", "the manifest and SHA256SUMS name different scenarios.jsonl files")
+    if _sha256(rows_path) != want:
+        raise DatasetError("dataset_digest_mismatch", f"{rows_path} is not the scenarios.jsonl the manifest names (sha256 {want})")
+    rows = [json.loads(line) for line in rows_path.read_text().splitlines() if line.strip()]
+    if len(rows) != manifest["counts"]["scenarios"]:
+        raise DatasetError("dataset_rows_count", f"{len(rows)} rows, the manifest counts {manifest['counts']['scenarios']}")
+    return manifest, rows
+
+
+def _checked(folder: Path, path: str, sha256: str | None, listed: dict[str, str]) -> Path:
+    """A file a row names, present with the digest the row and SHA256SUMS record."""
+    if sha256 is None or listed.get(path) != sha256:
+        raise DatasetError("dataset_row_unlisted", f"{path} is not listed in SHA256SUMS with the digest its row names")
+    source = _relative(folder, path)
+    if not source.is_file():
+        raise DatasetError("dataset_file_missing", f"{source} is missing")
+    if _sha256(source) != sha256:
+        raise DatasetError("dataset_digest_mismatch", f"{source} is not the file its row names (sha256 {sha256})")
+    return source
+
+
+def _laid_out(dest: Path, key: str) -> Path | None:
+    done = dest / MARKER
+    return dest if done.is_file() and done.read_text() == key else None
+
+
+def _install(partial: Path, dest: Path, key: str, manifest: dict[str, Any]) -> Path:
+    (partial / "manifest.json").write_text(json.dumps(manifest, sort_keys=True))
+    (partial / MARKER).write_text(key)
+    shutil.rmtree(dest, ignore_errors=True)
+    partial.rename(dest)
+    return dest
+
+
+def _files_workspace(folder: Path, row: dict[str, Any], listed: dict[str, str], dest: Path) -> Path:
+    """A pulled scenario as a package workspace whose episode is the input its simulation ran, once per scenario."""
+    key = row["id"]
+    if _laid_out(dest, key):
+        return dest
+    map_id, xodr_sha256 = row["map"]["sourceMapId"], row["map"]["xodrSha256"]
+    if map_id is None or xodr_sha256 is None:
+        raise DatasetError("dataset_map_unnamed", f"scenario {key} names no installable map (sourceMapId {map_id!r}, xodrSha256 {xodr_sha256!r})")
+    members = {m["role"]: m for m in row["members"]}
+    for role in ("document", "resolution"):
+        if role not in members:
+            raise DatasetError("dataset_row_unlisted", f"scenario {key} lists no {role}")
+    document = _checked(folder, members["document"]["path"], members["document"]["sha256"], listed)
+    resolution = _checked(folder, members["resolution"]["path"], members["resolution"]["sha256"], listed)
+    record = json.loads(gzip.decompress(resolution.read_bytes()))
+    if "resolvedInput" not in record:
+        raise DatasetError("dataset_input_missing", f"scenario {key}'s simulation record ({resolution.name}) holds no resolved input")
+    expected = row["simulation"]["resolvedInputDigest"]
+    if record.get("resolvedInputDigest") != expected:
+        raise DatasetError("dataset_input_mismatch", f"scenario {key}'s simulation record is of input {record.get('resolvedInputDigest')}, not the row's {expected}")
+    partial = dest.with_name(dest.name + ".partial")
+    shutil.rmtree(partial, ignore_errors=True)
+    (partial / "simulation").mkdir(parents=True)
+    shutil.copyfile(document, partial / "document.json")
+    shutil.copyfile(resolution, partial / "simulation" / "resolution.json.gz")
+    manifest = {"schema": "simforge.scenario-package/v1", "map": {"xodrSha256": xodr_sha256, "sourceMapId": map_id},
+                "timelines": [], "recordings": []}
+    return _install(partial, dest, key, manifest)
+
+
+def _public_workspace(folder: Path, row: dict[str, Any], listed: dict[str, str], dest: Path) -> Path:
+    """A public scenario's document and recordings as a package workspace on its map, once per scenario."""
+    key = row["id"]
+    if _laid_out(dest, key):
+        return dest
+    partial = dest.with_name(dest.name + ".partial")
+    shutil.rmtree(partial, ignore_errors=True)
+    recordings = []
+    for entry in row["files"]:
+        if entry["role"] not in ("document", "recording"):
+            continue
+        source = _checked(folder, entry["path"], entry["sha256"], listed)
+        name = "document.json" if entry["role"] == "document" else f"recordings/{source.name}"
+        (partial / name).parent.mkdir(parents=True, exist_ok=True)
+        shutil.copyfile(source, partial / name)
+        if entry["role"] == "recording":
+            recordings.append({"sha256": source.name.split(".", 1)[0], "path": name})
+    if not (partial / "document.json").is_file():
+        raise DatasetError("dataset_row_unlisted", f"scenario {key} lists no document")
+    manifest = {"schema": "simforge.scenario-package/v1", "map": {"xodrSha256": row["map"]["xodrSha256"], "sourceMapId": row["map"]["slug"]},
+                "timelines": [], "recordings": recordings}
+    return _install(partial, dest, key, manifest)
+
+
+def dataset_scenarios(
+    source: str | os.PathLike[str],
+    *,
+    split: str | None = "train",
+    map_name: str | None = None,
+    cache_dir: str | os.PathLike[str] | None = None,
+    binary: str | os.PathLike[str] | None = None,
+) -> list[Any]:
+    """One workspace :class:`~simforge_gym.scenarios.Scenario` per scenario of the dataset's ``split`` (all with ``None``).
+
+    ``source`` is a dataset folder, or a hosted dataset id (``usds_...``) pulled into the cache
+    first. ``map_name`` keeps the scenarios on that installed map id. Workspaces are laid out
+    under ``cache_dir`` (default :func:`cache_root`), keyed by the dataset's rows and the scenario.
+    """
+    from .scenarios import Scenario
+
+    cache = Path(cache_dir) if cache_dir is not None else cache_root()
+    folder = Path(source)
+    if not folder.exists():
+        if not DATASET_ID.match(os.fspath(source)):
+            raise DatasetError("dataset_file_missing", f"{source} is neither a dataset folder nor a dataset id (usds_...)")
+        folder = pull_dataset(os.fspath(source), cache / "datasets" / os.fspath(source), binary=binary)
+    manifest, rows = read_dataset(folder)
+    listed = _listed(folder, manifest)
+    public = manifest["schema"] == PUBLIC_SCHEMA
+    map_of = (lambda r: r["map"]["slug"]) if public else (lambda r: r["map"]["sourceMapId"])
+    chosen = [r for r in rows if (split is None or r["split"] == split) and (map_name is None or map_of(r) == map_name)]
+    if not chosen:
+        splits = manifest["counts"]["bySplit"] if public else manifest["splits"]
+        raise DatasetError("dataset_split_empty", f"no scenarios of split {split!r} on map {map_name!r}; the dataset's splits: {splits}")
+    root = cache / "workspaces" / listed["scenarios.jsonl"][:16]
+    layout = _public_workspace if public else _files_workspace
+    return [Scenario(row["id"], workspace=os.fspath(layout(folder, row, listed, root / row["id"]))) for row in sorted(chosen, key=lambda r: r["id"])]

simforge_gym/env.py

@@ -0,0 +1,229 @@
+"""Gymnasium environment over one SimForge episode session."""
+
+from __future__ import annotations
+
+import math
+from typing import Any, Literal, Mapping
+
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+
+from . import wire
+from .rewards import RewardFn, native
+from .transport import EnvTransport
+
+ActionMode = Literal["setpoint", "control"]
+ObservationMode = Literal["state", "sensors", "state+sensors"]
+CAMERA_PASSES = {"rgb": 4, "id": 4, "semantic": 4}
+#: The runtime's render warning for a step at which a camera's host actor is not in the world (the ego left the map).
+SENSOR_HOST_ABSENT = "sensor_host_absent"
+
+
+def action_space_for(mode: ActionMode) -> spaces.Box:
+    if mode == "setpoint":
+        return spaces.Box(np.array([0.0, -np.inf]), np.array([np.inf, np.inf]), (2,), np.float64)
+    if mode == "control":
+        return spaces.Box(np.array([0.0, 0.0, -1.0]), np.array([1.0, 1.0, 1.0]), (3,), np.float64)
+    raise ValueError(f"unknown action_mode {mode!r}; expected 'setpoint' or 'control'")
+
+
+def state_space(hello: Mapping[str, Any]) -> spaces.Dict:
+    """Privileged engine state: ego ``state_vector``, ``objects`` rows, optional ``bev``."""
+    max_objects, features = int(hello["maxObjects"]), int(hello["objectFeatures"])
+    members: dict[str, spaces.Space[Any]] = {}
+    if hello.get("stateVectorEnabled", True):
+        members["state_vector"] = spaces.Box(-np.inf, np.inf, (int(hello["stateVectorSize"]),), np.float64)
+    low = np.tile(np.array([0.0, -math.pi, -np.inf, 0.0, 0.0], dtype=np.float32), (max_objects, 1))
+    high = np.tile(np.array([np.inf, math.pi, np.inf, 1.0, 1.0], dtype=np.float32), (max_objects, 1))
+    members["objects"] = spaces.Box(low[:, :features], high[:, :features], (max_objects, features), np.float32)
+    if hello.get("bevShape"):
+        members["bev"] = spaces.Box(-np.inf, np.inf, tuple(hello["bevShape"]), np.float32)
+    return spaces.Dict(members)
+
+
+def sensor_space(hello: Mapping[str, Any]) -> spaces.Dict:
+    """Rendered camera passes, keyed ``"<sourceId>/<pass>"``; non-camera payloads stay in ``info``."""
+    members: dict[str, spaces.Space[Any]] = {}
+    for source in hello.get("sensors", []):
+        if source.get("modality") != "rgb":
+            continue
+        h, w = int(source["height"]), int(source["width"])
+        for name in source.get("passes", []):
+            if name == "depth":
+                members[f"{source['sourceId']}/depth"] = spaces.Box(-np.inf, np.inf, (h, w), np.float32)
+            elif name in CAMERA_PASSES:
+                members[f"{source['sourceId']}/{name}"] = spaces.Box(0, 255, (h, w, CAMERA_PASSES[name]), np.uint8)
+    if not members:
+        raise ValueError("observation mode needs rendered cameras: serve with --rig (not --no-sensors)")
+    return spaces.Dict(members)
+
+
+class SimForgeEnv(gym.Env[dict[str, Any], np.ndarray]):
+    """Reset/step one SimForge episode; ``terminated``/``truncated`` are the runtime's.
+
+    A step at which a camera's host actor is not in the world (the ego drove off the map) has no rendered
+    frame: the camera observation is the last rendered one, ``info["render_warnings"]`` carries the runtime's
+    ``sensor_host_absent`` warning, and the step is ``truncated``.
+    """
+
+    metadata: dict[str, Any] = {"render_modes": []}
+
+    def __init__(
+        self,
+        transport: EnvTransport,
+        *,
+        action_mode: ActionMode = "setpoint",
+        observation: ObservationMode = "state",
+        reward_fn: RewardFn = native,
+        info_channel: bool = True,
+        seed: int | str | None = None,
+    ) -> None:
+        super().__init__()
+        self.transport = transport
+        try:
+            self._configure(action_mode, observation)
+        except BaseException:
+            self.close()
+            raise
+        self.reward_fn = reward_fn
+        self.info_channel = info_channel
+        self._default_seed = seed
+        self._last_cameras: dict[str, np.ndarray] | None = None
+
+    def _configure(self, action_mode: ActionMode, observation: ObservationMode) -> None:
+        self.hello = self.transport.hello()
+        self.action_mode: ActionMode = action_mode
+        self.observation_mode: ObservationMode = observation
+        self._slot = {name: index for index, name in enumerate(self.hello["actionFields"])}
+        self._row = np.empty(int(self.hello["actionWidth"]), dtype=np.float64)
+        self.ego: str = self.hello["ego"]
+        self.decision_hz = int(self.hello["decisionHz"])
+        self.action_space = action_space_for(action_mode)
+        if observation not in ("state", "sensors", "state+sensors"):
+            raise ValueError(f"unknown observation mode {observation!r}")
+        self._state_space = dict(state_space(self.hello).spaces) if observation != "sensors" else {}
+        self._camera_space = dict(sensor_space(self.hello).spaces) if observation != "state" else {}
+        if observation == "state":
+            self.observation_space = spaces.Dict(self._state_space)
+        elif observation == "sensors":
+            self.observation_space = spaces.Dict(self._camera_space)
+        else:
+            self.observation_space = spaces.Dict({**self._state_space, "sensors": spaces.Dict(self._camera_space)})
+
+    def encode_action(self, action: Any) -> np.ndarray:
+        """``action`` in the mode's Box as the native row (NaN = unset)."""
+        return encode_row(self.action_mode, self._slot, action, self._row)
+
+    def _unpack(self, header: Mapping[str, Any], tail: bytes, *, restart: bool = False) -> tuple[dict[str, Any], dict[str, Any]]:
+        if restart:
+            self._last_cameras = None
+        if self.observation_mode == "sensors":
+            # A sensor policy sees no engine state or labels, whatever profile the runtime serves.
+            state, info = {}, {"t_s": header["info"]["t_s"]}
+        else:
+            state, info = unpack_state(header, tail, self._state_space, self.info_channel)
+        frames: dict[str, dict[str, np.ndarray]] = {}
+        for source, passes in header.get("sensors", {}).items():
+            frames[source] = {name: wire.array(tail, ref) for name, ref in passes.items()}
+        info["sensors"] = frames
+        if "renderMs" in header:
+            info["render_ms"] = header["renderMs"]
+        warnings = list(header.get("renderWarnings") or ())
+        if warnings:
+            info["render_warnings"] = warnings
+        if self.observation_mode == "state":
+            return state, info
+        if host_absent(header):
+            # No frame exists at this instant; the episode ends here (see `step`) on the last rendered one.
+            if self._last_cameras is None:
+                raise wire.ProtocolError(f"the runtime rendered no camera: {_absent_message(header)}, and no earlier frame exists")
+            cameras = self._last_cameras
+        else:
+            cameras = {}
+            for key, space in self._camera_space.items():
+                source, name = key.rsplit("/", 1)
+                cameras[key] = _declared(key, frames.get(source, {}).get(name), space)
+            self._last_cameras = cameras
+        if self.observation_mode == "sensors":
+            return cameras, info
+        return {**state, "sensors": cameras}, info
+
+    def reset(self, *, seed: int | None = None, options: Mapping[str, Any] | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
+        """``options["episode_seed"]`` overrides the runtime seed; ``options["scenario"]`` switches input."""
+        super().reset(seed=seed)
+        options = options or {}
+        scenario = options.get("scenario")
+        if scenario is not None and not isinstance(scenario, Mapping):
+            raise TypeError("reset options['scenario'] must be a compiled scenario input mapping")
+        episode_seed = options.get("episode_seed", seed if seed is not None else self._default_seed)
+        return self._unpack(*self.transport.reset(episode_seed, scenario=scenario), restart=True)
+
+    def step(self, action: Any) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
+        """Apply ``action`` for one decision; ``None`` keeps the authored control."""
+        header, tail = self.transport.step(None if action is None else self.encode_action(action))
+        observation, info = self._unpack(header, tail)
+        native_reward = float(header["reward"])
+        info["native_reward"] = native_reward
+        reward = float(self.reward_fn(native_reward, info))
+        terminated = bool(header["terminated"])
+        # A camera policy has nothing to act on once its host left the world: the episode ends there.
+        truncated = bool(header["truncated"]) or (not terminated and self.observation_mode != "state" and host_absent(header))
+        return observation, reward, terminated, truncated, info
+
+    def observe(self) -> tuple[dict[str, Any], dict[str, Any]]:
+        return self._unpack(*self.transport.observe())
+
+    def checkpoint(self) -> bytes:
+        return self.transport.checkpoint()
+
+    def restore(self, checkpoint: bytes) -> tuple[dict[str, Any], dict[str, Any]]:
+        return self._unpack(*self.transport.restore(checkpoint), restart=True)
+
+    def close(self) -> None:
+        transport = self.__dict__.pop("transport", None)
+        if transport is not None:
+            transport.close()
+
+
+def host_absent(header: Mapping[str, Any]) -> bool:
+    """The runtime rendered no camera at this step: a camera's host actor is not in the world."""
+    return any(warning.get("code") == SENSOR_HOST_ABSENT for warning in header.get("renderWarnings") or ())
+
+
+def _absent_message(header: Mapping[str, Any]) -> str:
+    return "; ".join(str(w.get("message") or w.get("actorIds")) for w in header.get("renderWarnings") or () if w.get("code") == SENSOR_HOST_ABSENT)
+
+
+def _declared(name: str, value: np.ndarray | None, space: spaces.Space[Any]) -> np.ndarray:
+    """``value`` if it has the declared shape and dtype; a runtime that breaks the space is a protocol error."""
+    if value is None:
+        raise wire.ProtocolError(f"the runtime sent no {name} observation")
+    if value.shape != space.shape or value.dtype != space.dtype:
+        raise wire.ProtocolError(f"{name} is {value.dtype}{list(value.shape)}, but the space declares {space.dtype}{list(space.shape or ())}")
+    return value
+
+
+def unpack_state(header: Mapping[str, Any], tail: bytes, state: Mapping[str, spaces.Space[Any]], info_channel: bool) -> tuple[dict[str, Any], dict[str, Any]]:
+    """The declared state arrays and the engine info of one native header."""
+    observation = header["observation"]
+    values = {name: _declared(name, wire.array(tail, observation[name]) if name in observation else None, space) for name, space in state.items()}
+    server_info = dict(header["info"])
+    info: dict[str, Any] = {k: server_info[k] for k in ("t_s", "ego", "object_ids", "reward_terms")}
+    if info_channel:
+        info.update({k: v for k, v in server_info.items() if k not in info})
+    return values, info
+
+
+def encode_row(mode: ActionMode, slots: Mapping[str, int], action: Any, row: np.ndarray) -> np.ndarray:
+    """Write ``action`` into the native ``row`` by field slot; every other field is NaN (unset)."""
+    values = np.asarray(action, dtype=np.float64).reshape(-1)
+    names = ("target_speed_mps", "target_acceleration_mps2") if mode == "setpoint" else ("throttle", "brake", "steer")
+    if values.shape != (len(names),):
+        raise ValueError(f"{mode} action must have shape ({len(names)},), got {values.shape}")
+    if not np.all(np.isfinite(values)):
+        raise ValueError(f"action contains non-finite values: {values}")
+    row.fill(np.nan)
+    for name, value in zip(names, values):
+        row[slots[name]] = value
+    return row

simforge_gym/examples/__init__.py

@@ -0,0 +1 @@
+"""Runnable examples: ``python -m simforge_gym.examples.ppo``."""