diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 043f8f5b3..89b62aed3 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -4871,14 +4871,6 @@ } ], "./positronic/policy/codec.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 16, - "lineCount": 1 - } - }, { "code": "reportIncompatibleMethodOverride", "range": { @@ -5032,16 +5024,6 @@ } } ], - "./positronic/policy/tests/test_layers.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 30, - "endColumn": 40, - "lineCount": 1 - } - } - ], "./positronic/policy/tests/test_policy_io.py": [ { "code": "reportArgumentType", @@ -5075,126 +5057,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalIterable", - "range": { - "startColumn": 28, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 17, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 37, - "lineCount": 1 - } - }, - { - "code": "reportCallIssue", - "range": { - "startColumn": 11, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 11, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportOptionalIterable", - "range": { - "startColumn": 28, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 17, - "lineCount": 1 - } - }, { "code": "reportCallIssue", "range": { @@ -5211,22 +5073,6 @@ "lineCount": 1 } }, - { - "code": "reportCallIssue", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -5340,46 +5186,6 @@ "endColumn": 58, "lineCount": 1 } - }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 17, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 19, - "endColumn": 24, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 30, - "endColumn": 40, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 30, - "endColumn": 50, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 30, - "endColumn": 50, - "lineCount": 1 - } } ], "./positronic/probe.py": [ @@ -5422,22 +5228,6 @@ "endColumn": 28, "lineCount": 1 } - }, - { - "code": "reportOptionalContextManager", - "range": { - "startColumn": 13, - "endColumn": 23, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 42, - "lineCount": 1 - } } ], "./positronic/replay_record.py": [ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 621a40dab..ba4aa3894 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -62,7 +62,7 @@ claim an outsider can check. ## Positronic owns the control loop The world runner and harness execute every episode: they drive the clock, deliver observations to -the policy, schedule and play back action chunks, and own resets and episode boundaries. This holds +the policy, emit the commands it answers, and own resets and episode boundaries. This holds for every evaluation and data-collection run — in simulation and on real hardware, for any embodiment, scene source, or scoring method. A foreign component never runs the loop and calls into Positronic; Positronic runs the loop and calls into it. @@ -116,21 +116,22 @@ control system its clock, and no component reads time at point of use. Trajector the same time frame the observations carry, so a virtual clock, a slowed sim, or a replayed episode changes nothing downstream. -**The layer owns the plan, the harness plays it, the driver executes.** A policy speaks in -trajectories — waypoints with absolute timestamps — because a model predicts a horizon, not an -instant. But a trajectory on the wire makes every driver buffer the future, and makes the recording -guess which prefix of that buffer actually ran. So the plan stops at the harness: a command channel -carries the single command due at the moment it is emitted, the driver executes the latest one and +**The layer owns the plan and plays it, the driver executes.** A policy predicts a horizon, not an +instant, so a chunk of waypoints is what a model answers. But a trajectory on the wire makes every +driver buffer the future, and makes the recording guess which prefix of that buffer actually ran. So +the plan stops inside the policy stack: a session answers the commands due at the moment it is +called, a command channel carries one command per emission, the driver executes the latest one and holds otherwise, and emission time *is* execution time. Continuous-update schemes (RTC, temporal -ensembling) therefore need no special mechanism: they are layers that hand back a new trajectory -more often, and the harness keeps playing the old one until they do. +ensembling) therefore need no special mechanism: they are layers that re-query before the chunk they +hold runs out. **The harness stays thin.** It is the one layer standing between any policy and any embodiment, so anything it encodes about either side breaks the any-to-any goal. It assembles the observation -dict, calls the session, plays the returned trajectory one command per channel per round, and runs -episode lifecycle — nothing else. Scheduling, blending, history stacking and error recovery live in -the layer stack around the policy; a session returning `None` means "keep executing the current -trajectory". +dict, calls the session, emits the commands it answers, and runs episode lifecycle — nothing else. +Scheduling, blending, history stacking and error recovery live in the layer stack around the policy. +A call answers `(commands, resume_at_ns)`: what to emit now, and the instant the session wants its +next call. That instant is what paces the loop. The harness may call earlier: the episode deadline +cuts a round short, and a floor and a ceiling of its own bound every one of them. **Inference cost is a fact of the trial, owned by the harness.** The policy declares its heavy work as functions, and the framework runs each one off the loop thread. That work costs the trial either diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 3bf6f124d..0ef115fe7 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -82,7 +82,7 @@ How the client fills the delay and merges successive predictions is a swappable Four small concepts make up the API. You meet them whether you use a built-in server or write your own. -**Policy and Session.** A `Policy` is your loaded model: it holds the weights and knows how to start an episode. `policy.new_session()` begins one episode and returns a `Session`. You call the session once per timestep with the latest observation and your clock reading, and it returns the next actions to run. Per-episode state (history, the trajectory in flight) lives in the session — so one `Policy` can serve several robots at once, each with its own `Session`. +**Policy, episode and Session.** A `Policy` is your loaded model: it holds the weights and knows how to start an episode. `policy.episode()` opens one episode and answers the work that episode runs, by name — the model call under `infer`, and whatever that call needs for as long as the episode lasts (a connection, a client, a reset model). The framework runs each function off the loop thread and hands the handles to a `Session`, which it opens with `policy.new_session(rt)` and calls once per timestep with the latest observation and its clock reading. A session answers the commands to run now and the instant it wants its next call. A policy served behind the `remote` marker answers a chunk from `infer`; the rig's `ChunkPlayer` is the session that holds that chunk and emits each waypoint at its own time. One `Policy` serves several robots at once, each with its own episode. **Codec.** Different models want different inputs: end-effector pose vs joint angles, absolute targets vs deltas, 224×224 vs 512×512 images. A `Codec` translates between the robot's raw data (what is on the wire) and your model's format — `encode` on the way in, `decode` on the way out. The same codec prepares the training data, so a model is served exactly the way it was trained. The full catalog is in the [Codecs Guide](codecs.md). @@ -172,16 +172,24 @@ Implement a `Policy`, close a pipeline over it with `PolicySource`, and hand the ```python from positronic.drivers.roboarm import command from positronic.offboard import PolicyServer -from positronic.policy import Policy, Session +from contextlib import contextmanager + +from positronic.drivers.roboarm import command +from positronic.policy import INFER, Policy from positronic.policy.spec import PolicySource, remote -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault -class MySession(Session): +class MyPolicy(Policy): def __init__(self, model): self._model = model - def __call__(self, obs, time_ns): + @contextmanager + def episode(self, context=None): + # Per-episode setup goes here, and the teardown after the yield. + yield {INFER: self._infer} + + def _infer(self, obs): # obs holds the raw keys from the wire table above. Pick what you need: images = obs['image.exterior'] ee = obs['robot_state.ee_pose'] @@ -192,33 +200,25 @@ class MySession(Session): for pose in predicted_poses ] - -class MyPolicy(Policy): - def __init__(self, model): - self._model = model - - def new_session(self, context=None, rt=None): - return MySession(self._model) # per-episode setup goes here - @property def meta(self): return {'type': 'my_model'} -pipeline = StopOnFault() | ChunkedSchedule() | remote | PolicySource(MyPolicy(load_my_model())) +pipeline = StopOnFault() | ChunkPlayer() | remote | PolicySource(MyPolicy(load_my_model())) PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() ``` -The pipeline reads left to right: everything left of the `remote` marker is the client-side stack the server declares in its handshake (here the standard `StopOnFault` and `ChunkedSchedule`); everything right of it runs on the server. `PolicySource` is the pipeline's terminal — a model source that serves one already-built policy. +The pipeline reads left to right: everything left of the `remote` marker is the client-side stack the server declares in its handshake (here the standard `StopOnFault` and `ChunkPlayer`); everything right of it runs on the server. `PolicySource` is the pipeline's terminal — a model source that serves one already-built policy. -The left side is not optional: a pipeline with nothing there is refused when the server starts, and a rig refuses a handshake that declares nothing. It needs a scheduler in particular, and `StopOnFault` outside that scheduler — an arm that is faulted or busy is not taking the plan it was given, so the layer answers the empty trajectory and the rig stops rather than resuming a chunk stamped before. Actions come back timestamped relative to their chunk, and `ChunkedSchedule` is what turns those into times on the rig's clock; a stack that leaves them relative — or anchors them twice — makes the harness reject the chunk at the first inference, since it schedules nothing more than `MAX_ACTION_SKEW_SEC` from now. +The left side is not optional: a pipeline with nothing there is refused when the server starts, and a rig refuses a handshake that declares nothing. It needs a player in particular, and `StopOnFault` outside that player — an arm that is faulted or busy is not taking the plan it was given, so the layer commands nothing and drops the chunk rather than playing one stamped before. Actions come back timestamped relative to their chunk, and `ChunkPlayer` is what holds the chunk and emits each waypoint at its own time on the rig's clock; a stack with no player answers a chunk where the harness expects commands, and the episode refuses to open. A chunk that reaches the player already anchored is refused too, since the player places nothing more than `MAX_ACTION_SKEW_SEC` from the call. -The session's `time_ns` argument is the caller's clock reading in nanoseconds, the same unit the observation's `obs_time_ns` carries. A session reads no clock of its own, and a policy that schedules nothing accepts the value and ignores it. +`infer` takes the observation and nothing else: the clock belongs to the session that plays what it answers, and each action carries its own offset from the call. A session's `time_ns` argument is the caller's clock reading in nanoseconds, the same unit the observation's `obs_time_ns` carries. A session reads no clock of its own. -If you put a `Codec` right of the marker (`ChunkedSchedule() | remote | codec | PolicySource(...)`), your session works entirely in *model space* — it receives encoded observations and returns model-native actions, and the codec handles the wire format. A codec that encodes images should also bound them on the rig, so full-resolution frames never cross the wire — that is what the built-in vendor pipelines do: +If you put a `Codec` right of the marker (`ChunkPlayer() | remote | codec | PolicySource(...)`), your `infer` works entirely in *model space* — it receives encoded observations and returns model-native actions, and the codec handles the wire format. A codec that encodes images should also bound them on the rig, so full-resolution frames never cross the wire — that is what the built-in vendor pipelines do: ```python -StopOnFault() | ChunkedSchedule() | RestrictImageSize() | remote | codec | source +StopOnFault() | ChunkPlayer() | RestrictImageSize() | remote | codec | source ``` Give it the geometry your codec encodes to — `RestrictImageSize(224, 224)` for a 224x224 model — so a frame is shrunk once, on the rig. The default is a loose 640x640, for a codec that resizes to nothing in particular. Leaving it out costs bandwidth, not correctness. @@ -253,15 +253,15 @@ Every message is msgpack. Numpy arrays use a custom extension: robot commands: ```python -import time - from positronic.offboard.protocol import serialise, deserialise - -session = policy.new_session() # one Session per episode/connection -async for message in websocket.iter_bytes(): - obs = deserialise(message) # dict with numpy arrays - actions = session(obs, time.time_ns()) # list of action dicts (or None) - await websocket.send_bytes(serialise({"result": actions})) +from positronic.policy import INFER + +with policy.episode() as fns: # one episode per connection + infer = fns[INFER] + async for message in websocket.iter_bytes(): + obs = deserialise(message) # dict with numpy arrays + actions = infer(obs) # list of action dicts (or None) + await websocket.send_bytes(serialise({"result": actions})) ``` A server written against another stack cannot import that module. Answer with the command as the plain diff --git a/docs/inference.md b/docs/inference.md index 2ab38c765..b44baf272 100644 --- a/docs/inference.md +++ b/docs/inference.md @@ -71,7 +71,7 @@ The model source (`checkpoints_dir`, `checkpoint`, device...) is fixed at server **Credentials stay out of the URL.** `--policy.headers='{"Modal-Key": "..."}'` passes auth headers for a fronted endpoint, so the URL itself is safe to paste around. Against a Nebius endpoint use `.authed_remote` or `.nebius_remote`, which build the bearer header for you (see [the Nebius workflow README](../workflows/nebius/README.md#authenticated-inference)). -**What crosses the wire is the server's call, not the client's.** A server that wants smaller frames declares `RestrictImageSize` in its rig-side stack (640x640 by default); one behind a proxy with a message-size cap declares `remote(compress_images=True)` and the rig JPEG-encodes frames before sending. A server whose checkpoint speaks a different end-effector frame declares `ChangeEEFrame` with the transform placing that frame relative to the rig's `default`, and the rig converts poses (see [End-effector frames](codecs.md#end-effector-frames)). The client builds whatever the handshake declares, and only that — connecting to a server that declares no stack fails with an error naming the version it runs. What the declared stack must achieve is checked where it matters: the harness refuses to emit an action scheduled further than `MAX_ACTION_SKEW_SEC` from now, which is what a stack that never anchored its chunk to the rig's clock produces. +**What crosses the wire is the server's call, not the client's.** A server that wants smaller frames declares `RestrictImageSize` in its rig-side stack (640x640 by default); one behind a proxy with a message-size cap declares `remote(compress_images=True)` and the rig JPEG-encodes frames before sending. A server whose checkpoint speaks a different end-effector frame declares `ChangeEEFrame` with the transform placing that frame relative to the rig's `default`, and the rig converts poses (see [End-effector frames](codecs.md#end-effector-frames)). The client builds whatever the handshake declares, and only that — connecting to a server that declares no stack fails with an error naming the version it runs. What the declared stack must achieve is checked where it matters: a stack with no `ChunkPlayer` answers a chunk where the harness expects commands, and the episode refuses to open. > **Recording inference I/O:** Pass `--policy.recording_dir=s3://bucket/path` to write a rerun `.rrd` file per episode capturing the raw and server-side observation/action boundaries. Useful for debugging codec behavior and visualizing what the policy actually received. diff --git a/positronic/cfg/layers.py b/positronic/cfg/layers.py index 481b5fc17..07698608c 100644 --- a/positronic/cfg/layers.py +++ b/positronic/cfg/layers.py @@ -1,9 +1,9 @@ import configuronic as cfn from positronic import keys as obs_keys -from positronic.policy.layers import ChunkedSchedule, StopOnFault, TemporalStack +from positronic.policy.layers import ChunkPlayer, StopOnFault, TemporalStack -chunked_schedule = cfn.Config(ChunkedSchedule) +chunk_player = cfn.Config(ChunkPlayer) temporal_stack = cfn.Config(TemporalStack) @@ -47,4 +47,4 @@ def video_context_layers(history_frames: int, stride: int, keys: tuple[str, ...] stack = TemporalStack( keys=tuple(keys), offsets_sec=_frame_offsets_sec(history_frames, stride, fps), pad_start=pad_start ) - return StopOnFault() | stack | ChunkedSchedule() + return StopOnFault() | stack | ChunkPlayer() diff --git a/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index 3aff56b0e..24ecd737a 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -20,7 +20,6 @@ from positronic.dataset.ds_writer_agent import TimeMode from positronic.dataset.local_dataset import LocalDatasetWriter from positronic.eval import Embodiment, Eval, Task -from positronic.policy.executor import blocking from positronic.policy.harness import Harness from positronic.simulator.env_server.telemetry import ATTR_RUN_ID, ENV_RUN_ID, ENV_TELEMETRY_DIR @@ -193,8 +192,9 @@ def main(policy, *, evals: list[Eval], output_dir: str | Path | None = None, tim # TODO: a policy with recording taps (recording_dir set) records this throwaway warmup session — an # empty .rrd plus a bump to the recorder's episode counter — but warmup is not a real episode. logger.info('Warming up policy endpoints') - # The session runs no inference, but a session that serves its model on a runtime needs one to open. - blocking(policy).new_session().close() + # The episode runs no inference; opening it is what pays for the connection and the model behind it. + with policy.episode(): + pass output_dir = prepare_output_dir(output_dir) try: diff --git a/positronic/offboard/README.md b/positronic/offboard/README.md index ef025b0c6..3400cd211 100644 --- a/positronic/offboard/README.md +++ b/positronic/offboard/README.md @@ -92,7 +92,7 @@ Upon connection, the server sends a ready packet with metadata: "action_fps": 15.0, "action_horizon_sec": 1.0, "local_stack": {"seq": [ - {"name": "chunked_schedule"}, + {"name": "chunk_player"}, {"name": "restrict_image_size", "args": {"width": 224, "height": 224}} ]}, "compress_images": false, @@ -113,8 +113,9 @@ This metadata tells the client: `positronic.policy.spec.WIRE_LAYERS` — an unknown entry fails at connect, before the robot moves. Never empty and never absent: a pipeline with nothing left of the marker is refused when the server starts, and a handshake declaring nothing is refused by the client. In practice it names at least a - `chunked_schedule`, which turns the chunk-relative timestamps a codec stamps into times on the rig's - clock — a stack that fails to leaves the harness rejecting the chunk at the first inference. + `chunk_player`, which holds the chunk a codec stamps and emits each waypoint at its own time on the + rig's clock — a stack without one answers a chunk where the harness expects commands, and the episode + refuses to open. - `compress_images` — the `remote` marker's own wire setting: whether the rig JPEG-encodes frames before sending, for an endpoint behind a proxy with a message-size cap - `positronic_version` — the server's positronic version, for diagnosing declaration mismatches @@ -214,9 +215,9 @@ The one server implementation behind every vendor. It serves a **policy pipeline ```python from positronic.offboard import PolicyServer from positronic.policy.spec import PolicySource, remote -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer -pipeline = ChunkedSchedule() | remote | PolicySource(my_policy) +pipeline = ChunkPlayer() | remote | PolicySource(my_policy) PolicyServer(pipeline, host='0.0.0.0', port=8000).serve() ``` diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index d4060a01c..2f437a4c0 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -18,9 +18,8 @@ from starlette.datastructures import QueryParams from positronic import keys -from positronic.policy import Policy, Recorder +from positronic.policy import INFER, Policy, Recorder from positronic.policy.base import Layer -from positronic.policy.executor import blocking from positronic.policy.spec import ModelSource, Pipeline, split from . import protocol @@ -174,7 +173,7 @@ def _declared_stack(local: Layer | None) -> dict[str, Any]: if local is None: raise ValueError( 'Nothing sits left of the `remote` marker, so the pipeline declares no rig-side stack. Put the ' - 'layers the rig runs there, starting with a scheduler such as ChunkedSchedule' + 'layers the rig runs there, starting with a scheduler such as ChunkPlayer' ) return local.to_spec() @@ -306,7 +305,7 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): self._active_sessions += 1 self._last_activity = time.monotonic() policy: Policy | None = None - session = None + episode = None try: pipeline = self._session_pipeline(_session_params(websocket.query_params)) local, border, remote_half = split(pipeline) @@ -315,33 +314,31 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): rid = self._source.resolve(model_id) if model_id is not None else self._default_id assert rid is not None policy = await self._manager.get_policy(rid, websocket) - # A request has no control loop to answer ``None`` to. This goes innermost, so every layer - # above it sees one call per answer rather than one per call the answer took. - answered = blocking(policy) if self._recording_dir is not None: # Tap both sides: 'raw' is the wire boundary, 'inference' the encoded obs and model output. rec = Recorder(self._recording_dir) if remote_half is not None: - served = (rec.tap('raw') | remote_half | rec.tap('inference')).wrap(answered) + served = (rec.chunk_tap('raw') | remote_half | rec.chunk_tap('inference')).wrap(policy) else: - served = rec.tap('inference').wrap(answered) + served = rec.chunk_tap('inference').wrap(policy) else: - served = remote_half.wrap(answered) if remote_half is not None else answered - # ``new_session`` resets the shared backend client, so it must not interleave with an in-flight - # inference. Keepalives here: queuing behind a peer would otherwise trip the handshake timeout. + served = remote_half.wrap(policy) if remote_half is not None else policy + # Opening an episode resets the shared backend client, so it must not interleave with an + # in-flight inference. Keepalives here: queuing behind a peer would otherwise trip the + # handshake timeout. await _acquire_with_keepalives(self._infer_lock, websocket, 'Waiting for inference slot') + opening = served.episode() try: - session = await asyncio.to_thread(served.new_session) + infer = (await asyncio.to_thread(opening.__enter__))[INFER] + episode = opening finally: self._infer_lock.release() - assert session is not None - # Later entries win: per-episode session facts over static ones, the server's own last. + # Later entries win: per-episode facts over static ones, the server's own last. meta = { **self.metadata, **self._source.meta(rid), keys.CHECKPOINT_ID: rid, **served.meta, - **session.meta, keys.LOCAL_STACK: local_spec, keys.COMPRESS_IMAGES: border.compress_images, keys.POSITRONIC_VERSION: _pkg_version('positronic'), @@ -357,8 +354,7 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): # Plain acquire, not the keepalive helper: the client is awaiting a ``result`` and # would mis-parse a ``waiting`` message. Its ``infer_timeout`` bounds the wait. async with self._infer_lock: - # The server's clock is not the rig's. - actions = await asyncio.to_thread(session, raw_obs, time.time_ns()) + actions = await asyncio.to_thread(infer, raw_obs) await websocket.send_bytes(serialise({protocol.RESULT: actions})) except Exception as e: logger.error(f'Error processing message: {e}', exc_info=True) @@ -379,12 +375,12 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): self._active_sessions = max(0, self._active_sessions - 1) self._last_activity = time.monotonic() try: - if session is not None: - # Both ends of a session's life touch the backend — close does a reset round-trip — so - # it takes the inference lock like ``new_session`` and runs off the event loop. The - # nesting keeps a failure here from swallowing the manager release. + if episode is not None: + # Both ends of an episode touch the backend — the close does a reset round trip — so it + # takes the inference lock like the open, and runs off the event loop. The nesting keeps + # a failure here from swallowing the manager release. async with self._infer_lock: - await asyncio.to_thread(session.close) + await asyncio.to_thread(episode.__exit__, None, None, None) finally: if policy is not None: await self._manager.release_session() diff --git a/positronic/offboard/server_utils.py b/positronic/offboard/server_utils.py index 6e27d2ff4..a4122c783 100644 --- a/positronic/offboard/server_utils.py +++ b/positronic/offboard/server_utils.py @@ -9,10 +9,10 @@ import threading import time from collections.abc import Callable +from functools import partial from typing import Any -from positronic.policy import Policy -from positronic.policy.executor import blocking +from positronic.policy import INFER, Policy logger = logging.getLogger(__name__) @@ -42,11 +42,8 @@ def warmup(policy: Policy, obs: dict[str, Any], on_progress: Callable[[str], Non ``obs`` has to be an observation the loaded backend accepts. """ - session = blocking(policy).new_session() - try: - run_with_progress(lambda: session(obs, time.time_ns()), 'Running warmup inference', on_progress) - finally: - session.close() + with policy.episode() as fns: + run_with_progress(partial(fns[INFER], obs), 'Running warmup inference', on_progress) def wait_for_subprocess_ready( diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 5a371ff75..ff6e59213 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -3,15 +3,16 @@ import threading import time from collections.abc import Callable, Generator, Mapping -from unittest.mock import MagicMock +from contextlib import ExitStack, contextmanager +from typing import Any import pytest import uvicorn from positronic.offboard.server import PolicyServer -from positronic.policy import Policy, Session +from positronic.policy import INFER, Policy, Session from positronic.policy.executor import Executor -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.spec import ModelSource, PolicySource, remote @@ -57,52 +58,84 @@ async def _run(): @pytest.fixture -def open_session() -> Generator[Callable[..., tuple[Session, Executor]], None, None]: - """Opens a policy's session against a runtime that serves its functions, as the harness does.""" - runtimes: list[Executor] = [] +def open_episode() -> Generator[Callable[[Policy], Mapping[str, Callable[..., Any]]], None, None]: + """Opens a policy's episode; every one it opened is closed at teardown.""" + closing = ExitStack() + + def make(policy: Policy) -> Mapping[str, Callable[..., Any]]: + return closing.enter_context(policy.episode()) + + yield make + closing.close() + + +@pytest.fixture +def open_session() -> Generator[Callable[[Policy], tuple[Session, Executor]], None, None]: + """Opens a policy's session against a runtime that serves its episode, as the harness does.""" + closing = ExitStack() def make(policy: Policy) -> tuple[Session, Executor]: - runtimes.append(Executor(policy.functions)) - return policy.new_session(None, runtimes[-1]), runtimes[-1] + rt = Executor(closing.enter_context(policy.episode())) + closing.callback(rt.close) + session = policy.new_session(rt) + closing.callback(session.close) + return session, rt yield make - for runtime in runtimes: - runtime.close() + closing.close() # How long a round trip against a local server may take before a test calls it lost. ANSWER_SEC = 5.0 -def round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> list[dict] | None: - """What ``session`` answers for ``obs``, over the two calls one round trip takes. +def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> Mapping[str, Any]: + """What ``session`` commands for ``obs`` once the call under its ``ChunkPlayer`` has come back. - Both calls get the same ``time_ns``, so a chunk comes back anchored at the value the test passed. + Every call gets the same ``time_ns``, so the chunk comes back anchored at the value the test passed and + the answer is the waypoints due at it. """ - assert session(obs, time_ns) is None, 'a round-trip was already in flight' + commands, _ = session(obs, time_ns) + if commands: + return commands rt.wait(ANSWER_SEC) assert not rt.in_flight, 'the round-trip never came back' - return session(obs, time_ns) + return session(obs, time_ns)[0] + + +class MockPolicy(Policy): + """Answers one fixed chunk, or raises ``failure``. Records its episodes and the observations it took.""" + + def __init__(self, action, meta: dict[str, Any], failure: Exception | None = None): + self._action = action + self._meta = meta + self._failure = failure + self.episodes = 0 + self.closed = 0 + self.observations: list[Any] = [] + @contextmanager + def episode(self, context=None): + self.episodes += 1 + try: + yield {INFER: self._infer} + finally: + self.closed += 1 -def _make_mock_policy(action, meta): - """Create a mock policy with session-based API.""" - session = MagicMock() - session.return_value = action - session.meta = meta - session.close = MagicMock() + def _infer(self, obs): + self.observations.append(obs) + if self._failure is not None: + raise self._failure + return self._action - policy = MagicMock() - policy.new_session.return_value = session - policy.meta = meta - policy.functions = {} # `Policy.functions` is a mapping, and MagicMock's stand-in is not - policy._mock_session = session # expose for assertions - return policy + @property + def meta(self) -> dict[str, Any]: + return dict(self._meta) @pytest.fixture -def make_mock_policy() -> Callable[..., MagicMock]: - return _make_mock_policy +def make_mock_policy() -> Callable[..., MockPolicy]: + return MockPolicy class _DictSource(ModelSource): @@ -126,33 +159,32 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) @pytest.fixture -def mock_policy() -> MagicMock: - """Mock policy for testing.""" - return _make_mock_policy({'action_data': [1, 2, 3]}, {'model_name': 'test_model'}) +def mock_policy() -> MockPolicy: + return MockPolicy({'action_data': [1, 2, 3]}, {'model_name': 'test_model'}) @pytest.fixture -def mock_policy_registry() -> dict[str, MagicMock]: +def mock_policy_registry() -> dict[str, MockPolicy]: return { - 'alpha': _make_mock_policy({'action_data': ['alpha']}, {'model_name': 'alpha'}), - 'beta': _make_mock_policy({'action_data': ['beta']}, {'model_name': 'beta'}), + 'alpha': MockPolicy({'action_data': ['alpha']}, {'model_name': 'alpha'}), + 'beta': MockPolicy({'action_data': ['beta']}, {'model_name': 'beta'}), } @pytest.fixture -def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple[str, int]: +def inference_server(start_server: StartServer, mock_policy: MockPolicy) -> tuple[str, int]: """A served single-policy pipeline. Returns: tuple[str, int]: (host, port) """ - host, port, _server = start_server(ChunkedSchedule() | remote | PolicySource(mock_policy)) + host, port, _server = start_server(ChunkPlayer() | remote | PolicySource(mock_policy)) return host, port @pytest.fixture def multi_policy_server( - start_server: StartServer, mock_policy_registry: dict[str, MagicMock] -) -> tuple[str, int, dict[str, MagicMock]]: - host, port, _server = start_server(ChunkedSchedule() | remote | _DictSource(mock_policy_registry)) + start_server: StartServer, mock_policy_registry: dict[str, MockPolicy] +) -> tuple[str, int, dict[str, MockPolicy]]: + host, port, _server = start_server(ChunkPlayer() | remote | _DictSource(mock_policy_registry)) return host, port, mock_policy_registry diff --git a/positronic/offboard/tests/test_offboard.py b/positronic/offboard/tests/test_offboard.py index 6967f70a3..121322f2d 100644 --- a/positronic/offboard/tests/test_offboard.py +++ b/positronic/offboard/tests/test_offboard.py @@ -1,5 +1,4 @@ from types import MappingProxyType -from unittest.mock import ANY import numpy as np import pytest @@ -34,13 +33,13 @@ def test_inference_client_connect_and_infer(inference_server, mock_policy): action = session.infer(obs) assert action['action_data'] == [1, 2, 3] - mock_policy._mock_session.assert_called_with(obs, ANY) + assert mock_policy.observations[-1] == obs finally: session.close() def test_inference_client_new_session(inference_server, mock_policy): - """Test that starting a new session calls new_session on the policy.""" + """Test that starting a new session opens an episode on the policy.""" host, port = inference_server client = InferenceClient(f'{host}:{port}') @@ -52,7 +51,7 @@ def test_inference_client_new_session(inference_server, mock_policy): session = client.new_session() session.close() - assert mock_policy.new_session.call_count == 2 + assert mock_policy.episodes == 2 def test_session_url_selects_the_model(multi_policy_server): @@ -83,9 +82,8 @@ def test_session_url_selects_the_model(multi_policy_server): finally: beta_session.close() - policies['alpha']._mock_session.assert_any_call({'obs': 'alpha'}, ANY) - policies['beta']._mock_session.assert_any_call({'obs': 'beta'}, ANY) - policies['alpha']._mock_session.assert_any_call({'obs': 'default'}, ANY) + assert policies['alpha'].observations == [{'obs': 'default'}, {'obs': 'alpha'}] + assert policies['beta'].observations == [{'obs': 'beta'}] def test_wire_serialisation_accepts_mappingproxy(): diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index c1a385296..899026a7d 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -1,4 +1,3 @@ -import threading import time from http import HTTPStatus from unittest.mock import MagicMock, patch @@ -12,17 +11,17 @@ from positronic import keys, telemetry, telemetry_keys from positronic.drivers.roboarm import command from positronic.offboard.client import DEFAULT_INFER_TIMEOUT, InferenceClient, _ConnectRetries -from positronic.offboard.tests.conftest import ANSWER_SEC, round_trip -from positronic.policy import RemotePolicy +from positronic.offboard.tests.conftest import played_round_trip +from positronic.policy import INFER, RemotePolicy from positronic.policy.codec import ActionHorizon -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.remote import _prepare_obs from positronic.policy.spec import PolicySource, remote # These fixtures stand in for a server, so they spell the handshake fields rather than importing the # ``keys`` constants the client reads: sharing a constant makes the two agree whatever its value, which # would leave nothing pinning the client to the wire. -CHUNKED_STACK = {'local_stack': {'name': 'chunked_schedule'}} +CHUNKED_STACK = {'local_stack': {'name': 'chunk_player'}} def _mock_ws_session(metadata=None): @@ -277,7 +276,7 @@ def test_remote_policy_hands_the_url_and_headers_to_the_client(): class TestActionHorizonWrapping: - def test_truncates_action_chunks(self, open_session): + def test_truncates_action_chunks(self, open_episode): actions = [ {'a': 1, 'timestamp': 0.0}, {'a': 2, 'timestamp': 0.25}, @@ -285,159 +284,62 @@ def test_truncates_action_chunks(self, open_session): {'a': 4, 'timestamp': 0.75}, ] endpoint, _ = _mock_endpoint(infer_return=actions) - session, rt = open_session(ActionHorizon(0.5).wrap(endpoint)) + fns = open_episode(ActionHorizon(0.5).wrap(endpoint)) - actions = round_trip(session, rt, {keys.OBS_TIME_NS: 0}) + actions = fns[INFER]({keys.OBS_TIME_NS: 0}) assert actions is not None assert len(actions) == 3 # 2 within-horizon actions + horizon sentinel assert actions[0]['timestamp'] == 0.0 assert actions[1]['timestamp'] == 0.25 assert actions[2] == {'timestamp': 0.5} # horizon sentinel (timestamp = horizon_sec) - def test_no_truncation_without_horizon(self, open_session): + def test_no_truncation_without_horizon(self, open_episode): endpoint, _ = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}, {'a': 2, 'timestamp': 1.0}]) - session, rt = open_session(endpoint) - - actions = round_trip(session, rt, {}) + actions = open_episode(endpoint)[INFER]({}) assert actions is not None assert len(actions) == 2 -def test_remote_session_normalizes_single_dict(open_session): - """Server returning a single action dict is wrapped into a 1-element list.""" +def test_the_wire_answers_a_single_dict_as_it_came(open_episode): + """A server answer of one action is what the player normalizes, so the wire passes it through.""" endpoint, _ = _mock_endpoint(infer_return={keys.ROBOT_COMMAND: 'X', 'timestamp': 0.0}) - session, rt = open_session(endpoint) - assert round_trip(session, rt, {}) == [{keys.ROBOT_COMMAND: 'X', 'timestamp': 0.0}] + assert open_episode(endpoint)[INFER]({}) == {keys.ROBOT_COMMAND: 'X', 'timestamp': 0.0} -def test_remote_session_passes_through_none(open_session): +def test_remote_inference_passes_through_none(open_episode): endpoint, mock_ws = _mock_endpoint() mock_ws.infer.return_value = None - session, rt = open_session(endpoint) - - assert round_trip(session, rt, {}) is None - - -def test_a_call_while_a_round_trip_is_in_flight_answers_none(open_session): - """A session never waits. Every call while the round trip is in flight answers ``None``, and none of - them starts a second round trip.""" - chunk = [{'a': 1, 'timestamp': 0.0}] - endpoint, mock_ws = _mock_endpoint() - started, release = threading.Event(), threading.Event() - def blocked(obs): - started.set() - assert release.wait(ANSWER_SEC), 'the test never released the round-trip' - return chunk + assert open_episode(endpoint)[INFER]({}) is None - mock_ws.infer.side_effect = blocked - session, rt = open_session(endpoint) - assert session({}, 0) is None - assert started.wait(ANSWER_SEC), 'the round-trip never started' - assert session({}, 0) is None - assert mock_ws.infer.call_count == 1 - - release.set() - rt.wait(ANSWER_SEC) - assert session({}, 0) == chunk - - -def test_opening_a_session_without_a_runtime_is_refused(): - """Nothing serves the round trip without a runtime, so the session is refused where it is opened, and not - at the first observation it is given.""" +def test_a_chunk_policy_has_no_session_of_its_own(): + """A chunk policy answers ``INFER`` and nothing else; a ``ChunkPlayer`` above it is what plays one.""" endpoint, _ = _mock_endpoint() - with pytest.raises(ValueError, match='runs its inference on a runtime'): - endpoint.new_session() - - -def test_cancel_drops_the_chunk_of_the_round_trip_in_flight(open_session): - """A cancelled session drops the chunk it waited for, because that chunk applies to a world the cancel - says has gone, and it asks for a new one.""" - endpoint, mock_ws = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) - session, rt = open_session(endpoint) - - assert session({}, 0) is None - rt.wait(ANSWER_SEC) - session.cancel() - - assert session({}, 0) is None # the cancelled answer, read and thrown away - assert session({}, 0) is None # a round-trip of its own - rt.wait(ANSWER_SEC) - assert mock_ws.infer.call_count == 2 - - -def test_a_cancelled_round_trip_still_raises_what_it_failed_with(open_session): - """A dropped chunk drops no failure. The session reads a cancelled answer, so a stalled server raises - to the caller that asked for the episode.""" - endpoint, mock_ws = _mock_endpoint() - mock_ws.infer.side_effect = TimeoutError('server stalled') - session, rt = open_session(endpoint) - - assert session({}, 0) is None - rt.wait(ANSWER_SEC) - session.cancel() + with pytest.raises(NotImplementedError, match='ChunkPlayer'): + endpoint.new_session(MagicMock()) - with pytest.raises(TimeoutError, match='server stalled'): - session({}, 0) - -def test_a_cancel_dies_with_the_answer_it_was_made_against(open_session): - """A cancel ends with the round trip it was made against, even when that round trip fails. A caller - that catches the failure and keeps the session gets the next chunk.""" - endpoint, mock_ws = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) - mock_ws.infer.side_effect = [TimeoutError('server stalled'), [{'a': 1, 'timestamp': 0.0}]] - session, rt = open_session(endpoint) - - assert session({}, 0) is None - rt.wait(ANSWER_SEC) - session.cancel() - with pytest.raises(TimeoutError, match='server stalled'): - session({}, 0) - - assert round_trip(session, rt, {}) == [{'a': 1, 'timestamp': 0.0}] - - -def test_closing_a_session_with_a_round_trip_in_flight_is_refused(open_session): - """A runtime closes before the session it serves. A caller that closes the websocket under a round trip - gets an error that names the order, and not a failure on a dead socket.""" - endpoint, mock_ws = _mock_endpoint() - release = threading.Event() - - def blocked(obs): - assert release.wait(ANSWER_SEC), 'the test never released the round-trip' - return None - - mock_ws.infer.side_effect = blocked - session, _rt = open_session(endpoint) - - assert session({}, 0) is None - with pytest.raises(AssertionError, match='close the runtime'): - session.close() - - release.set() - - -def test_records_infer_span_without_scheduling_layer(tmp_path, open_session): +def test_records_infer_span_without_scheduling_layer(tmp_path, open_episode): """The ``policy.infer`` span is recorded at the remote inference boundary itself, not by a layer in front of it.""" endpoint, _ = _mock_endpoint(infer_return=[{'a': 1, 'timestamp': 0.0}]) - session, rt = open_session(endpoint) + infer = open_episode(endpoint)[INFER] with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-span'): - assert round_trip(session, rt, {keys.OBS_TIME_NS: 0}) is not None + assert infer({keys.OBS_TIME_NS: 0}) is not None spans = list(telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS))) assert [s.name for s in spans] == [telemetry_keys.SPAN_POLICY_INFER] -def test_infer_span_excludes_client_side_image_preparation(tmp_path, open_session): +def test_infer_span_excludes_client_side_image_preparation(tmp_path, open_episode): """``policy.infer`` is the remote round-trip, so JPEG-encoding the observation stays outside it: folding client CPU work into the span would inflate the inference percentiles and the policy-server capacity estimate the report derives from them.""" endpoint, _ = _mock_endpoint({'compress_images': True}, infer_return=[]) - session, rt = open_session(endpoint) + infer = open_episode(endpoint)[INFER] encoded_at: list[int] = [] def _stamp_encode(image): @@ -446,7 +348,7 @@ def _stamp_encode(image): with patch('positronic.policy.remote.encode_jpeg', side_effect=_stamp_encode): with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-prep'): - round_trip(session, rt, {'cam': _make_image(48, 64), keys.OBS_TIME_NS: 0}) + infer({'cam': _make_image(48, 64), keys.OBS_TIME_NS: 0}) (span,) = telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS)) assert span.name == telemetry_keys.SPAN_POLICY_INFER @@ -454,15 +356,14 @@ def _stamp_encode(image): assert span.start_ns >= encoded_at[-1] # every encode finishes before the span opens, not inside it -def test_records_infer_span_when_inference_raises(tmp_path, open_session): - """A round trip that raises still records the time it took to fail, and the answer raises it again at - the call that reads it.""" +def test_records_infer_span_when_inference_raises(tmp_path, open_episode): + """A round trip that raises still records the time it took to fail.""" endpoint, mock_ws = _mock_endpoint() mock_ws.infer.side_effect = TimeoutError('server stalled') - session, rt = open_session(endpoint) + infer = open_episode(endpoint)[INFER] with telemetry.bind(tmp_path, telemetry_keys.HARNESS_PROCESS, 'run-infer-raise'): with pytest.raises(TimeoutError): - round_trip(session, rt, {keys.OBS_TIME_NS: 0}) + infer({keys.OBS_TIME_NS: 0}) spans = list(telemetry.read_spans(telemetry.spans_path(tmp_path, telemetry_keys.HARNESS_PROCESS))) assert [s.name for s in spans] == [telemetry_keys.SPAN_POLICY_INFER] @@ -478,47 +379,45 @@ def test_remote_policy_meta_exposes_server_fields(): def test_missing_declaration_fails_before_motion(): - """A handshake carrying no ``local_stack`` leaves nothing to build, so no session opens.""" + """A handshake carrying no ``local_stack`` leaves nothing to build, so no episode opens.""" policy, _ = _mock_remote_policy({'positronic_version': '0.1.0'}) with pytest.raises(ValueError, match='0.1.0'): - policy.new_session() + policy.episode().__enter__() def test_empty_declaration_fails_before_motion(): """An empty stack declares nothing to build, so it is refused like an absent one.""" policy, _ = _mock_remote_policy({'local_stack': {'seq': []}}) with pytest.raises(ValueError, match='declares no rig-side stack'): - policy.new_session() + policy.episode().__enter__() -def test_declared_stack_built_at_session_open(open_session): +def test_declared_stack_built_at_episode_open(open_session): """The server-declared local stack runs in front of the connection.""" policy, mock_ws = _mock_remote_policy(CHUNKED_STACK, infer_return=[{'a': 1, 'timestamp': 0.0}]) session, rt = open_session(policy) - assert round_trip(session, rt, {keys.OBS_TIME_NS: 0}, int(1e9)) == [{'a': 1, 'timestamp': 1.0}] + assert played_round_trip(session, rt, {keys.OBS_TIME_NS: 0}, int(1e9)) == {'a': 1} def test_unknown_declared_entry_fails_before_motion(): policy, _ = _mock_remote_policy({'local_stack': {'name': 'run_arbitrary_code'}, keys.POSITRONIC_VERSION: '9.9.9'}) with pytest.raises(ValueError, match='9.9.9'): - policy.new_session() + policy.episode().__enter__() -def test_compression_follows_the_server_declaration(open_session): +def test_compression_follows_the_server_declaration(open_episode): """A server behind a message-size cap declares ``remote(compress_images=True)`` and the rig obeys.""" endpoint, mock_ws = _mock_endpoint({'compress_images': True}, infer_return=[]) - session, rt = open_session(endpoint) - round_trip(session, rt, {'cam': _make_image(48, 64)}) + open_episode(endpoint)[INFER]({'cam': _make_image(48, 64)}) assert isinstance(mock_ws.infer.call_args.args[0]['cam'], dict) -def test_frames_stay_raw_where_the_server_declares_no_compression(open_session): +def test_frames_stay_raw_where_the_server_declares_no_compression(open_episode): endpoint, mock_ws = _mock_endpoint({'compress_images': False}, infer_return=[]) - session, rt = open_session(endpoint) - round_trip(session, rt, {'cam': _make_image(48, 64)}) + open_episode(endpoint)[INFER]({'cam': _make_image(48, 64)}) assert isinstance(mock_ws.infer.call_args.args[0]['cam'], np.ndarray) @@ -531,32 +430,32 @@ def test_a_command_crossing_a_live_websocket_arrives_typed(start_server, make_mo pose = [0.4, 0.0, 0.6, 1, 0, 0, 0, 1, 0, 0, 0, 1] # translation + a 3x3 rotation, the wire's own layout wire_action = [{keys.ROBOT_COMMAND: {'type': 'cartesian_pos', 'pose': pose}, 'timestamp': 0.0}] served = make_mock_policy(wire_action, {'model_name': 'm'}) - host, port, _ = start_server(ChunkedSchedule() | remote | PolicySource(served)) + host, port, _ = start_server(ChunkPlayer() | remote | PolicySource(served)) session, rt = open_session(RemotePolicy(f'{host}:{port}')) - actions = round_trip(session, rt, {keys.OBS_TIME_NS: 0}) + commands = played_round_trip(session, rt, {keys.OBS_TIME_NS: 0}) - assert actions is not None, 'the chunk was swallowed before any command reached a driver' - decoded = actions[0][keys.ROBOT_COMMAND] + assert commands, 'the chunk was swallowed before any command reached a driver' + decoded = commands[keys.ROBOT_COMMAND] assert isinstance(decoded, command.CartesianPosition), f'the driver would be handed {decoded!r}' np.testing.assert_allclose(decoded.pose.translation, [0.4, 0.0, 0.6], atol=1e-6) def test_remote_policy_lifecycle(inference_server, mock_policy, open_session): - """RemotePolicy against a live server whose pipeline declares a chunked_schedule local stack.""" + """RemotePolicy against a live server whose pipeline declares a chunk_player local stack.""" host, port = inference_server policy = RemotePolicy(f'{host}:{port}') session, rt = open_session(policy) - meta = session.meta + meta = policy.meta assert meta['server.model_name'] == 'test_model' assert meta['type'] == 'remote' - action = round_trip(session, rt, {'dataset': 'test'}) - # Single-dict server response is normalized to a 1-element list (Session contract) and - # anchored to absolute time by the declared ChunkedSchedule. - assert action == [{'action_data': [1, 2, 3], 'timestamp': 0.0}] + commands = played_round_trip(session, rt, {'dataset': 'test'}) + # A single-dict server answer is one waypoint, anchored to the call by the declared ChunkPlayer and + # commanded at once. + assert commands == {'action_data': [1, 2, 3]} session.close() @@ -565,13 +464,11 @@ def test_remote_policy_lifecycle(inference_server, mock_policy, open_session): session2.close() -def test_remote_session_meta(inference_server, open_session): - """Session meta must include server metadata.""" +def test_remote_policy_meta_carries_the_server_handshake(inference_server): + """A remote policy's meta must include the server metadata.""" host, port = inference_server - session, _ = open_session(RemotePolicy(f'{host}:{port}')) - meta = session.meta + meta = RemotePolicy(f'{host}:{port}').meta + assert meta['type'] == 'remote' assert meta['server.model_name'] == 'test_model' - - session.close() diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index e5ef71f01..fe970f370 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -3,8 +3,9 @@ import time import urllib.parse from collections.abc import Callable, Generator +from contextlib import contextmanager from typing import Any -from unittest.mock import ANY, MagicMock +from unittest.mock import MagicMock import configuronic as cfn import httpx @@ -17,11 +18,10 @@ from positronic.offboard.protocol import deserialise from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer from positronic.offboard.server_utils import warmup -from positronic.offboard.tests.conftest import round_trip -from positronic.policy import Codec, Policy, RemotePolicy, Session -from positronic.policy.base import Runtime +from positronic.offboard.tests.conftest import ANSWER_SEC +from positronic.policy import INFER, Codec, Policy, RemotePolicy from positronic.policy.codec import ActionTimestamp -from positronic.policy.layers import ChunkedSchedule, TemporalStack +from positronic.policy.layers import ChunkPlayer, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, inline, remote @@ -48,7 +48,7 @@ def meta(self, model_id: str) -> dict[str, Any]: @pytest.fixture def stub_server(start_server, make_mock_policy) -> tuple[str, int, PolicyServer, MagicMock]: policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, server = start_server(ChunkedSchedule() | remote | _StubSource(policy)) + host, port, server = start_server(ChunkPlayer() | remote | _StubSource(policy)) return host, port, server, policy @@ -59,13 +59,13 @@ def test_full_inference_cycle(stub_server): try: assert session.metadata['model_name'] == 'stub' assert session.metadata['type'] == 'stub' - assert session.metadata['local_stack'] == {'name': 'chunked_schedule'} + assert session.metadata['local_stack'] == {'name': 'chunk_player'} assert keys.POSITRONIC_VERSION in session.metadata obs = {'image': 'test'} result = session.infer(obs) assert result == [{'action': [1, 2, 3]}] - policy._mock_session.assert_called_with(obs, ANY) + assert policy.observations[-1] == obs finally: session.close() @@ -125,7 +125,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) def test_latest_checkpoint_pinned_once_at_startup(start_server, make_mock_policy): source = _LatestSource(make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'})) - host, port, _server = start_server(ChunkedSchedule() | remote | source) + host, port, _server = start_server(ChunkPlayer() | remote | source) # A newer checkpoint lands after startup (e.g. a training job writes it)... source.latest = '200' client = InferenceClient(f'{host}:{port}') @@ -154,7 +154,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) def test_load_progress_frames_reach_the_client(start_server, make_mock_policy): policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _ProgressSource(policy)) + host, port, _server = start_server(ChunkPlayer() | remote | _ProgressSource(policy)) # Requesting a non-pinned id forces a load inside the handshake; the source's progress # callbacks must arrive as ``loading`` frames before ``ready``. ws = connect(f'ws://{host}:{port}/api/v1/session/other') @@ -183,7 +183,7 @@ def meta(self): @pytest.fixture def codec_server(start_server, make_mock_policy) -> tuple[str, int, MagicMock]: policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _IdentityCodec() | _StubSource(policy)) + host, port, _server = start_server(ChunkPlayer() | remote | _IdentityCodec() | _StubSource(policy)) return host, port, policy @@ -199,34 +199,33 @@ def test_codec_wrapping(codec_server): session.close() -def test_warmup_runs_one_inference_and_ends_its_session(make_mock_policy): +def test_warmup_runs_one_inference_and_ends_its_episode(make_mock_policy): policy = make_mock_policy([{'action': [1, 2, 3]}], {}) obs = {'obs': 'zeros'} warmup(policy, obs) - policy._mock_session.assert_called_once_with(obs, ANY) - policy._mock_session.close.assert_called_once() + assert policy.observations == [obs] + assert policy.closed == 1 -def test_a_backend_that_cannot_answer_its_warmup_raises_and_still_ends_its_session(make_mock_policy): - policy = make_mock_policy([], {}) - policy._mock_session.side_effect = RuntimeError('shape mismatch') +def test_a_backend_that_cannot_answer_its_warmup_raises_and_still_ends_its_episode(make_mock_policy): + policy = make_mock_policy([], {}, failure=RuntimeError('shape mismatch')) with pytest.raises(RuntimeError, match='shape mismatch'): warmup(policy, {}) - policy._mock_session.close.assert_called_once() + assert policy.closed == 1 def test_local_stack_declared_in_handshake(start_server, make_mock_policy): stub = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - pipeline = ChunkedSchedule() | remote | _IdentityCodec() | _StubSource(stub) + pipeline = ChunkPlayer() | remote | _IdentityCodec() | _StubSource(stub) host, port, _server = start_server(pipeline) client = InferenceClient(f'{host}:{port}') session = client.new_session() try: - assert session.metadata['local_stack'] == {'name': 'chunked_schedule'} + assert session.metadata['local_stack'] == {'name': 'chunk_player'} finally: session.close() @@ -238,67 +237,56 @@ def test_pipeline_with_no_rig_side_half_refused_at_startup(make_mock_policy): PolicyServer(remote | _StubSource(stub)) -_INFER = 'infer' - - -class _ScriptedSession(Session): - def __init__(self, rt: Runtime): - self._rt = rt - self._answer = None - - def __call__(self, obs, time_ns): - if self._answer is None: - self._answer = self._rt.fns[_INFER](obs) - return None - if not self._answer.done(): - return None - answer, self._answer = self._answer, None - return answer.result() - - class _ScriptedPolicy(Policy): - """Deterministic base policy: every session serves the same untimestamped chunk from its runtime.""" - - def new_session(self, context=None, rt=None) -> Session: - assert rt is not None - return _ScriptedSession(rt) + """Deterministic base policy: every episode answers the same untimestamped chunk.""" - @property - def functions(self): - return {_INFER: lambda obs: [{'a': 1.0}, {'a': 2.0}, {'a': 3.0}]} + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: [{'a': 1.0}, {'a': 2.0}, {'a': 3.0}]} def test_in_process_equals_remote_for_same_pipeline(start_server, open_session): """The same pipeline must behave identically served in-process and over the wire.""" def pipeline(): - return ChunkedSchedule() | remote | ActionTimestamp(fps=10.0) | PolicySource(_ScriptedPolicy()) + return ChunkPlayer() | remote | ActionTimestamp(fps=10.0) | PolicySource(_ScriptedPolicy()) host, port, _server = start_server(pipeline()) remote_session, rt = open_session(RemotePolicy(f'{host}:{port}')) local_session, local_rt = open_session(inline(pipeline())) - remote_actions = round_trip(remote_session, rt, {keys.OBS_TIME_NS: 0}, int(100e9)) - local_actions = round_trip(local_session, local_rt, {keys.OBS_TIME_NS: 0}, int(100e9)) - assert remote_actions == local_actions - # Three scripted actions plus the chunk-closing validity sentinel ActionTimestamp appends. - assert local_actions == [ - {'a': 1.0, 'timestamp': 100.0}, - {'a': 2.0, 'timestamp': 100.1}, - {'a': 3.0, 'timestamp': 100.2}, - {'timestamp': 100.3}, - ] - - # Both gate identically while the chunk plays out. - assert remote_session({keys.OBS_TIME_NS: 0}, int(100.15e9)) is None - assert local_session({keys.OBS_TIME_NS: 0}, int(100.15e9)) is None - + def play_out(session, runtime): + """Every waypoint the session plays, and how long after the chunk loaded each one went out.""" + at_ns = int(100e9) + commands, resume_at_ns = session({keys.OBS_TIME_NS: 0}, at_ns) + while not commands: # the round trip has to land before the chunk plays + runtime.wait(ANSWER_SEC) + at_ns = resume_at_ns + commands, resume_at_ns = session({keys.OBS_TIME_NS: 0}, at_ns) + loaded_at_ns, played = at_ns, [(0, commands)] + for _ in range(3): + at_ns = resume_at_ns + commands, resume_at_ns = session({keys.OBS_TIME_NS: 0}, at_ns) + played.append((at_ns - loaded_at_ns, commands)) + return played + + played = play_out(local_session, local_rt) + assert play_out(remote_session, rt) == played + # Three scripted actions, then the chunk-closing validity sentinel ActionTimestamp appends, which + # commands nothing. + assert [commands for _, commands in played] == [{'a': 1.0}, {'a': 2.0}, {'a': 3.0}, {}] + gaps = [round((b - a) / 1e9, 6) for (a, _), (b, _) in zip(played, played[1:], strict=False)] + assert gaps == [0.1, 0.1, 0.1], 'the waypoints are one action period apart' + + # The call that drained the chunk asked for the next one, so the runtime closes first, as the harness + # closes an episode. + rt.close() remote_session.close() def _tunable_pipe(source: ModelSource, offsets: tuple[float, ...] = (-0.1, 0.0), pad_start: bool = True): - return TemporalStack(keys=('x',), offsets_sec=offsets, pad_start=pad_start) | ChunkedSchedule() | remote | source + return TemporalStack(keys=('x',), offsets_sec=offsets, pad_start=pad_start) | ChunkPlayer() | remote | source def _param_session(host: str, port: int, query: list[tuple[str, str]]) -> InferenceSession: @@ -335,7 +323,7 @@ def test_session_params_coerce_json_values(param_server): def _fps_pipe(source: ModelSource, fps: float = 10.0): - return ChunkedSchedule() | remote | ActionTimestamp(fps=fps) | source + return ChunkPlayer() | remote | ActionTimestamp(fps=fps) | source def test_session_param_retunes_the_served_remote_half(start_server): @@ -430,7 +418,7 @@ def authed_endpoint(start_server, make_mock_policy) -> tuple[str, str]: if _LIVE_ENDPOINT: return _LIVE_ENDPOINT, os.environ[AUTH_TOKEN_ENV] policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) + host, port, _server = start_server(ChunkPlayer() | remote | _StubSource(policy), auth_token=_TOKEN) return f'{host}:{port}', _TOKEN @@ -505,14 +493,14 @@ def test_server_without_a_token_serves_open(stub_server): ) def test_a_token_that_could_never_gate_fails_closed_at_startup(make_mock_policy, token): with pytest.raises(ValueError, match='ASCII'): - PolicyServer(ChunkedSchedule() | remote | _StubSource(make_mock_policy([], {})), auth_token=token) + PolicyServer(ChunkPlayer() | remote | _StubSource(make_mock_policy([], {})), auth_token=token) def test_a_non_ascii_authorization_header_is_refused_rather_than_crashing(start_server, make_mock_policy): """A header carries bytes, and Starlette hands them over latin-1 decoded, so a peer can put a non-ASCII ``str`` in front of the token comparison.""" policy = make_mock_policy([{'action': [1, 2, 3]}], {'model_name': 'stub'}) - host, port, _server = start_server(ChunkedSchedule() | remote | _StubSource(policy), auth_token=_TOKEN) + host, port, _server = start_server(ChunkPlayer() | remote | _StubSource(policy), auth_token=_TOKEN) with socket.create_connection((host, port), timeout=5.0) as sock: sock.sendall( b'GET /api/v1/models HTTP/1.1\r\nHost: localhost\r\n' diff --git a/positronic/policy/__init__.py b/positronic/policy/__init__.py index 297e11662..03acf0b13 100644 --- a/positronic/policy/__init__.py +++ b/positronic/policy/__init__.py @@ -1,14 +1,17 @@ -from .base import DelegatingPolicy, DelegatingSession, Layer, Policy, Session +from .base import INFER, Answer, ChunkLayer, DelegatingPolicy, DelegatingSession, Layer, Policy, Runtime, Session from .codec import ActionHorizon, ActionTimestamp, ActionTiming, Codec, is_action from .recording import Recorder from .remote import RemotePolicy __all__ = [ + 'INFER', 'Policy', 'Session', + 'Runtime', 'DelegatingPolicy', 'DelegatingSession', 'Layer', + 'ChunkLayer', 'RemotePolicy', 'Codec', 'ActionTimestamp', @@ -16,4 +19,5 @@ 'ActionTiming', 'is_action', 'Recorder', + 'Answer', ] diff --git a/positronic/policy/base.py b/positronic/policy/base.py index db88af250..ff0567c42 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,13 +1,19 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager from typing import Any, ClassVar # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. SEQ = 'seq' PAR = 'par' +# The name a policy's inference travels under. It takes the observation and answers a chunk: a list of +# actions, each naming command channels and carrying its own ``keys.ACTION_TIMESTAMP`` in seconds from the +# call. An empty list drops what is playing, and one action may come back bare. +INFER = 'infer' + class NotAnswered(RuntimeError): """The call has not answered yet. A read of an ``Answer`` raises this rather than waiting.""" @@ -32,66 +38,52 @@ def result(self) -> Any: class Runtime(ABC): - """What the framework offers one session. Every session gets its own. + """What the framework offers one episode: its work, started off the loop thread. - Closed before the session it serves: a call still in flight is using what that session holds. + Closed before the work it serves is released: a call still in flight is using what that work holds. """ @property @abstractmethod def fns(self) -> Mapping[str, Fn]: - """The policy's functions, under the names it declared them by.""" + """The episode's work, under the names the policy declared it by.""" class Session(ABC): - """Per-episode inference session. Created by ``Policy.new_session()``. - - Sessions hold per-episode state (trajectory buffers, latency tracking, etc.) - and are the primary interface for running inference. Call the session like - a function to get actions:: + """Per-episode control session, created by ``Policy.new_session``. - session = policy.new_session(context) - trajectory = session(obs, time_ns) + A session holds what one episode plays — the chunk in flight, the history a stack samples — and is + what a caller drives a robot through:: - **Plain-data contract**: sessions accept and return only plain data - (dicts, lists, numpy arrays, scalars). No tensors or custom objects. + session = policy.new_session(rt) + commands, resume_at_ns = session(obs, time_ns) - **Return contract**: ``list[dict] | None``. ``None`` means "no new - trajectory, keep executing the current one" — what a scheduling layer - answers while its chunk plays, and what a session answers while the - function it asked is still in flight. - An empty list means "stop whatever is executing now". A non-empty list is - a new trajectory. Single-action returns must be wrapped into a 1-element - list by the producer. + **Plain-data contract**: sessions accept and return only plain data (dicts, lists, numpy arrays, + scalars). No tensors or custom objects. """ @abstractmethod - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | None: - """Predict actions for the given observation, without waiting: heavy work belongs in - ``Policy.functions``. + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> tuple[Mapping[str, Any], int]: + """The commands to run now and the time this session wants its next call, without waiting: + heavy work belongs in ``Policy.episode``. ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. + ``commands`` names a command channel per entry; an empty mapping asks for nothing this call. + ``resume_at_ns`` is after ``time_ns``, on the same clock. The caller aims at it and may call + earlier. A session that waits for work it cannot time names a poll period of its own. """ - @property - def meta(self) -> dict[str, Any]: - """Session metadata (may include policy meta + per-session info).""" - return {} - def cancel(self): - """Drop any in-flight trajectory state. Layers that buffer/schedule a - trajectory (e.g. ``ChunkedSchedule``) should reset so the next call - triggers a fresh inference. Override and propagate via ``super().cancel()``. - """ + """Drop what is in flight, so the next call plans afresh. Propagate with ``super().cancel()``.""" return None def close(self): - """End this session and release per-episode resources.""" + """End this session and release what it holds.""" return None class DelegatingSession(Session): - """Session that delegates all methods to an inner session. Subclass and override what you need.""" + """Session that delegates all methods to an inner ``Session``. Subclass and override what you need.""" def __init__(self, inner: Session): self._inner = inner @@ -99,10 +91,6 @@ def __init__(self, inner: Session): def __call__(self, obs, time_ns): return self._inner(obs, time_ns) - @property - def meta(self): - return self._inner.meta - def cancel(self): self._inner.cancel() @@ -110,35 +98,33 @@ def close(self): self._inner.close() -class Policy(ABC): - """Factory for inference sessions. +class Policy: + """Factory for episodes. - A Policy holds shared resources (model weights, connections) and creates - per-episode ``Session`` instances. One Policy can serve multiple robots - by creating independent sessions. + A Policy holds what every episode shares — model weights, a connection, a subprocess — and opens the + per-episode work each one runs. One Policy serves several robots through independent episodes. It + declares the work of an episode, the session that plays one, or both. """ - @abstractmethod - def new_session(self, context: dict[str, Any] | None = None, rt: Runtime | None = None) -> Session: - """Create a new inference session for an episode. + @contextmanager + def episode(self, context: dict[str, Any] | None = None) -> Iterator[Mapping[str, Callable[..., Any]]]: + """The work of one episode, by name, and what it holds for as long as the episode runs. - Args: - context: The episode's task description. - rt: This session's runtime, serving ``functions``. ``None`` only where no caller supplied one. - A session that needs one refuses to open without it. + The framework serves the work as ``Runtime.fns`` and closes it after the episode. A policy that + answers a chunk declares that work under ``INFER``. ``context`` is the episode's task description. """ + yield {} - @property - def functions(self) -> Mapping[str, Callable[..., Any]]: - """The work this policy runs off the session's thread, by name. The framework serves it as ``rt.fns``.""" - return {} + def new_session(self, rt: Runtime) -> Session: + """The session that plays this episode, over the work ``rt`` serves.""" + raise NotImplementedError(f'{type(self).__name__} answers a chunk; put a ChunkPlayer above it to play one') @property def meta(self) -> dict[str, Any]: """Static metadata about this policy/model.""" return {} - def close(self): # noqa: B027 + def close(self): """Release shared resources (model weights, connections, etc.).""" @@ -148,12 +134,11 @@ class DelegatingPolicy(Policy): def __init__(self, inner: Policy): self._inner = inner - def new_session(self, context=None, rt=None): - return self._inner.new_session(context, rt) + def episode(self, context=None): + return self._inner.episode(context) - @property - def functions(self): - return self._inner.functions + def new_session(self, rt): + return self._inner.new_session(rt) @property def meta(self): @@ -174,12 +159,13 @@ class Layer: ``codec | layer`` all produce a Layer pipeline that ``wrap(policy)`` applies right-to-left:: - pipeline = TemporalStack(...) | ChunkedSchedule() | codec + pipeline = TemporalStack(...) | ChunkPlayer() | codec wrapped = pipeline.wrap(RemotePolicy(...)) **Extension points**: subclasses override *one* of ``make_session`` (the common case — transform one session's ``__call__``) or ``wrap`` (for - policy-level state across sessions, like composition). + policy-level state across sessions, like composition). A layer that sits under a ``ChunkPlayer`` + subclasses ``ChunkLayer`` instead: there is no session there, only the work ``INFER`` names. """ def wrap(self, policy: Policy) -> Policy: @@ -194,6 +180,10 @@ def make_session(self, inner: Session) -> Session: # it, so the name is written once and both sides of the wire read the same attribute. WIRE_NAME: ClassVar[str] + # Whether this layer turns the chunk work below it into a session. A ``ChunkLayer`` wraps that work, so + # a composition puts every one of them under the layer that plays it. + PLAYS_CHUNKS: ClassVar[bool] = False + def to_spec(self) -> dict[str, Any]: """Plain-data wire spec of this layer, for a server's local-stack declaration. @@ -218,24 +208,52 @@ def _layers(self) -> tuple: return (self,) -class _LayerPolicy(DelegatingPolicy): - """Policy produced by ``Layer.wrap()``. +class ChunkLayer(Layer): + """Layer under a ``ChunkPlayer``: it wraps the work ``INFER`` names, for one episode. - Delegates session creation to the layer's ``make_session`` and merges meta. + Below the player there is no session, so a chunk layer has nothing to cancel and no round to pace. It + sees the observation on the way down and the chunk on the way up. """ + def wrap(self, policy: Policy) -> Policy: + return _ChunkLayerPolicy(policy, self) + + @contextmanager + def episode_fn(self, infer: Callable[..., Any]) -> Iterator[Callable[..., Any]]: + """``infer`` with this layer's work around it, and what that takes for as long as the episode runs.""" + raise NotImplementedError('Override episode_fn or wrap') + + +class _WrappedPolicy(DelegatingPolicy): + """Policy produced by ``Layer.wrap()``: the layer's own metadata joins what it wraps.""" + def __init__(self, inner: Policy, layer: Layer): super().__init__(inner) self._layer = layer - def new_session(self, context=None, rt=None): - return self._layer.make_session(self._inner.new_session(context, rt)) - @property def meta(self): return self._inner.meta | self._layer.meta +class _LayerPolicy(_WrappedPolicy): + """Every session it creates goes through the layer's ``make_session``.""" + + def new_session(self, rt): + return self._layer.make_session(self._inner.new_session(rt)) + + +class _ChunkLayerPolicy(_WrappedPolicy): + """Its episode serves the layer's ``INFER`` in place of the one below it.""" + + _layer: ChunkLayer + + @contextmanager + def episode(self, context=None): + with self._inner.episode(context) as fns, self._layer.episode_fn(fns[INFER]) as infer: + yield {**fns, INFER: infer} + + class _ComposedLayer(Layer): """Composed pipeline of layers. Applies right-to-left.""" @@ -243,7 +261,12 @@ def __init__(self, components: tuple): self._components = components def wrap(self, policy: Policy) -> Policy: + played = False for component in reversed(self._components): + assert not (played and isinstance(component, ChunkLayer)), ( + f'compose {type(component).__name__} under the layer that plays the chunk, not above it' + ) + played = played or component.PLAYS_CHUNKS policy = component.wrap(policy) return policy diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index c701dc24c..31a2dc774 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -10,6 +10,7 @@ """ import collections.abc as cabc +from contextlib import contextmanager from dataclasses import replace from functools import partial from typing import Any, final @@ -24,7 +25,7 @@ from positronic.drivers.roboarm import command from positronic.drivers.roboarm.ik import assert_default_frame, change_frame, ee_frame from positronic.drivers.roboarm.models import DEFAULT_FRAME -from positronic.policy.base import PAR, SEQ, DelegatingSession, Layer, Session, _ComposedLayer +from positronic.policy.base import PAR, SEQ, ChunkLayer, Layer, _ComposedLayer from positronic.utils import merge_dicts _QUAT = geom.Rotation.Representation.QUAT @@ -48,7 +49,7 @@ def lerobot_action(dim: int) -> dict[str, Any]: return {'shape': (dim,), 'names': ['actions'], 'dtype': 'float32'} -class Codec(Layer): +class Codec(ChunkLayer): """Base class for observation/action codecs. Subclasses override ``encode`` (observation encoding) and/or ``_decode_single`` @@ -81,8 +82,13 @@ def training_encoder(self) -> EpisodeTransform: def meta(self) -> dict: return {} - def make_session(self, inner: Session): - return _CodecSession(inner, self) + @contextmanager + def episode_fn(self, infer): + yield partial(self._coded, infer) + + def _coded(self, infer, obs): + """The chunk ``infer`` answers for the encoded observation, decoded.""" + return self.decode(infer(self.encode(obs))) @final def __or__(self, other): @@ -100,25 +106,6 @@ def __and__(self, other): return NotImplemented -class _CodecSession(DelegatingSession): - """Session wrapped with a codec: encodes observations, decodes actions.""" - - def __init__(self, inner: Session, codec: 'Codec'): - super().__init__(inner) - self._codec = codec - - def __call__(self, obs, time_ns): - encoded = self._codec.encode(obs) - action = self._inner(encoded, time_ns) - if action is None: - return None - return self._codec.decode(action) - - @property - def meta(self): - return self._inner.meta | self._codec.meta - - def _meta_conflicts(left: dict, right: dict, prefix: str = '') -> list[str]: """``key: left != right`` for every leaf the two metas declare differently. Nested dicts merge per key, so only leaves can conflict.""" @@ -447,7 +434,7 @@ class RestrictImageSize(Codec): Declared left of the ``remote`` marker, so the rig applies it before sending:: - ChunkedSchedule() | RestrictImageSize() | remote | codec | source + ChunkPlayer() | RestrictImageSize() | remote | codec | source """ WIRE_NAME = 'restrict_image_size' diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index 93a7f4952..daab032f1 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -9,16 +9,7 @@ from functools import partial from typing import Any -from positronic.policy.base import ( - Answer, - DelegatingPolicy, - DelegatingSession, - Fn, - NotAnswered, - Policy, - Runtime, - Session, -) +from positronic.policy.base import Answer, Fn, NotAnswered, Runtime class Executor(Runtime): @@ -68,8 +59,6 @@ def in_flight(self) -> bool: @property def owes_an_answer(self) -> bool: """Whether any call's answer has still to be read, whether or not that call has landed.""" - # TODO(#661): a caller polls this because a session cannot say when it wants the next call. Rung 7 - # gives the session ``resume_at``, and the poll goes with it. with self._lock: return bool(self._unread) @@ -112,46 +101,3 @@ def close(self) -> None: # and the log is the only place the failure can go. if (exc := answer.failure()) is not None: logging.error(f'The function {answer.name} failed and no caller read its answer: {exc}') - - -class _BlockingPolicy(DelegatingPolicy): - class _Session(DelegatingSession): - def __init__(self, inner: Session, rt: Executor): - super().__init__(inner) - self._rt = rt - - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | None: - # The inner session reads an answer only on a later call. A test of ``in_flight`` would exit - # on a call that lands while the session call runs, leaving its answer unread. - while (actions := self._inner(obs, time_ns)) is None and self._rt.owes_an_answer: - self._rt.wait() - return actions - - def close(self): - # The runtime closes first: a call in flight is still using what the session holds. - self._rt.close() - self._inner.close() - - def new_session(self, context=None, rt=None) -> Session: - assert rt is None, 'a blocking policy serves its own functions; nothing above it runs them' - own = Executor(self._inner.functions) - try: - return _BlockingPolicy._Session(self._inner.new_session(context, own), own) - except BaseException: - own.close() - raise - - @property - def functions(self) -> Mapping[str, Callable[..., Any]]: - return {} - - -def blocking(policy: Policy) -> Policy: - """``policy`` with its heavy work waited out: a session answers in the call that asked. - - For a caller with no control loop to give the time back to — a server request, a warmup, a probe. - Layers wrap the result rather than the other way round, so each sees one call per answer. Layers that - ``policy`` composes itself are inside, so those still run once per call, each with the ``time_ns`` of - the call that asked. - """ - return _BlockingPolicy(policy) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 9aabb524e..7a35cf20b 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,6 +1,6 @@ +import contextlib import time -from collections import deque -from collections.abc import Generator, Iterator +from collections.abc import Generator, Iterator, Mapping from typing import Any import numpy as np @@ -16,34 +16,32 @@ from positronic.policy.executor import Executor from positronic.utils import flatten_dict, frozen_view -# How far from now an action may be scheduled: past any real chunk, short of the decades a rig-side stack is -# off by when it leaves timestamps chunk-relative or anchors them twice. -MAX_ACTION_SKEW_SEC = 60.0 - -# How long a real-time round may last when no waypoint is due sooner. It bounds how late a call is noticed, -# and with it the granularity every command timestamp is quantized to. -POLL_PERIOD_SEC = 0.01 +# How long the harness waits between looks at a call it made. +WAIT_PERIOD_SEC = 0.01 +# The shortest and the longest real-time round. The floor caps how fast a session drives the loop, the +# ceiling how late the harness reads ``done`` and the stop signal. +MIN_ROUND_SEC = 0.001 +MAX_ROUND_SEC = 1.0 class _EpisodeInference: - """One episode's policy session, and the runtime that serves the policy's functions to it.""" + """One episode: the work the policy opens, the runtime that serves it, and the session that plays it.""" def __init__(self, policy: Policy, context: dict[str, Any], charges_wall_time: bool, clock: pimm.Clock) -> None: self._charges_wall_time = charges_wall_time self._clock = clock # One instant on two clocks, so ``wait`` adds a wall duration to a world instant. self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() - self._runtime = Executor(policy.functions) + self._closing = contextlib.ExitStack() try: - self._session = policy.new_session(context, self._runtime) + self._runtime = Executor(self._closing.enter_context(policy.episode(context))) + self._closing.callback(self._runtime.close) + self._session = policy.new_session(self._runtime) + self._closing.callback(self._session.close) except BaseException: - self._runtime.close() + self._closing.close() raise - @property - def meta(self) -> dict[str, Any]: - return self._session.meta - @staticmethod def _owned(obs: dict[str, Any]) -> dict[str, Any]: """The observation with its arrays copied, so nothing rewrites what a function is still reading. @@ -53,12 +51,14 @@ def _owned(obs: dict[str, Any]) -> dict[str, Any]: """ return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: dict[str, Any]) -> tuple[Mapping[str, Any], int]: now_ns = self._clock.now_ns() # A call that joins work already in flight keeps its anchor, so the trial pays for that work one time. if not self._runtime.in_flight: self._t0_ns, self._wall_t0 = now_ns, time.monotonic() - return self._session(frozen_view(self._owned(obs)), now_ns) + commands, resume_at_ns = self._session(frozen_view(self._owned(obs)), now_ns) + assert resume_at_ns > now_ns, f'resume time must be in the future: {resume_at_ns} <= {now_ns}' + return commands, resume_at_ns def wait(self, should_stop: pimm.SignalReceiver[bool]) -> None: """Wait for the function in flight, for as long as the trial charges the loop for it.""" @@ -70,15 +70,14 @@ def wait(self, should_stop: pimm.SignalReceiver[bool]) -> None: # A trial that pays nothing waits the function out. It waits in steps, because a model that never # answers must not also keep the world from coming down. while self._runtime.in_flight and not should_stop.value: - self._runtime.wait(POLL_PERIOD_SEC) + self._runtime.wait(WAIT_PERIOD_SEC) def close(self) -> None: - """Close the runtime, then the session it was serving. + """End the session, then the runtime, then the work the episode opened. - Until ``Executor.close`` returns, the function in flight still holds the session's websocket or model. + Until ``Executor.close`` returns, the function in flight still holds the episode's websocket or model. """ - self._runtime.close() - self._session.close() + self._closing.close() class _EpisodeTelemetry: @@ -109,6 +108,7 @@ def start_rollout(self, virtual_now: float) -> None: self._virtual_start = virtual_now def step(self) -> None: + """Count one control round: the harness read an observation, called the session and emitted its answer.""" self._steps += 1 def end(self, virtual_now: float) -> None: @@ -144,12 +144,12 @@ def _end_span(self) -> None: class Harness(pimm.ControlSystem): - """Control system that runs the episode lifecycle and plays the policy's trajectory to the drivers. + """Control system that runs the episode lifecycle and emits what the policy commands. - The layer owns the trajectory, the harness plays it, one command per channel per round. The session is - called on the loop thread and answers inside the round; the functions it starts run on the runtime's own, - so playing continues while the model runs. That work costs the trial either the wall time it took or - nothing, with the world held still for it. + The session owns the trajectory and answers the commands to run now; the harness emits them and comes + back when the session asked. The session is called on the loop thread and answers inside the round; the + functions it starts run on the runtime's own, so playing continues while the model runs. That work costs + the trial either the wall time it took or nothing, with the world held still for it. An episode runs one ``Task``, asked for by a ``perform_task`` call and answered with the terminal payload it ended on. The task's ``timeout_sec`` bounds it and a truthy ``done`` within budget ends it @@ -169,14 +169,14 @@ def __init__(self, policy: Policy, embodiment: Embodiment, *, static_meta: dict[ self._retired: _EpisodeInference | None = None # ``task.timeout_sec``, armed per episode; a task without one has no deadline and ends on ``done`` alone. self._deadline_ns: int | None = None + # When the live session asked for its next call. ``None`` while no session has answered. + self._resume_at_ns: int | None = None # Wall-clock telemetry for the live rollout, opened under ``--timing`` and inert otherwise. self._telemetry = _EpisodeTelemetry() self.observations = pimm.ReceiverDict(self, names=embodiment.observations) self.commands = pimm.EmitterDict(self, names=embodiment.commands) self.prepare = pimm.calls.CallerDict[Any, None](self, names=embodiment.prepare_handlers) - # Each channel's waypoints not yet played, stamped with absolute clock ns and ascending. - self._schedules: dict[str, deque[tuple[int, Any]]] = {name: deque() for name in embodiment.commands} # One episode per call, answered with the terminal payload it ended on. self.perform_task = pimm.calls.ControlSystemHandler[Task, dict[str, Any]](self) @@ -210,10 +210,7 @@ def _build_episode_meta(self) -> dict[str, Any]: meta[keys.EVAL_CHARGE_INFERENCE_TIME] = self._charges_wall_time if self._task.timeout_sec is not None: # the recorder takes no nulls, and an unbounded episode has none meta[keys.EVAL_TIMEOUT] = self._task.timeout_sec - # ``policy.meta`` is the static baseline; the session overlays per-episode specifics (e.g. the - # sampled sub-policy) and wins on conflict. - session_meta = self.policy.meta | (self._inference.meta if self._inference else {}) - for k, v in flatten_dict(session_meta).items(): + for k, v in flatten_dict(self.policy.meta).items(): meta[f'{keys.POLICY_META}.{k}'] = v meta.update(self._task.meta) meta[keys.TASK] = self._task.instruction @@ -231,28 +228,29 @@ def _ready(self, should_stop: pimm.SignalReceiver, args: dict[str, Any]) -> Gene raise ValueError(f'{unknown} is not something {rig} readies; it readies {sorted(self.prepare)}') ready = pimm.calls.all_of([self.prepare[name](arg) for name, arg in args.items()]) while not ready.done() and not should_stop.value: - yield pimm.Yield() if self._embodiment.simulated else pimm.Sleep(POLL_PERIOD_SEC) + yield pimm.Yield() if self._embodiment.simulated else pimm.Sleep(WAIT_PERIOD_SEC) # An episode must not open on a rig that never got ready, and its asker must hear that rather than wait if not ready.done(): raise RuntimeError('The world stopped before every device was ready') ready.result() def _pace(self, clock: pimm.Clock) -> pimm.Command: - """Sim: yield, so the simulator's control-period sleep is the sole time-master and the policy reads - each observation instantly. Real: sleep to the next waypoint, capped at the poll period, so a - waypoint is emitted at its own time and a round rarely finds more than one due.""" + """The command that ends this round: a yield in sim, where the simulator's control-period sleep is + the sole time-master, and a sleep to the moment the live session asked for on a real rig.""" if self._embodiment.simulated: return pimm.Yield() - due = min((sched[0][0] for sched in self._schedules.values() if sched), default=None) - if due is None: - return pimm.Sleep(POLL_PERIOD_SEC) - return pimm.Sleep(min(POLL_PERIOD_SEC, max(due - clock.now_ns(), 1) / 1e9)) + if self._resume_at_ns is None: + return pimm.Sleep(WAIT_PERIOD_SEC) + until_ns = self._resume_at_ns if self._deadline_ns is None else min(self._resume_at_ns, self._deadline_ns) + due_sec = (until_ns - clock.now_ns()) / 1e9 + return pimm.Sleep(min(max(due_sec, MIN_ROUND_SEC), MAX_ROUND_SEC)) def _retire_inference(self) -> None: """Let go of this episode's inference, keeping it for ``_reap_inference``: ending an episode must not wait for a model that hangs.""" if self._inference is not None: self._retired, self._inference = self._inference, None + self._resume_at_ns = None def _reap_inference(self) -> None: if self._retired is not None: @@ -262,13 +260,12 @@ def _reap_inference(self) -> None: def _finalize_recording( self, clock: pimm.Clock, payload: dict[str, Any] | None = None ) -> Generator[pimm.Command, None, None]: - """Commit the live episode: cancel the in-flight chunk, stop the recorder — stamping the + """Commit the live episode: end the session playing it, stop the recorder — stamping the episode's full static meta (plus any terminal payload) — then close its span.""" self._set_deadline(None) # Stamped before the inference is retired: the meta overlays what its session reports. stop = DsWriterCommand.STOP({**self._build_episode_meta(), **(payload or {})}) - for schedule in self._schedules.values(): # devices hold their last commanded position - schedule.clear() + # Retiring the session ends the chunk it was playing; devices hold their last commanded position. self._retire_inference() self.ds_command.emit(stop) virtual_now = clock.now() # before the round below, whose sim-clock advance belongs to no rollout @@ -318,8 +315,8 @@ def _end_episode( """Close the live episode: finalize the recording, put the rig back, retire the session, hand the terminal back to whoever asked for the episode. - The inference is retired rather than closed here, so a ``RemoteSession``'s websocket outlives the - function still using it. + The inference is retired rather than closed here, so the episode's websocket outlives the function + still using it. """ yield from self._finalize_recording(clock, payload) # A powered arm holds the policy's last setpoint until the next trial, so each device the trial placed @@ -376,49 +373,14 @@ def _infer(self, inference: _EpisodeInference, clock: pimm.Clock, should_stop: p obs = self._build_obs(clock) except pimm.NoValueException: return # no function is in flight yet, so this skips no wait - if (trajectory := inference(obs)) is not None: - self._reschedule(trajectory, clock) - inference.wait(should_stop) - - @staticmethod - def _assert_anchored(trajectory: list[dict[str, Any]], now: float) -> None: - """Reject a chunk whose timestamps are not times on the harness clock.""" - skew = max((abs(action[keys.ACTION_TIMESTAMP] - now) for action in trajectory), default=0.0) - if skew > MAX_ACTION_SKEW_SEC: - raise ValueError( - f'Action scheduled {skew:.0f}s from now, over the {MAX_ACTION_SKEW_SEC:.0f}s bound: the ' - f'rig-side stack is not anchoring chunks to the harness clock' - ) - - def _reschedule(self, trajectory: list[dict[str, Any]], clock: pimm.Clock) -> None: - """Replace the schedule being played with the session's trajectory. Every channel it names gets that - channel's waypoints; one it omits is cleared and holds. The timestamps are already absolute, stamped - by the scheduling layer against the harness clock. - """ - if self._deadline_ns is not None and clock.now_ns() >= self._deadline_ns: - # The world reached the deadline while the function was in flight, so its chunk is dropped rather - # than placed past the point the trial advertises it stops at; ``_run`` finishes the trial next round. - return - self._assert_anchored(trajectory, clock.now()) + commands, self._resume_at_ns = inference(obs) self._telemetry.step() - # Layers time actions in float seconds; the schedules and every pimm channel are in ns. - for name, schedule in self._schedules.items(): - schedule.clear() - schedule.extend((int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in trajectory if name in a) - - def _play(self, clock: pimm.Clock) -> None: - """Emit each channel's command due this round, and nothing on a channel with none. - - A channel with several waypoints due emits the last: exact for an absolute setpoint, lossy for a - relative one. Pacing keeps one due per round wherever a round is shorter than the waypoint spacing. - """ - now_ns = clock.now_ns() - for name, schedule in self._schedules.items(): - value = None - while schedule and schedule[0][0] <= now_ns: - value = schedule.popleft()[1] - if value is not None: - self.commands[name].emit(value) + # The world can reach the deadline while the session runs, and a command placed after the point the + # trial advertises it stops at outlives the trial; ``_run`` finishes it next round. + if self._deadline_ns is None or clock.now_ns() < self._deadline_ns: + # The key-filtered demux: a command this rig declares no channel for reaches no driver. + self._emit({name: commands[name] for name in self.commands if name in commands}) + inference.wait(should_stop) def _trial_terminal(self, done: pimm.Message[dict] | None, clock: pimm.Clock) -> dict[str, Any] | None: """The terminal static payload if the live trial has ended this round, else ``None``. @@ -472,7 +434,6 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ yield from self._begin_episode(clock, should_stop, call) elif manual is not None: self._emit(manual) - self._play(clock) yield self._pace(clock) if self._inference is not None: diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 530039e06..be0c253a3 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -1,17 +1,22 @@ -"""Composable policy layers — scheduling, fault handling and temporal frame stacking. +"""Composable policy layers — chunk playback, fault handling and temporal frame stacking. -Layers are serving-time concerns wrapped around a policy with ``|`` (left is outermost). Most read time -from the observation (``obs_time_ns``); ``ChunkedSchedule`` anchors a chunk with the ``time_ns`` of the -call, which is the caller's clock reading in nanoseconds. +Layers are serving-time concerns wrapped around a policy with ``|`` (left is outermost). ``ChunkPlayer`` +turns the chunk the policy answers into the commands a caller runs; the layers above it read time from the +observation (``obs_time_ns``) or from the call's ``time_ns``, both in nanoseconds. """ from collections import deque +from typing import Any, NamedTuple import numpy as np from positronic import keys from positronic.drivers.roboarm import RobotStatus -from positronic.policy.base import DelegatingSession, Layer, Session +from positronic.policy.base import INFER, Answer, DelegatingPolicy, DelegatingSession, Fn, Layer, Policy, Session + +# How far from the call a waypoint may sit: past any real chunk, short of the decades a chunk is off by +# when it reaches the player already anchored. +MAX_ACTION_SKEW_SEC = 60.0 def _obs_time(obs) -> float: @@ -37,71 +42,122 @@ def _arms_available(obs) -> bool: class StopOnFault(Layer): """Stop the arm while it will not take a command, and plan afresh once it will. - An arm the driver has taken, or that is faulted, is not tracking the plan it was given: this answers the - empty trajectory and resets the sessions below. It goes outside the scheduling layer, which would - otherwise answer "keep playing" without seeing the status. Every arm is checked, so a bimanual rig stops - on either. + An arm the driver has taken, or that is faulted, is not tracking the plan it was given: this commands + nothing and resets the sessions below, which drops the chunk being played. It goes outside the player, + which would otherwise keep emitting waypoints without seeing the status. Every arm is checked, so a + bimanual rig stops on either. """ WIRE_NAME = 'stop_on_fault' + POLL_SEC = 0.01 class _Session(DelegatingSession): def __call__(self, obs, time_ns): if _arms_available(obs): return self._inner(obs, time_ns) self.cancel() - return [] + return {}, time_ns + int(StopOnFault.POLL_SEC * 1e9) - def make_session(self, inner: Session): + def make_session(self, inner: Session) -> Session: return StopOnFault._Session(inner) def to_spec(self): return {'name': self.WIRE_NAME} -class ChunkedSchedule(Layer): - """Wait for the current trajectory to finish before calling the inner policy again. - - Owns relative→absolute time conversion: the sessions below (codecs, models) emit relative timestamps; - this layer anchors them to the call's ``time_ns``, in the seconds a timestamp is written in. Returns - ``None`` ("keep executing the current trajectory") until the last action's timestamp is reached, then - calls the inner policy. +class ChunkPlayer(Layer): + """Hold the chunk ``INFER`` answers, anchor it to the call that received it, and emit each waypoint at + its own time. - The call's ``time_ns`` and the observation's ``obs_time_ns`` must be readings of one clock: the anchor - comes from the first, and the test for a complete chunk uses the second. + The player is the bottom session of a chunk policy: it turns the work below into the commands a caller + runs, so the layers under it wrap that work rather than a session. """ - WIRE_NAME = 'chunked_schedule' + WIRE_NAME = 'chunk_player' + PLAYS_CHUNKS = True + POLL_SEC = 0.01 - class _Session(DelegatingSession): - """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" + class _Policy(DelegatingPolicy): + """The player is the session of a chunk policy, so nothing below it opens one.""" - def __init__(self, inner: Session): - super().__init__(inner) - self._trajectory_end: float | None = None + def new_session(self, rt): + return ChunkPlayer._Session(rt.fns[INFER]) + + class _Session(Session): + class _Waypoint(NamedTuple): + cmd: dict[str, Any] + time_ns: int + + def __init__(self, infer: Fn): + self._infer = infer + self._waypoints: deque[ChunkPlayer._Session._Waypoint] = deque() + # The one call this session keeps in flight, and whether a ``cancel`` has orphaned the chunk it + # will bring back. + self._answer: Answer | None = None + self._orphaned = False def __call__(self, obs, time_ns): - if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: - return None - result = self._inner(obs, time_ns) - if result is not None: - # A single-action session may return a bare dict, and a no-codec path may omit - # ``timestamp`` (servers can stamp/truncate themselves); normalize both so an - # immediate action executes instead of raising. - if isinstance(result, dict): - result = [result] - # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). - anchor = time_ns / 1e9 - result = [{**r, keys.ACTION_TIMESTAMP: anchor + r.get(keys.ACTION_TIMESTAMP, 0.0)} for r in result] - self._trajectory_end = result[-1][keys.ACTION_TIMESTAMP] if result else None - return result + """Plays the chunk it holds; asks for the next one in the call that drains it.""" + if not self._waypoints or self._waypoints[-1].time_ns <= time_ns: + if self._answer is None: + self._answer = self._infer(obs) + if self._answer.done(): + chunk = self._take() + if chunk is not None: + self._load(chunk, time_ns) + commands: dict[str, Any] = {} + while self._waypoints and self._waypoints[0].time_ns <= time_ns: + commands.update(self._waypoints.popleft().cmd) + if not self._waypoints: + return commands, time_ns + int(ChunkPlayer.POLL_SEC * 1e9) + return commands, self._waypoints[0].time_ns + + def _take(self) -> list[dict[str, Any]] | dict[str, Any] | None: + """The chunk the call brought back, and ``None`` for one a ``cancel`` orphaned. + + An orphaned call is read too, so its failure reaches the caller. The state clears before that + read, so a cancel ends with the call it was made against and never drops the chunk after it. + """ + assert self._answer is not None, 'only a call this session made brings a chunk back' + answer, orphaned = self._answer, self._orphaned + self._answer, self._orphaned = None, False + chunk = answer.result() + return None if orphaned else chunk + + def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> None: + """Anchor ``chunk`` to ``time_ns`` and hold it. + + A single-action policy may answer a bare dict, and a no-codec path may omit ``timestamp`` + (servers can stamp and truncate themselves); both are normalized here so an immediate action + plays instead of raising. + """ + if isinstance(chunk, dict): + chunk = [chunk] + skew = max((abs(action.get(keys.ACTION_TIMESTAMP, 0.0)) for action in chunk), default=0.0) + if skew > MAX_ACTION_SKEW_SEC: + raise ValueError( + f'Action scheduled {skew:.0f}s from the call, over the {MAX_ACTION_SKEW_SEC:.0f}s bound: ' + f'the work below is timing actions against a clock of its own' + ) + # The single explicit seconds->ns seam: the work below times actions in float seconds, and + # every pimm channel is in ns. The offset converts, not the sum, so a waypoint at 0.0 lands on the + # call itself whatever the clock reads. A waypoint naming no channel — the codecs' end-of-chunk + # sentinel — commands nothing and states where the chunk ends. + self._waypoints = deque( + self._Waypoint( + {name: value for name, value in action.items() if name != keys.ACTION_TIMESTAMP}, + time_ns + int(action.get(keys.ACTION_TIMESTAMP, 0.0) * 1e9), + ) + for action in chunk + ) def cancel(self): - self._trajectory_end = None - super().cancel() + self._waypoints.clear() + # The chunk in flight describes a world that has gone. The call is still read, for its failure. + self._orphaned = self._answer is not None - def make_session(self, inner: Session): - return ChunkedSchedule._Session(inner) + def wrap(self, policy: Policy) -> Policy: + return ChunkPlayer._Policy(policy) def to_spec(self): return {'name': self.WIRE_NAME} @@ -156,7 +212,7 @@ class TemporalStack(Layer): A model that conditions on a short window of history (e.g. DreamZero's video context) needs several samples spanning the just-executed chunk at the cadence seen in training, but the harness only forwards an observation to the policy at re-query boundaries. This layer sits outside the - scheduling layer so it sees every control tick: it records the named ``keys`` and substitutes, for + player so it sees every control tick: it records the named ``keys`` and substitutes, for each, a ``(len(offsets_sec), ...)`` stack sampled at ``offsets_sec`` (ascending negative seconds relative to now), so every stacked step carries its own value at that time rather than the current one repeated across history. @@ -196,7 +252,7 @@ def __init__(self, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_st 'in-range targets and the stack would be empty' ) - def make_session(self, inner: Session): + def make_session(self, inner: Session) -> Session: return TemporalStack._Session(inner, self._keys, self._offsets_sec, self._pad_start) def to_spec(self): diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index a49d39374..990485970 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -14,8 +14,8 @@ pipeline = rec.tap('raw') | codec | rec.tap('server') policy = pipeline.wrap(remote_policy) -- the ``raw`` tap (outermost) logs the observation as received and the final action - chunk; +- the ``raw`` tap (outermost) logs the observation as received and, above a ``ChunkPlayer``, + the command of each round rather than a chunk; - the ``server`` tap (innermost, next to the remote policy) logs the observation as sent to the server and the chunk as received back. @@ -42,16 +42,18 @@ both sides, the later write overwrites at that timestamp. TODO(#661): this module becomes the recording the framework offers a session, and stops being a layer a -pipeline composes. The framework then records every boundary it carries, a session appends series of its -own, and a served function records against the call it answers. A tap around a session cannot do the -last of those: it sees an answer on a later call than the observation that produced it. +pipeline composes. The framework then records every boundary it carries, and a session appends series of +its own. """ +import contextvars import itertools -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from contextlib import contextmanager from datetime import datetime +from functools import partial from pathlib import Path -from typing import Any +from typing import Any, NamedTuple import numpy as np import rerun as rr @@ -60,12 +62,18 @@ from positronic import geom from positronic import keys as obs_keys from positronic.drivers.roboarm import command as roboarm_command -from positronic.policy.base import DelegatingSession, Layer, Session +from positronic.policy.base import ChunkLayer, DelegatingSession, Layer, Session from positronic.policy.codec import is_action from positronic.utils.rerun_compat import log_numeric_series, set_timeline_sequence, set_timeline_time DEFAULT_TIMELINES = {'wall_time': obs_keys.WALL_TIME_NS, 'obs_time': obs_keys.OBS_TIME_NS} +# The timeline values of the inference being logged, set by the tap that opens the round. A runtime runs a +# function under a copy of the context the call was made in, so a tap on a worker thread reads them too. +_TIMELINE_VALUES: contextvars.ContextVar[dict[str, Any] | None] = contextvars.ContextVar( + 'recording_timeline_values', default=None +) + # Process-wide episode counter so files stay unique even across concurrent # ``Recorder`` instances (e.g. one per websocket session on a server). _EPISODE_COUNTER = itertools.count(1) @@ -303,11 +311,11 @@ class Recorder: Taps share the recorder's current episode stream, so taps placed at different points in one pipeline write to the same recording. The episode boundary is - tracked by a live-session counter: the first tap session to start (when none are - active) opens a fresh ``.rrd``; later taps in the same episode write to it; each - ``close()`` decrements, and the next session opened after the count returns to - zero starts the next file. Episodes are assumed to run one at a time (sessions on - one recorder do not overlap). + tracked by a live-tap counter: the first tap to open (when none are active) + opens a fresh ``.rrd``; later taps in the same episode write to it; each tap + that closes decrements, and the next tap opened after the count returns to zero + starts the next file. Episodes are assumed to run one at a time (taps on one + recorder do not overlap). ``timelines`` maps rerun timeline names to observation keys. The values are read once per inference at the outermost tap and reused by inner taps so every tap @@ -332,8 +340,13 @@ def __init__( self._path3d_paths: list[str] = [] self._series_paths: list[str] = [] - def tap(self, name: str) -> '_RecordingTap': - return _RecordingTap(self, name) + def tap(self, name: str) -> '_CommandTap': + """A tap above a ``ChunkPlayer``, logging the command of each round.""" + return _CommandTap(self, name) + + def chunk_tap(self, name: str) -> '_ChunkTap': + """A tap under a ``ChunkPlayer``, logging the observation the model sees and the chunk it answers.""" + return _ChunkTap(self, name) @property def stream(self) -> rr.RecordingStream | None: @@ -354,38 +367,38 @@ def _release_stream(self) -> None: self._live -= 1 -class _RecordingTap(Layer): - """A named tap. Wraps a single session to log its observations and actions.""" +class _Frame(NamedTuple): + """What one call is logged against: its observation, the rerun timelines it set, and its step.""" - def __init__(self, rec: Recorder, name: str): - self._rec = rec - self._name = name + obs: Mapping[str, Any] + timelines: dict[str, Any] + step: int - def make_session(self, inner: Session) -> Session: - stream = self._rec._open_stream() - return _RecordingTapSession(inner, self._rec, self._name, stream) +class _TapLog: + """Logs the observation flowing down and the actions flowing up at one point. -class _RecordingTapSession(DelegatingSession): - """Logs the observation flowing down and the action chunk flowing up at one point.""" + The timeline values are read once per inference, by the tap that opens the round, and every tap under + it stamps that inference identically. A tap whose work runs on a worker thread reads them too: the + runtime copies the calling context, and the values travel in it. + """ - def __init__(self, inner: Session, rec: Recorder, name: str, stream: rr.RecordingStream): - super().__init__(inner) + def __init__(self, rec: 'Recorder', name: str): self._rec = rec self._name = name - self._stream = stream + self._stream = rec._open_stream() self._step = 0 - @property - def meta(self): - return {**self._inner.meta, 'recording.rrd': str(self._rec._rrd_path)} + def close(self) -> None: + """Give up this tap's hold on the episode's recording.""" + self._rec._release_stream() - def _set_timelines(self) -> None: - for timeline, value in self._rec._timeline_values.items(): + def _set_timelines(self, frame: _Frame) -> None: + for timeline, value in frame.timelines.items(): set_timeline_time(timeline, value) - set_timeline_sequence('step', self._step) + set_timeline_sequence('step', frame.step) - def _log(self, prefix: str, data: dict) -> None: + def _log(self, prefix: str, data: Mapping[str, Any]) -> None: """Recursively log obs *data* under *prefix*, recording entity paths on the Recorder.""" for key, value in data.items(): if key.endswith('_time_ns') or isinstance(value, str): @@ -400,20 +413,21 @@ def _log(self, prefix: str, data: dict) -> None: log_numeric_series(path, num) self._rec._numeric_paths.append(path) - def _log_action_chunk(self, prefix: str, actions: list[dict], obs: dict) -> None: # noqa: C901 + def _log_action_chunk(self, prefix: str, actions: list[dict], frame: _Frame) -> None: # noqa: C901 """Log the action chunk as an enriched 3D trajectory + ``action_time`` time series.""" # Skip validity sentinels: they carry no command to plot and would flip the ``all(... in a)`` checks below. actions = [a for a in actions if is_action(a)] horizon = _horizon(actions) - tv = self._rec._timeline_values + tv = frame.timelines base_ns = int(tv.get('obs_time', tv.get('wall_time', next(iter(tv.values()), 0)))) grip = ( _stack_numeric([a[obs_keys.TARGET_GRIP] for a in actions]) if all(obs_keys.TARGET_GRIP in a for a in actions) else None ) - actual_pos = obs.get(obs_keys.EE_POSE) if isinstance(obs, Mapping) else None - actual_grip = obs.get(obs_keys.GRIP) if isinstance(obs, Mapping) else None + obs = frame.obs + actual_pos = obs.get(obs_keys.EE_POSE) + actual_grip = obs.get(obs_keys.GRIP) # Under TemporalStack these arrive as (T, 7) / (T,) stacks; the overlay draws the current pose, # which is the last frame (offsets end at 0 = now), mirroring the image collapse in `_as_image`. if actual_pos is not None: @@ -458,38 +472,92 @@ def _send_blueprint(self) -> None: if bp is not None: rr.send_blueprint(bp) - def __call__(self, obs, time_ns): - rec = self._rec - outermost = rec._depth == 0 - if outermost: - rec._timeline_values = {t: obs[k] for t, k in rec._timelines.items() if k in obs} - rec._depth += 1 + @contextmanager + def round(self, obs: Mapping[str, Any]) -> Iterator[_Frame]: + """The frame this call is logged against, with its observation already logged.""" + values = _TIMELINE_VALUES.get() + token = None + if values is None: + values = {t: obs[k] for t, k in self._rec._timelines.items() if k in obs} + token = _TIMELINE_VALUES.set(values) try: + frame = _Frame(obs, dict(values), self._step) with self._stream: - self._set_timelines() + self._set_timelines(frame) self._log(self._name, obs) - - actions = self._inner(obs, time_ns) - - if actions is not None: - with self._stream: - self._set_timelines() - # TODO(#661): the chunk is logged against the observation and the timelines of this - # call, which a session that answers from a served function makes a later call than the - # one the chunk was computed from. See the module docstring. - self._log_action_chunk(self._name, actions, obs) - # Send a combined blueprint (all taps' paths) once, from the outermost - # tap, after inner taps have logged their first obs. - if outermost and self._step == 0: - with self._stream: - self._send_blueprint() - self._step += 1 - return actions + yield frame finally: - rec._depth -= 1 - if rec._depth == 0: - rec._timeline_values = {} + if token is not None: + _TIMELINE_VALUES.reset(token) + + def actions(self, actions: list[dict], frame: _Frame) -> None: + if actions: + with self._stream: + self._set_timelines(frame) + self._log_action_chunk(self._name, actions, frame) + + def end_step(self) -> None: + """Close the round: send the combined blueprint once, from the tap that opened the round.""" + if _TIMELINE_VALUES.get() is None and self._step == 0: + with self._stream: + self._send_blueprint() + self._step += 1 + + +def _log_infer(log: _TapLog, infer, obs): + """``infer`` with its observation and the chunk it answers logged.""" + with log.round(obs) as frame: + chunk = infer(obs) + log.actions([dict(chunk)] if isinstance(chunk, Mapping) else (chunk or []), frame) + log.end_step() + return chunk + + +class _CommandTapSession(DelegatingSession): + def __init__(self, inner: Session, log: _TapLog): + super().__init__(inner) + self._log = log + + def __call__(self, obs, time_ns): + with self._log.round(obs) as frame: + answer = self._inner(obs, time_ns) + commands, _ = answer + self._log.actions([dict(commands)] if commands else [], frame) + self._log.end_step() + return answer def close(self): - super().close() - self._rec._release_stream() + self._inner.close() + self._log.close() + + +class _Tap: + """What a tap holds: the recorder and the name it logs under.""" + + def __init__(self, rec: 'Recorder', name: str): + self._rec = rec + self._name = name + + @property + def meta(self) -> dict[str, Any]: + return {} if self._rec._rrd_path is None else {'recording.rrd': str(self._rec._rrd_path)} + + +class _CommandTap(_Tap, Layer): + """A tap above a ``ChunkPlayer``: it logs the command of each round, so a plot reads a point per round + rather than a chunk per inference.""" + + def make_session(self, inner: Session) -> Session: + return _CommandTapSession(inner, _TapLog(self._rec, self._name)) + + +class _ChunkTap(_Tap, ChunkLayer): + """A tap under a ``ChunkPlayer``: it logs the observation the model sees and the chunk it answers.""" + + @contextmanager + def episode_fn(self, infer): + log = _TapLog(self._rec, self._name) + try: + yield partial(_log_infer, log, infer) + finally: + log.close() diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index a051fded6..effe4f14e 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -1,5 +1,7 @@ import collections.abc as cabc import time +from contextlib import contextmanager +from functools import partial from typing import Any import numpy as np @@ -10,13 +12,10 @@ from positronic.utils import flatten_dict from positronic.utils.serialization import encode_jpeg -from .base import Answer, Layer, Policy, Runtime, Session +from .base import INFER, Layer, Policy, Session from .recording import Recorder from .spec import from_spec -# The name the wire round trip is served under. A policy whose sessions are ``RemoteSession``s declares it. -INFER = 'infer' - def _prepare_value(value: Any) -> Any: # Codecs nest images inside dicts and lists (e.g. GR00T), so recurse to reach every image array. @@ -37,12 +36,12 @@ def _prepare_obs(obs: cabc.Mapping[str, Any], compress_images: bool) -> dict[str def round_trip( - ws_session: InferenceSession, obs: cabc.Mapping[str, Any], compress_images: bool + ws_session: InferenceSession, compress_images: bool, obs: cabc.Mapping[str, Any] ) -> list[dict[str, Any]] | dict[str, Any]: """One inference over the wire, timed as the ``policy.infer`` span. - The observation is prepared here rather than in the session, because a JPEG encode of an HD frame - stack must not run on the thread that calls the session. The span starts after it, because that + The observation is prepared here rather than in the caller, because a JPEG encode of an HD frame + stack must not run on the thread that drives the control loop. The span starts after it, because that encode is not inference. """ prepared = _prepare_obs(obs, compress_images) @@ -53,60 +52,6 @@ def round_trip( telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns()) -class RemoteSession(Session): - """Per-episode session that forwards observations to a remote inference server. - - One round trip is in flight at a time. The call that starts it answers ``None``, and so does every - call until the round trip comes back. The call that finds it answered returns its trajectory, or drops - that trajectory after a ``cancel``. - - ``compress_images`` comes from what the server declared (see ``RemoteMarker``). - """ - - def __init__(self, ws_session: InferenceSession, rt: Runtime, compress_images: bool = False): - self._session = ws_session - self._rt = rt - self._compress_images = compress_images - self._answer: Answer | None = None - self._cancelled = False - - def __call__(self, obs: cabc.Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | None: - """The trajectory of a round trip that has come back, and ``None`` while one is in flight. - - A server answer of one action becomes a 1-element list, which is the form ``Session.__call__`` - returns. - """ - if self._answer is None: - self._answer = self._rt.fns[INFER](self._session, obs, self._compress_images) - return None - if not self._answer.done(): - return None - answer, cancelled = self._answer, self._cancelled - # The answer and the flag are cleared before the read, because ``result`` raises what the round - # trip raised. A cancel then ends with the answer it was made against, and never drops the next - # chunk. - self._answer, self._cancelled = None, False - result = answer.result() - if cancelled: - return None - return [result] if isinstance(result, dict) else result - - def cancel(self): - # The cancel says the world the chunk applies to has gone. The session still reads the round trip - # for its failure, and drops the chunk that comes with it. - self._cancelled = self._answer is not None - - @property - def meta(self) -> dict[str, Any]: - return flatten_dict({keys.TYPE: 'remote', keys.SERVER: self._session.metadata}) - - def close(self): - assert self._answer is None or self._answer.done(), ( - 'close the runtime serving this session first: the round trip in flight uses the websocket that this closes' - ) - self._session.close() - - class _Endpoint(Policy): """The wire connection to one inference server: sessions forward observations under the border's settings. @@ -127,16 +72,15 @@ def server_meta(self) -> dict[str, Any]: ws_session.close() return self._server_meta - def new_session(self, context=None, rt=None) -> RemoteSession: - if rt is None: - raise ValueError('A remote session runs its inference on a runtime: pass rt to new_session.') + @contextmanager + def episode(self, context=None): + """One connection to the server, open for as long as the episode runs.""" compress = bool(self.server_meta().get(keys.COMPRESS_IMAGES)) ws_session = self._client.new_session() - return RemoteSession(ws_session, rt, compress_images=compress) - - @property - def functions(self) -> cabc.Mapping[str, cabc.Callable[..., Any]]: - return {INFER: round_trip} + try: + yield {INFER: partial(round_trip, ws_session, compress)} + finally: + ws_session.close() @property def meta(self) -> dict[str, Any]: @@ -197,12 +141,11 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, rt=None) -> Session: - return self._policy().new_session(context, rt) + def episode(self, context=None): + return self._policy().episode(context) - @property - def functions(self) -> cabc.Mapping[str, cabc.Callable[..., Any]]: - return self._policy().functions + def new_session(self, rt) -> Session: + return self._policy().new_session(rt) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/spec.py b/positronic/policy/spec.py index 99dcfa3e5..d45a249c6 100644 --- a/positronic/policy/spec.py +++ b/positronic/policy/spec.py @@ -3,7 +3,7 @@ A policy pipeline is one layer chain with a ``remote`` marker naming the client/server border, closed by a ``ModelSource`` terminal:: - pipeline = TemporalStack(...) | ChunkedSchedule() | remote | codec | source + pipeline = TemporalStack(...) | ChunkPlayer() | remote | codec | source Everything left of the marker is the *local* half — the stack the rig runs in front of the connection; everything right of it is the *remote* half — what the inference server runs around @@ -40,7 +40,7 @@ FlipGrip, RestrictImageSize, ) -from positronic.policy.layers import ChunkedSchedule, StopOnFault, TemporalStack +from positronic.policy.layers import ChunkPlayer, StopOnFault, TemporalStack from positronic.policy.observation import ObservationCodec @@ -53,7 +53,7 @@ class RemoteMarker(Layer): ``remote`` is the plain border; call it to describe the wire:: - ChunkedSchedule() | remote(compress_images=True) | codec | source + ChunkPlayer() | remote(compress_images=True) | codec | source """ def __init__(self, compress_images: bool = False): @@ -140,7 +140,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) WIRE_LAYERS: dict[str, type[Layer]] = { layer.WIRE_NAME: layer for layer in ( - ChunkedSchedule, + ChunkPlayer, StopOnFault, TemporalStack, ActionTimestamp, diff --git a/positronic/policy/tests/conftest.py b/positronic/policy/tests/conftest.py new file mode 100644 index 000000000..76f620d6d --- /dev/null +++ b/positronic/policy/tests/conftest.py @@ -0,0 +1,52 @@ +"""Test doubles for the policy API: an answer that is already answered, and a runtime that answers inline.""" + +from collections.abc import Callable, Generator, Mapping +from contextlib import ExitStack +from functools import partial +from typing import Any + +import pytest + +from positronic.policy.base import Answer, Fn, Policy, Runtime, Session + + +class Done(Answer): + """The answer to a call whose work ran inside it.""" + + def __init__(self, value: Any): + self._value = value + + def done(self) -> bool: + return True + + def result(self) -> Any: + return self._value + + +class InlineRuntime(Runtime): + """Runs each call on the calling thread, so the answer it gives is already answered.""" + + def __init__(self, functions: Mapping[str, Callable[..., Any]]): + self._fns: Mapping[str, Fn] = {name: partial(self._call, fn) for name, fn in functions.items()} + + @staticmethod + def _call(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Answer: + return Done(fn(*args, **kwargs)) + + @property + def fns(self) -> Mapping[str, Fn]: + return self._fns + + +@pytest.fixture +def open_session() -> Generator[Callable[[Policy], Session], None, None]: + """Opens a policy's session over an inline runtime; every one it opened is closed at teardown.""" + closing = ExitStack() + + def make(policy: Policy) -> Session: + session = policy.new_session(InlineRuntime(closing.enter_context(policy.episode()))) + closing.callback(session.close) + return session + + yield make + closing.close() diff --git a/positronic/policy/tests/test_executor.py b/positronic/policy/tests/test_executor.py index 7e5ff780b..1530eb555 100644 --- a/positronic/policy/tests/test_executor.py +++ b/positronic/policy/tests/test_executor.py @@ -10,8 +10,8 @@ import pytest -from positronic.policy.base import Answer, DelegatingSession, Layer, NotAnswered, Policy, Session -from positronic.policy.executor import Executor, blocking +from positronic.policy.base import Answer, NotAnswered +from positronic.policy.executor import Executor # How long a test waits for the worker threads before calling the call lost. TIMEOUT_SEC = 5.0 @@ -234,142 +234,6 @@ def test_close_stays_quiet_about_a_call_that_answered(serve, caplog): assert caplog.text == '' -class _PlainPolicy(Policy): - """Serves nothing: its session answers inside the call that asked.""" - - class _Session(Session): - def __init__(self): - self.calls = 0 - - def __call__(self, obs, time_ns): - self.calls += 1 - return [{'action': obs}] - - def __init__(self): - self.session = _PlainPolicy._Session() - - def new_session(self, context=None, rt=None) -> Session: - return self.session - - -_ECHO = 'echo' - - -class _EchoPolicy(Policy): - """Serves ``echo``, and makes sessions that take ``rounds`` calls of it to answer.""" - - class _Session(Session): - def __init__(self, rt, rounds: int): - self._rt = rt - self._answer = None - self._left = rounds - self.calls = 0 - - def __call__(self, obs, time_ns): - self.calls += 1 - result = None - if self._answer is not None: - result, self._answer = self._answer.result(), None - if self._left > 0: - self._left -= 1 - self._answer = self._rt.fns[_ECHO](obs) - return None - return result - - def __init__(self, rounds: int): - self._rounds = rounds - self.session: _EchoPolicy._Session - - def new_session(self, context=None, rt=None) -> Session: - assert rt is not None - self.session = _EchoPolicy._Session(rt, self._rounds) - return self.session - - @property - def functions(self): - return {_ECHO: lambda obs: obs} - - -class _CountingLayer(Layer): - """Counts the calls that reach the session it wraps.""" - - def __init__(self): - self.calls = 0 - - class _Session(DelegatingSession): - def __init__(self, inner: Session, layer: '_CountingLayer'): - super().__init__(inner) - self._layer = layer - - def __call__(self, obs, time_ns): - self._layer.calls += 1 - return self._inner(obs, time_ns) - - def make_session(self, inner): - return self._Session(inner, self) - - -@pytest.fixture -def opened(): - """Opens the sessions a test asks for, and closes every one at teardown.""" - sessions = [] - - def make(policy: Policy) -> Session: - sessions.append(policy.new_session()) - return sessions[-1] - - yield make - for session in sessions: - session.close() - - -class TestBlocking: - """A policy whose sessions answer in the call that asked.""" - - def test_a_session_that_answers_in_its_own_call_is_called_one_time(self, opened): - policy = _PlainPolicy() - - assert opened(blocking(policy))({'x': 1}, 0.0) == [{'action': {'x': 1}}] - assert policy.session.calls == 1 - - @pytest.mark.parametrize(('rounds', 'calls'), [(1, 2), (2, 3)]) - def test_a_session_is_called_again_for_every_function_it_starts(self, opened, rounds, calls): - policy = _EchoPolicy(rounds) - - assert opened(blocking(policy))({'x': 1}, 0.0) == {'x': 1} - assert policy.session.calls == calls - - def test_a_session_that_starts_nothing_and_answers_none_is_called_one_time(self, opened): - policy = _EchoPolicy(rounds=0) - - assert opened(blocking(policy))({'x': 1}, 0.0) is None - assert policy.session.calls == 1 - - def test_a_layer_above_it_is_called_one_time_for_one_answer(self, opened): - """A layer above ``blocking`` is called once for one answer. That is why ``blocking`` wraps the - policy and not the chain: a layer that encodes the observation, or records it, would otherwise do - that work once per call the answer took.""" - layer, policy = _CountingLayer(), _EchoPolicy(rounds=2) - - assert opened(layer.wrap(blocking(policy)))({'x': 1}, 0.0) == {'x': 1} - assert (layer.calls, policy.session.calls) == (1, 3) - - def test_it_serves_its_functions_itself(self): - """A blocking policy runs its own functions, so nothing above it builds a runtime for them.""" - assert blocking(_EchoPolicy(rounds=1)).functions == {} - - def test_closing_the_session_closes_the_runtime_it_made(self): - """The session owns the runtime it was made with, and closing the session is the only way to - close it.""" - policy = _EchoPolicy(rounds=1) - session = blocking(policy).new_session() - session.close() - - # The session's own runtime is closed, so the function it would start is gone. - with pytest.raises(RuntimeError): - policy.session({'x': 1}, 0) - - class _Weights: """Stands in for what a function is declared with: model weights, a socket.""" diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index 1c4a36cfb..e47cae0c1 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -25,6 +25,7 @@ import gzip import json import os +from contextlib import contextmanager from functools import partial from pathlib import Path @@ -41,10 +42,10 @@ from positronic.drivers.roboarm.tests.fakes import make_robot_state from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation, Task from positronic.geom import Rotation, Transform3D -from positronic.policy.base import DelegatingPolicy, DelegatingSession, Policy, Session +from positronic.policy.base import INFER, Answer, DelegatingPolicy, Fn, Policy, Runtime from positronic.policy.codec import ActionTiming from positronic.policy.harness import Harness -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.tests.testing_coutils import ManualDriver, drive_scheduler GOLDEN_FILE = Path(__file__).parent / 'golden_pipeline.json.gz' @@ -64,8 +65,18 @@ CAPTURED_SIGNALS = (keys.EE_POSE, keys.JOINTS, keys.GRIP) -class _ScriptedSession(Session): - def __call__(self, obs, time_ns): +class ScriptedProportionalPolicy(Policy): + """Pure proportional controller toward ``TARGET_POS``. + + Reads ``robot_state.ee_pose`` only; returns a 10-action chunk. No RNG, no + clock, no images. The codec stamps and truncates; ``ChunkPlayer`` anchors the chunk and plays it. + """ + + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} + + def _infer(self, obs): current = np.asarray(obs[keys.EE_POSE][:3], dtype=np.float32) delta = TARGET_POS - current chunk = [] @@ -76,52 +87,46 @@ def __call__(self, obs, time_ns): return chunk -class ScriptedProportionalPolicy(Policy): - """Pure proportional controller toward ``TARGET_POS``. +class _SimulatedLatency(DelegatingPolicy): + """A fixed inference latency in world time: every call answers ``latency_sec`` after it was made, + whatever the machine took. It holds back the runtime the session runs on, so a chunk that is ready + early still reaches the player at the instant the world says.""" - Reads ``robot_state.ee_pose`` only; returns a 10-action chunk. No RNG, no - clock, no images. Codec stamps/truncates; ``ChunkedSchedule`` anchors and the harness plays. - """ + def __init__(self, inner: Policy, latency_sec: float, clock: pimm.Clock): + super().__init__(inner) + self._latency_ns = round(latency_sec * 1e9) + self._clock = clock - def new_session(self, context=None, rt=None): - return _ScriptedSession() + class _Held(Answer): + """The answer under it, held back until the world reaches ``release_at_ns``.""" + def __init__(self, inner: Answer, release_at_ns: int, clock: pimm.Clock): + self._inner = inner + self._release_at_ns = release_at_ns + self._clock = clock -class _SimulatedLatency(DelegatingPolicy): - """A fixed inference latency in world time: every chunk is stamped for, and reaches the harness at, - ``latency_sec`` after the call that produced it, whatever the machine took. Sits outermost, so nothing - below sees an observation while a chunk is held. A skip or a stop passes through at once.""" + def done(self) -> bool: + return self._clock.now_ns() >= self._release_at_ns and self._inner.done() - def __init__(self, inner: Policy, latency_sec: float): - super().__init__(inner) - self._latency_ns = round(latency_sec * 1e9) + def result(self): + return self._inner.result() - class _Session(DelegatingSession): - def __init__(self, inner: Session, latency_ns: int): - super().__init__(inner) + class _Runtime(Runtime): + def __init__(self, inner: Runtime, latency_ns: int, clock: pimm.Clock): + self._fns = {name: partial(self._held, fn) for name, fn in inner.fns.items()} self._latency_ns = latency_ns - self._held: list | None = None - self._release_at_ns = 0 - - def __call__(self, obs, time_ns): - if self._held is not None: - if time_ns < self._release_at_ns: - return None - held, self._held = self._held, None - return held - # The chunk is held for ``latency_ns``, so the sessions below stamp it from the release instant. - result = self._inner(obs, time_ns + self._latency_ns) - if not result: - return result - self._held, self._release_at_ns = result, time_ns + self._latency_ns - return None - - def cancel(self): - self._held = None - super().cancel() - - def new_session(self, context=None, rt=None): - return _SimulatedLatency._Session(self._inner.new_session(context, rt), self._latency_ns) + self._clock = clock + + def _held(self, fn: Fn, *args, **kwargs) -> Answer: + release_at_ns = self._clock.now_ns() + self._latency_ns + return _SimulatedLatency._Held(fn(*args, **kwargs), release_at_ns, self._clock) + + @property + def fns(self): + return self._fns + + def new_session(self, rt): + return self._inner.new_session(_SimulatedLatency._Runtime(rt, self._latency_ns, self._clock)) class FakeRobot(pimm.ControlSystem): @@ -208,7 +213,10 @@ def _run_pipeline(tmp_path: Path) -> dict: simulated=True, ) harness = Harness( - _SimulatedLatency((StopOnFault() | ChunkedSchedule()).wrap(policy), INFERENCE_LATENCY_S), embodiment + # The latency wraps the runtime the whole stack runs on, so it goes outside the player that + # reads that runtime. + _SimulatedLatency((StopOnFault() | ChunkPlayer()).wrap(policy), INFERENCE_LATENCY_S, world.clock), + embodiment, ) ds_agent = wire.wire_embodiment(world, harness, embodiment, ds_writer, TimeMode.MESSAGE) world.connect(harness.ds_command, ds_agent.command) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 096558368..679fa9340 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,5 +1,6 @@ import threading import time +from collections import deque from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -21,14 +22,14 @@ from positronic.eval import Command, Embodiment, Observation, Task from positronic.geom import Rotation, Transform3D from positronic.offboard.client import InferenceSession -from positronic.policy.base import DelegatingSession, Layer, Policy, Session +from positronic.policy.base import INFER, Answer, DelegatingPolicy, DelegatingSession, Fn, Layer, Policy, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import POLL_PERIOD_SEC, Harness, _EpisodeInference -from positronic.policy.layers import ChunkedSchedule, StopOnFault -from positronic.policy.remote import INFER, RemoteSession, round_trip +from positronic.policy.harness import MAX_ROUND_SEC, MIN_ROUND_SEC, WAIT_PERIOD_SEC, Harness, _EpisodeInference +from positronic.policy.layers import ChunkPlayer, StopOnFault +from positronic.policy.remote import round_trip from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler -POLL_PERIOD_NS = round(POLL_PERIOD_SEC * 1e9) +WAIT_PERIOD_NS = round(WAIT_PERIOD_SEC * 1e9) @contextmanager @@ -79,15 +80,6 @@ def make_embodiment( ) -class _SpySession(Session): - def __init__(self, policy): - self._policy = policy - - def __call__(self, obs, time_ns): - self._policy.last_obs = obs - return [{keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.target_grip, 'timestamp': 0.0}] - - class SpyPolicy(Policy): def __init__(self, command: roboarm.command.CommandType | None = None, target_grip: float = 0.33) -> None: if command is None: @@ -99,25 +91,15 @@ def __init__(self, command: roboarm.command.CommandType | None = None, target_gr self.reset_calls: int = 0 self.last_reset_context = None - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): self.reset_calls += 1 self.last_reset_context = context - return _SpySession(self) - - -class _StubSession(Session): - def __init__(self, policy): - self._policy = policy - self._meta = dict(policy._meta) - - def __call__(self, obs, time_ns): - self._policy.last_obs = obs - self._policy.observations.append(obs) - return [{keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.target_grip, 'timestamp': 0.0}] + yield {INFER: self._infer} - @property - def meta(self): - return self._meta + def _infer(self, obs): + self.last_obs = obs + return [{keys.ROBOT_COMMAND: self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] class StubPolicy(Policy): @@ -144,27 +126,16 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, rt=None) -> Session: + @contextmanager + def episode(self, context=None): self.reset_calls += 1 self.last_reset_context = context - return _StubSession(self) + yield {INFER: self._infer} - -class _ChunkSession(Session): - def __init__(self, policy): - self._policy = policy - - def __call__(self, obs, time_ns): - self._policy.counter += 1 - dt = 0.005 - return [ - { - keys.ROBOT_COMMAND: self._policy.command, - 'target_grip': self._policy.counter * 100.0 + i, - 'timestamp': i * dt, - } - for i in range(10) - ] + def _infer(self, obs): + self.last_obs = obs + self.observations.append(obs) + return [{keys.ROBOT_COMMAND: self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] class ChunkPolicy(StubPolicy): @@ -174,16 +145,18 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 - def new_session(self, context=None, rt=None): - self.reset_calls += 1 - self.last_reset_context = context - return _ChunkSession(self) + def _infer(self, obs): + self.counter += 1 + dt = 0.005 + return [ + {keys.ROBOT_COMMAND: self.command, 'target_grip': self.counter * 100.0 + i, 'timestamp': i * dt} + for i in range(10) + ] class _FakeInferenceSession(InferenceSession): - """A stub ``InferenceSession`` returning a canned action after ``wall_sec`` of real time, so a - ``RemoteSession`` over it round-trips ``RemoteSession.__call__`` — the real inference boundary that records - the ``policy.infer`` span.""" + """A stub ``InferenceSession`` returning a canned action after ``wall_sec`` of real time, so the wire + round trip over it is the real inference boundary that records the ``policy.infer`` span.""" def __init__(self, action: list[dict[str, Any]], wall_sec: float = 0.0) -> None: self._action = action @@ -202,9 +175,8 @@ def close(self) -> None: class ServedPolicy(Policy): - """A policy whose model runs in a served function: a real ``RemoteSession`` over the ``InferenceSession`` - it is given, so its inference round-trips ``RemoteSession.__call__`` and records the ``policy.infer`` - span independent of any layer. + """A policy whose model runs in a served function: a real wire round trip over the ``InferenceSession`` + it is given, so its inference records the ``policy.infer`` span independent of any layer. A model that costs wall time, hangs or raises belongs in that session's ``infer``. """ @@ -212,13 +184,16 @@ class ServedPolicy(Policy): def __init__(self, session: InferenceSession) -> None: self._session = session - def new_session(self, context=None, rt=None) -> RemoteSession: - assert rt is not None, 'the harness supplies the runtime' - return RemoteSession(self._session, rt) + @contextmanager + def episode(self, context=None): + try: + yield {INFER: partial(round_trip, self._session, False)} + finally: + self._session.close() - @property - def functions(self): - return {INFER: round_trip} + +def _boom(obs): + raise RuntimeError('inference boom') def slow_chunk(span_sec: float = 0.2, steps: int = 10) -> list[dict[str, Any]]: @@ -295,6 +270,9 @@ def run(self, should_stop, clock): yield pimm.Sleep(0.001) +_NEVER_STOPS: pimm.SignalReceiver[bool] = pimm.NoOpReceiver() + + def _pair_all(world, harness): """Pair all harness signals and return a dict of test handles.""" ds_recorder = RecordingEmitter() @@ -356,7 +334,7 @@ def _emitted_grips(recorder): def test_harness_emits_cartesian_move(world): pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) policy = SpyPolicy(command=CartesianPosition(pose=pose), target_grip=0.33) - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) @@ -417,7 +395,7 @@ def test_harness_emits_cartesian_move(world): def test_harness_passes_descriptor_to_policy(world): """The embodiment descriptor reaches the policy on every call (stateless policy).""" policy = SpyPolicy() - harness = Harness(policy, make_embodiment(descriptor='mujoco.franka')) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(descriptor='mujoco.franka')) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands['target_grip']._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) @@ -447,7 +425,7 @@ def test_robot_model_stays_out_of_the_observation(world): policy = SpyPolicy() model = bundled_franka_model() statics = {keys.URDF: model[keys.URDF], keys.CONTROL_FRAME: model[keys.CONTROL_FRAME]} - harness = Harness(policy, make_embodiment(static_meta=statics)) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(static_meta=statics)) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands['target_grip']._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) @@ -474,7 +452,7 @@ def test_robot_model_stays_out_of_the_observation(world): @pytest.mark.timeout(3.0) def _run_with_model(world, model, static_meta=None): """Drive one episode with ``model`` published on ``robot_meta_in``, or baked into embodiment statics.""" - harness = Harness(SpyPolicy(), make_embodiment(static_meta=static_meta)) + harness = Harness(ChunkPlayer().wrap(SpyPolicy()), make_embodiment(static_meta=static_meta)) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands['target_grip']._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) @@ -516,7 +494,7 @@ def test_rejects_a_control_frame_a_late_model_declares(world): def test_harness_waits_for_complete_inputs(world): pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) policy = SpyPolicy(command=CartesianPosition(pose=pose), target_grip=0.33) - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) @@ -563,7 +541,7 @@ def assert_no_inference(): @pytest.mark.timeout(3.0) def test_episode_meta_stamped_at_finalize(world): policy = StubPolicy(meta={'type': 'stub', 'checkpoint': 'v1'}) - harness = Harness(policy, make_embodiment(), static_meta={keys.JOINT_SIGNALS: [keys.JOINTS]}) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(), static_meta={keys.JOINT_SIGNALS: [keys.JOINTS]}) p = _pair_all(world, harness) driver = ManualDriver([ @@ -592,26 +570,20 @@ def test_episode_meta_includes_policy_static_meta(world): """Static fields exposed only via ``Policy.meta`` (empty ``Session.meta``) must still reach episode metadata once the policy is wrapped.""" - class _StaticMetaSession(Session): - def __init__(self, command): - self._command = command - - def __call__(self, obs, time_ns): - return [{keys.ROBOT_COMMAND: self._command, 'target_grip': 0.0, 'timestamp': 0.0}] - class _StaticMetaPolicy(Policy): def __init__(self): pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) self._command = CartesianPosition(pose=pose) - def new_session(self, context=None, rt=None): - return _StaticMetaSession(self._command) # Session.meta defaults to {} + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: [{keys.ROBOT_COMMAND: self._command, 'target_grip': 0.0, 'timestamp': 0.0}]} @property def meta(self): return {'checkpoint': 'v1', 'type': 'static'} - harness = Harness(_StaticMetaPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(_StaticMetaPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([ @@ -633,7 +605,7 @@ def meta(self): @pytest.mark.timeout(3.0) def test_finish_emits_ds_stop_with_data(world): policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) driver = ManualDriver([ @@ -654,7 +626,7 @@ def test_finish_emits_ds_stop_with_data(world): @pytest.mark.timeout(3.0) def test_the_call_is_answered_with_the_terminal_the_episode_ended_on(world): """A task with no timeout ends on ``done`` alone, and its caller gets what it ended on.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) scheduler = world.start([harness]) @@ -671,7 +643,7 @@ def test_the_call_is_answered_with_the_terminal_the_episode_ended_on(world): @pytest.mark.timeout(3.0) def test_the_world_stopping_under_a_live_episode_fails_the_call(world): """The caller hears that its episode will never answer, rather than holding a handle that never completes.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) scheduler = world.start([harness, ManualDriver([(None, 0.02)])]) @@ -689,23 +661,11 @@ def test_an_uncharged_wait_ends_when_the_world_comes_down(world): never_answers = threading.Event() class _HangingPolicy(Policy): - class _Session(Session): - def __init__(self, rt): - self._rt = rt - - def __call__(self, obs, time_ns): - self._rt.fns[INFER]() - return None - - def new_session(self, context=None, rt=None): - assert rt is not None - return _HangingPolicy._Session(rt) + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: never_answers.wait()} - @property - def functions(self): - return {INFER: never_answers.wait} - - inference = _EpisodeInference(_HangingPolicy(), {}, charges_wall_time=False, clock=world.clock) + inference = _EpisodeInference(ChunkPlayer().wrap(_HangingPolicy()), {}, charges_wall_time=False, clock=world.clock) try: inference({}) # starts the function, which never answers world.request_stop() @@ -719,15 +679,12 @@ def test_a_session_that_raises_fails_the_call_that_asked_for_the_episode(world): """The session is called on the loop thread, so its failure reaches whoever asked for the episode rather than the log.""" - class _RaisingSession(Session): - def __call__(self, obs, time_ns): - raise RuntimeError('inference boom') - class _RaisingPolicy(Policy): - def new_session(self, context=None, rt=None): - return _RaisingSession() + @contextmanager + def episode(self, context=None): + yield {INFER: _boom} - harness = Harness(_RaisingPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(_RaisingPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([ @@ -748,7 +705,7 @@ def new_session(self, context=None, rt=None): def test_trial_ends_at_its_timeout(world): """Nothing ever lands on ``done``, yet the trial still ends at ``task.timeout_sec``: terminated=False.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) scheduler = world.start([harness]) @@ -766,7 +723,7 @@ def test_trial_budget_starts_when_the_rig_is_ready(world): """The 0.05 budget is measured from the end of the prepare, not from the ask: the 0.2 the scene takes to draw is not the trial's to spend.""" scene = _Scene(lambda _: None, draw_s=0.2) - harness = Harness(StubPolicy(), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) @@ -783,7 +740,7 @@ def test_trial_budget_starts_when_the_rig_is_ready(world): def test_trial_stop_signal_terminates(world): """Delivering the privileged ``done`` ends a trial early: terminated=True, payload recorded.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) scheduler = world.start([harness]) @@ -809,7 +766,7 @@ def test_stale_done_does_not_terminate_next_trial(world): latched value is ignored — no producer ``reset`` clears it here (``reset`` is ``None``, as on a real embodiment). A falsy payload never terminates; trial 1 runs until its own fresh terminal lands.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) task = Task(instruction_source='t', timeout_sec=100.0) @@ -880,7 +837,7 @@ def test_the_policy_opens_on_the_frame_the_reset_published(world): simulated=True, ) policy = StubPolicy() - harness = Harness(policy, embodiment) + harness = Harness(ChunkPlayer().wrap(policy), embodiment) perform_task = world.pair(harness.perform_task) wire.wire_embodiment(world, harness, embodiment, None) @@ -924,7 +881,7 @@ def run(self, should_stop, clock): ) # Termination is independent of the policy layers; the minimal embodiment has no # ``robot_state``, so run the stub policy bare. - harness = Harness(StubPolicy(), embodiment) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), embodiment) ds_recorder = RecordingEmitter() harness.ds_command._bind(ds_recorder) perform_task = world.pair(harness.perform_task) @@ -945,7 +902,7 @@ def test_done_after_deadline_is_a_timeout(world): """The deadline is hard: a ``done`` delivered past it records as a timeout — ``eval.terminated`` False, payload dropped — not a late stop-signal success.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -971,7 +928,7 @@ def test_a_handler_the_trial_does_not_name_is_left_alone(world): drawn, moved = [], [] scene, arm = _Scene(drawn.append), _Scene(moved.append) handlers = {keys.SCENE: scene.env_reset, keys.ARM: arm.env_reset} - harness = Harness(StubPolicy(), make_embodiment(prepare_handlers=handlers)) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment(prepare_handlers=handlers)) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) wire_call(world, harness.prepare[keys.ARM], arm.env_reset) @@ -993,7 +950,7 @@ def test_every_rig_is_put_back_where_the_trial_placed_it(world, simulated): placed = [] arm = _Scene(placed.append) handlers = {keys.ARM: arm.env_reset} - harness = Harness(StubPolicy(), make_embodiment(simulated=simulated, prepare_handlers=handlers)) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment(simulated=simulated, prepare_handlers=handlers)) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.ARM], arm.env_reset) @@ -1029,7 +986,7 @@ def test_a_trial_does_not_end_until_the_rig_is_back_where_it_started(world): """The terminal waits on the return move. Handed back sooner, the next trial's scene draw goes ahead of a move still travelling and rebuilds the model under it, leaving nothing but its timeout to end it.""" arm = _PlacesOnce() - harness = Harness(StubPolicy(), make_embodiment(prepare_handlers={keys.ARM: arm.env_reset})) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment(prepare_handlers={keys.ARM: arm.env_reset})) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.ARM], arm.env_reset) @@ -1046,7 +1003,7 @@ def test_the_deadline_is_published_once_the_rig_is_ready(world): """An idle harness publishes nothing: ``deadline_ns`` states the instant the harness will stop at, and between episodes there is none to state. The first one goes out when the episode's prepare has answered — this embodiment readies nothing, so that is the round the task is taken.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) driver = ManualDriver([(None, 100.0)]) # outlives the pumping, so the world stays up between phases @@ -1059,8 +1016,8 @@ def test_the_deadline_is_published_once_the_rig_is_ready(world): drive_scheduler(scheduler, steps=20) # The world is already past zero here, so the published instant and a bare ``timeout_sec`` are # different numbers, which is what this pins. - assert _deadlines(p) == [pytest.approx(asked_at_ns + 5e9, abs=2 * POLL_PERIOD_NS)] - assert asked_at_ns > 2 * POLL_PERIOD_NS, 'the two would be indistinguishable at a clock still near zero' + assert _deadlines(p) == [pytest.approx(asked_at_ns + 5e9, abs=2 * WAIT_PERIOD_NS)] + assert asked_at_ns > 2 * WAIT_PERIOD_NS, 'the two would be indistinguishable at a clock still near zero' @pytest.mark.timeout(3.0) @@ -1086,7 +1043,7 @@ def run(self, should_stop, clock): yield pimm.Sleep(0.001) scene = _SlowScene() - harness = Harness(StubPolicy(), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) @@ -1104,7 +1061,7 @@ def run(self, should_stop, clock): def test_the_deadline_clears_when_the_episode_ends(world): """``None`` at the close is what tells a display the countdown is over, and it is published only there: cleared mid-episode it would stop a countdown the harness is still enforcing.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1125,7 +1082,7 @@ def test_the_deadline_clears_when_the_episode_ends(world): def test_an_episode_with_no_timeout_publishes_no_deadline(world): """A task with no ``timeout_sec`` publishes ``None``, which corrects a reader still holding the previous episode's deadline. Silence would leave it counting down against one that has lapsed.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1143,7 +1100,7 @@ def test_an_episode_with_no_timeout_publishes_no_deadline(world): def test_a_world_stopping_mid_episode_withdraws_the_deadline(world): """A run that ends with an episode still live withdraws its deadline like any other close: a receiver latches what it last got, and there is no later episode to correct it.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1162,15 +1119,12 @@ def test_a_world_stopping_mid_episode_withdraws_the_deadline(world): def test_an_episode_abandoned_by_a_raise_withdraws_the_deadline(world): """An episode a raise abandons withdraws its deadline like any other close.""" - class _BoomSession(Session): - def __call__(self, obs): - raise RuntimeError('inference boom') - class _BoomPolicy(Policy): - def new_session(self, context=None, now=None, rt=None): - return _BoomSession() + @contextmanager + def episode(self, context=None): + yield {INFER: _boom} - harness = Harness(_BoomPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(_BoomPolicy()), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1192,7 +1146,7 @@ def test_a_trial_asking_to_ready_what_the_rig_has_not_got_fails_loudly(world): trial would open on a rig nothing readied.""" scene = _Scene(lambda _params: None) embodiment = make_embodiment(descriptor='yam', prepare_handlers={keys.SCENE: scene.env_reset}) - harness = Harness(StubPolicy(), embodiment) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), embodiment) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) @@ -1221,7 +1175,7 @@ def reset(params): p['meta_em'].emit({}) # the producer publishes fresh scene meta, recorded into the episode at finalize scene = _Scene(reset) - harness = Harness(policy, make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) @@ -1252,7 +1206,7 @@ def test_timeout_during_inference_drops_the_chunk(world): point.""" harness = Harness( # the function runs well past the deadline - ChunkedSchedule().wrap(RemoteStubPolicy(wall_sec=0.3, chunk=slow_chunk())), + ChunkPlayer().wrap(RemoteStubPolicy(wall_sec=0.3, chunk=slow_chunk())), make_embodiment(simulated=True), ) cmd_recorder = RecordingEmitter() @@ -1287,7 +1241,7 @@ def test_timeout_during_inference_drops_the_chunk(world): def test_a_terminal_landing_while_idle_does_not_end_the_next_episode(world): """A finish pressed with nothing running belongs to no episode, so the one asked for next runs on.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1309,7 +1263,7 @@ def test_a_terminal_landing_while_idle_does_not_end_the_next_episode(world): def test_a_call_arriving_mid_episode_is_refused(world): """The live episode runs on and the second caller is told why, rather than its ask being dropped.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) scheduler = world.start([harness]) @@ -1345,7 +1299,7 @@ def test_a_producer_reusing_its_buffer_cannot_rewrite_a_pending_observation(worl """A camera renders into the array behind the adapter it re-emits, and a wall-charged trial keeps the loop stepping while the function runs — so the observation handed to that function has to be its own copy.""" watcher = _FrameWatchingSession(wall_sec=0.3) - harness = Harness(ServedPolicy(watcher), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(ServedPolicy(watcher)), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) frame = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) @@ -1389,9 +1343,11 @@ def __init__(self, wall_sec: float): self.events: list[str] = [] super().__init__(_AbandonedCallPolicy._Infer(self.events, wall_sec)) - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): self.events.append('open') - return super().new_session(context, rt) + with super().episode(context) as fns: + yield fns @pytest.mark.timeout(10.0) @@ -1399,7 +1355,7 @@ def test_a_new_episode_waits_out_the_call_the_last_one_abandoned(world): """An in-process policy is one model across episodes, so opening a session must not overtake a function still inside the previous one — ``new_session`` resets the object that function is using.""" policy = _AbandonedCallPolicy(wall_sec=0.4) - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) emit_obs = partial(emit_ready_payload, p['frame_em'], p['robot_em'], p['grip_em'], robot_state) @@ -1423,7 +1379,7 @@ def test_a_new_episode_waits_out_the_call_the_last_one_abandoned(world): @pytest.mark.timeout(3.0) def test_run_calls_policy_reset_with_context(world): policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) driver = ManualDriver([ @@ -1449,7 +1405,7 @@ def reset(_params): scene['task'] = 'resolved-on-reset' # the env reports its task only here drawing = _Scene(reset) - harness = Harness(policy, make_embodiment(prepare_handlers={keys.SCENE: drawing.env_reset})) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(prepare_handlers={keys.SCENE: drawing.env_reset})) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], drawing.env_reset) @@ -1478,7 +1434,7 @@ def test_finish_stops_playing_the_live_chunk(world): devices, so nothing is emitted past the recorder's STOP.""" policy = ChunkPolicy() wrapped = ActionTimestamp(fps=5.0).wrap(policy) # 1.8 s chunk — won't drain before the episode ends - harness = Harness(wrapped, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(wrapped), make_embodiment()) events: list[tuple[str, object]] = [] harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) harness.commands['target_grip']._bind(_LabeledRecorder('target_grip', events)) @@ -1512,15 +1468,12 @@ def test_empty_trajectory_leaves_every_channel_holding(world): """A trajectory with no waypoints schedules nothing on any channel, so every device holds where it already is rather than one channel draining on while another stops.""" - class _EmptyChunkSession(Session): - def __call__(self, obs, time_ns): - return [] - class EmptyChunkPolicy(Policy): - def new_session(self, context=None, rt=None): - return _EmptyChunkSession() + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: []} - harness = Harness(EmptyChunkPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(EmptyChunkPolicy()), make_embodiment()) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) @@ -1548,7 +1501,7 @@ def new_session(self, context=None, rt=None): @pytest.mark.timeout(3.0) def test_harness_clears_trajectory_on_finish(world): policy = ChunkPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1575,7 +1528,7 @@ def test_harness_clears_trajectory_on_finish(world): @pytest.mark.timeout(3.0) def test_harness_clears_trajectory_on_run(world): policy = ChunkPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1593,7 +1546,7 @@ def test_harness_clears_trajectory_on_run(world): drive_scheduler(scheduler, steps=1) emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) - drive_scheduler(scheduler, steps=4) + drive_scheduler(scheduler, steps=20) assert _last_grip(p) >= 200.0, 'Expected chunk 2; trajectory clearing on a new episode failed' @@ -1606,7 +1559,7 @@ def test_the_stack_keeps_the_model_away_from_an_unavailable_arm(world, unavailab asked again, on a fresh chunk.""" policy = ChunkPolicy() - harness = Harness((StopOnFault() | ChunkedSchedule()).wrap(policy), make_embodiment()) + harness = Harness((StopOnFault() | ChunkPlayer()).wrap(policy), make_embodiment()) p = _pair_all(world, harness) state_ok = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.AVAILABLE) @@ -1677,7 +1630,7 @@ def test_shutdown_stops_playing_the_live_chunk(world): the devices after the recorder's STOP.""" events: list[tuple[str, object]] = [] wrapped = ActionTimestamp(fps=5.0).wrap(ChunkPolicy()) # 1.8 s chunk — won't drain before shutdown - harness = Harness(wrapped, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(wrapped), make_embodiment()) harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) harness.commands[keys.TARGET_GRIP]._bind(_LabeledRecorder(keys.TARGET_GRIP, events)) harness.ds_command._bind(_LabeledRecorder('ds_command', events)) @@ -1732,7 +1685,7 @@ def test_stop_mid_episode_keeps_episode_open_for_recorder_flush(world, tmp_path) shutdown-flush ``record.io`` span parents to the episode, not the pass. Driven straight through the generator protocol: the yield after the queued STOP is the recorder's flush slot.""" policy = StubPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) ds_recorder = RecordingEmitter() harness.ds_command._bind(ds_recorder) # Never ends within the drive: the stop, not the deadline, is what winds this episode down. @@ -1775,8 +1728,8 @@ def test_timing_spans_recorded_with_taxonomy(world, tmp_path): episode parents to the pass, and reset + policy.infer parent to the episode, with the episode carrying its index, step count, and virtual duration. Read back from the file so the OTLP encoding is exercised. The ``policy.infer`` span is recorded at the remote inference boundary, so the terminal is a ``RemoteStubPolicy`` - (a real ``RemoteSession`` over a fake inference session).""" - policy = ChunkedSchedule().wrap(RemoteStubPolicy()) + (a real wire round trip over a fake inference session).""" + policy = ChunkPlayer().wrap(RemoteStubPolicy()) harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1810,9 +1763,9 @@ def test_timing_spans_recorded_with_taxonomy(world, tmp_path): assert all(r.parent_id == episode.span_id for r in by_name[telemetry_keys.SPAN_POLICY_INFER]) assert episode.attrs[telemetry_keys.ATTR_EPISODE_INDEX] == 0 - # Every answered round trip is one step. The trial can end on a round trip in flight, which has no step. + # Every round the harness calls the session is one step, so the count covers the round trips it made. infers = len(by_name[telemetry_keys.SPAN_POLICY_INFER]) - assert infers - 1 <= episode.attrs[telemetry_keys.ATTR_EPISODE_STEPS] <= infers + assert episode.attrs[telemetry_keys.ATTR_EPISODE_STEPS] >= infers > 0 assert episode.attrs[telemetry_keys.ATTR_EPISODE_VIRTUAL_S] >= 0.0 @@ -1823,7 +1776,7 @@ def test_an_inference_outliving_its_episode_parents_to_it(world, tmp_path): the episode that asked for it rather than to the pass: charging wall time is the mode that measures real inference cost, so that span is the one a reader most wants attributed.""" harness = Harness( - ChunkedSchedule().wrap(RemoteStubPolicy(wall_sec=0.3)), # the call runs well past the deadline + ChunkPlayer().wrap(RemoteStubPolicy(wall_sec=0.3)), # the call runs well past the deadline make_embodiment(simulated=True), ) p = _pair_all(world, harness) @@ -1857,7 +1810,7 @@ def test_failed_pass_seals_open_episode_span(world, tmp_path): policy = StubPolicy() scene = pimm.calls.ControlSystemHandler[Any, None](Passive()) - harness = Harness(policy, make_embodiment(prepare_handlers={keys.SCENE: scene})) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(prepare_handlers={keys.SCENE: scene})) wire_call(world, harness.prepare[keys.SCENE], scene) harness.ds_command._bind(RecordingEmitter()) task = Task( @@ -1899,7 +1852,9 @@ def test_episode_virtual_duration_starts_when_the_rig_is_ready(world, tmp_path): The rollout's virtual duration measures from the end of the prepare, so that stretch stays reset work instead of inflating the real-time factor the report derives from it.""" scene = _Scene(lambda _: None, draw_s=0.2) - harness = Harness(ChunkPolicy(), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset})) + harness = Harness( + ChunkPlayer().wrap(ChunkPolicy()), make_embodiment(prepare_handlers={keys.SCENE: scene.env_reset}) + ) p = _pair_all(world, harness) wire_call(world, harness.prepare[keys.SCENE], scene.env_reset) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1928,62 +1883,119 @@ def test_episode_virtual_duration_starts_when_the_rig_is_ready(world, tmp_path): assert virtual_s < draw_s -def test_unanchored_chunk_is_refused(): - """A stack that never anchored leaves chunk-relative stamps, which read as decades before now.""" - with pytest.raises(ValueError, match='not anchoring'): - Harness._assert_anchored([{'timestamp': 0.0}], now=1.7e9) +@pytest.mark.parametrize( + ('resume_in_sec', 'sleep_sec'), + [(0.002, 0.002), (5.0, MAX_ROUND_SEC), (0.0002, MIN_ROUND_SEC), (-0.001, MIN_ROUND_SEC), (None, WAIT_PERIOD_SEC)], + ids=['asked_for', 'over_the_ceiling', 'under_the_floor', 'already_passed', 'no_answer_yet'], +) +def test_a_real_rig_wakes_at_the_moment_the_session_asked_for(world, resume_in_sec, sleep_sec): + """The harness sleeps to the instant the session named, inside its own floor and ceiling. A round the + live session has not answered waits on the call the harness made.""" + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) + now_ns = world.clock.now_ns() + harness._resume_at_ns = None if resume_in_sec is None else now_ns + int(resume_in_sec * 1e9) + command = harness._pace(world.clock) -def test_doubly_anchored_chunk_is_refused(): - """Two schedulers each add the clock, putting the chunk a lifetime ahead.""" - with pytest.raises(ValueError, match='not anchoring'): - Harness._assert_anchored([{'timestamp': 3.4e9}], now=1.7e9) + assert isinstance(command, pimm.Sleep) + assert command.seconds == pytest.approx(sleep_sec, abs=1e-4) -def test_anchored_chunk_passes(): - """A real chunk spans seconds around now, and a late action sits just behind it.""" - Harness._assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) +def test_a_real_rig_never_sleeps_past_the_deadline(world): + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) + now_ns = world.clock.now_ns() + harness._resume_at_ns = now_ns + 500_000_000 + harness._deadline_ns = now_ns + 20_000_000 + command = harness._pace(world.clock) + + assert isinstance(command, pimm.Sleep) + assert command.seconds == pytest.approx(0.02, abs=1e-4) + + +@pytest.mark.parametrize(('expired', 'emitted'), [(True, False), (False, True)]) +def test_a_command_is_emitted_only_while_the_trial_still_has_budget(world, expired, emitted): + """A trial advertises the instant it stops at. A session call the world passes that instant during + commands nothing, and ``_run`` finishes the trial on the next round.""" + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) + p = _pair_all(world, harness) + scheduler = world.start([harness]) + p['perform_task'](Task(instruction_source='t', timeout_sec=None)) + emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], make_robot_state([0.1] * 3, [0.2] * 7)) + drive_scheduler(scheduler, steps=6) + assert harness._inference is not None, 'the episode never opened' + + # Armed after the round's terminal check, which is the window a slow session call opens. + harness._deadline_ns = world.clock.now_ns() + (-1_000_000_000 if expired else 1_000_000_000) + p['command_rx'].read() # drop what the opening round played + harness._infer(harness._inference, world.clock, _NEVER_STOPS) + + message = p['command_rx'].read() + assert (message is not None and message.updated) is emitted -@pytest.mark.parametrize(('expired', 'scheduled'), [(True, False), (False, True)]) -def test_a_reply_is_scheduled_only_while_the_trial_still_has_budget(world, expired, scheduled): - """A trial advertises the instant it stops at. A chunk answered after the world passed that instant is - dropped instead of placed, and ``_run`` finishes the trial on the next round.""" - harness = Harness(StubPolicy(), make_embodiment()) - now_ns = world.clock.now_ns() - harness._deadline_ns = now_ns - 1_000_000_000 if expired else now_ns + 1_000_000_000 - harness._reschedule(slow_chunk(), world.clock) +def test_a_policy_that_answers_chunks_refuses_to_open(world): + """Nothing turns a chunk into commands without a player, and an episode must not start on one.""" + harness = Harness(StubPolicy(), make_embodiment()) + p = _pair_all(world, harness) + scheduler = world.start([harness]) + call = p['perform_task'](Task(instruction_source='t', timeout_sec=None)) + emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], make_robot_state([0.1] * 3, [0.2] * 7)) - assert bool(harness._schedules[keys.ROBOT_COMMAND]) is scheduled + with pytest.raises(NotImplementedError, match='put a ChunkPlayer above it'): + drive_scheduler(scheduler, steps=20) + assert call.done() class _ReplanEarly(Layer): - """Infers on the first observation and again halfway through the chunk it returned. + """Plays a chunk and asks the model again halfway through it. The re-query-before-exhaustion shape (RTC, temporal ensembling) that the substrate exists for: unlike - ``ChunkedSchedule`` it leaves waypoints to play while a call is in flight. + ``ChunkPlayer`` it starts the next call while waypoints of the chunk it holds are still due. """ - class _Session(DelegatingSession): - def __init__(self, inner: Session): - super().__init__(inner) - self._replan_at: float | None = None + class _Session(Session): + POLL_SEC = 0.001 + + def __init__(self, infer: Fn): + self._infer = infer + self._waypoints: deque[tuple[int, dict[str, Any]]] = deque() + self._replan_at_ns: int | None = None + self._answer: Answer | None = None def __call__(self, obs, time_ns): - t0 = obs[keys.OBS_TIME_NS] / 1e9 - if self._replan_at is not None and t0 < self._replan_at: - return None - result = self._inner(obs, time_ns) - if result is None: # the function it asked has still to answer - return None + commands: dict[str, Any] = {} + while self._waypoints and self._waypoints[0][0] <= time_ns: + commands.update(self._waypoints.popleft()[1]) + if self._replan_at_ns is None or time_ns >= self._replan_at_ns: + self._replan(obs, time_ns) + next_waypoint_ns = self._waypoints[0][0] if self._waypoints else None + due = [at_ns for at_ns in (self._replan_at_ns, next_waypoint_ns) if at_ns is not None and at_ns > time_ns] + return commands, min(due, default=time_ns + int(self.POLL_SEC * 1e9)) + + def _replan(self, obs, time_ns: int) -> None: + if self._answer is None: + self._answer = self._infer(obs) + if not self._answer.done(): # the call it made has still to answer + return + chunk, self._answer = self._answer.result(), None anchor = time_ns / 1e9 - result = [{**action, keys.ACTION_TIMESTAMP: anchor + action[keys.ACTION_TIMESTAMP]} for action in result] - self._replan_at = t0 + (result[-1][keys.ACTION_TIMESTAMP] - t0) / 2 - return result + self._waypoints = deque( + (int((anchor + action[keys.ACTION_TIMESTAMP]) * 1e9), {keys.TARGET_GRIP: action[keys.TARGET_GRIP]}) + for action in chunk + if keys.TARGET_GRIP in action + ) + end_ns = self._waypoints[-1][0] if self._waypoints else time_ns + self._replan_at_ns = time_ns + (end_ns - time_ns) // 2 - def make_session(self, inner: Session): - return _ReplanEarly._Session(inner) + class _Policy(DelegatingPolicy): + def new_session(self, rt): + return _ReplanEarly._Session(rt.fns[INFER]) + + PLAYS_CHUNKS = True + + def wrap(self, policy: Policy) -> Policy: + return _ReplanEarly._Policy(policy) class _TimedRecorder(pimm.SignalEmitter): @@ -2035,7 +2047,7 @@ def test_an_uncharged_call_pauses_the_world(world): """Sim's default charges nothing: the world does not advance while the model runs, so the chunk is anchored at the observation's own instant however long the function really took.""" policy = RemoteStubPolicy(wall_sec=0.05, chunk=slow_chunk()) - played = _run_episode(world, policy, ChunkedSchedule(), charge_inference_time=False) + played = _run_episode(world, policy, ChunkPlayer(), charge_inference_time=False) assert played, 'no command was played' assert played[0][0] < 0.05, f'the world paid for the function: first command at {played[0][0]}s' @@ -2046,7 +2058,7 @@ def test_a_charged_call_costs_its_own_wall_duration(world): """``charge_inference_time=True`` charges the world what the model really took, so a slow server is scored as slow — at the cost of a trace that inherits the machine's noise.""" policy = RemoteStubPolicy(wall_sec=0.2, chunk=slow_chunk()) - played = _run_episode(world, policy, ChunkedSchedule(), charge_inference_time=True) + played = _run_episode(world, policy, ChunkPlayer(), charge_inference_time=True) assert played, 'no command was played' assert played[0][0] >= 0.2, f'first command at {played[0][0]}s, under the 0.2s the function took' @@ -2057,7 +2069,7 @@ def test_a_real_rig_pays_wall_time_whatever_the_trial_asks_for(world): """The knob is sim-only: a real rig pays what its functions take, so a task leaving ``charge_inference_time`` unset does not hold the world for them.""" policy = RemoteStubPolicy(wall_sec=0.2, chunk=slow_chunk()) - played = _run_episode(world, policy, ChunkedSchedule(), charge_inference_time=False, simulated=False) + played = _run_episode(world, policy, ChunkPlayer(), charge_inference_time=False, simulated=False) assert played, 'no command was played' assert played[0][0] >= 0.2, f'first command at {played[0][0]}s, under the 0.2s the function took' @@ -2093,7 +2105,7 @@ def test_the_layers_see_every_tick(world): _run_episode( world, RemoteStubPolicy(wall_sec=0.01, chunk=slow_chunk(0.3, 15)), - ticks | ChunkedSchedule(), + ticks | ChunkPlayer(), charge_inference_time=True, run_sec=1.0, ) @@ -2123,7 +2135,7 @@ def test_an_unavailable_arm_reaches_the_policy_with_its_pose(world, unavailable) """An unavailable arm is not swallowed as "no observation": its measurements reach the policy beside the status. The harness here runs no ``StopOnFault``, so nothing filters it on the way.""" policy = SpyPolicy() - harness = Harness(policy, make_embodiment()) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment()) p = _pair_all(world, harness) driver = ManualDriver([ @@ -2167,7 +2179,7 @@ def test_every_arm_of_a_bimanual_rig_reports_its_own_status(world): pimm.ControlSystemEmitter(Passive()), ) policy = SpyPolicy() - harness = Harness(policy, embodiment) + harness = Harness(ChunkPlayer().wrap(policy), embodiment) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) left_em = world.pair(harness.observations[left]) @@ -2197,7 +2209,7 @@ def emit_states(): def test_a_stop_clears_the_chunk_in_the_round_the_fault_is_seen(world): """A stop has no waypoints to place: an arm that faults mid-chunk stops in the round its fault is seen.""" fault_at, period = 0.5, 0.005 - stack = StopOnFault() | ChunkedSchedule() + stack = StopOnFault() | ChunkPlayer() harness = Harness(stack.wrap(RemoteStubPolicy(chunk=slow_chunk(1.0, 50))), make_embodiment(simulated=True)) grip_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) @@ -2239,7 +2251,7 @@ def infer(self, obs): raise RuntimeError('inference boom') with pimm.World() as world: - harness = Harness(ServedPolicy(_HangingInfer([], hang_sec)), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(ServedPolicy(_HangingInfer([], hang_sec))), make_embodiment()) cmd_recorder = RecordingEmitter() ds_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) @@ -2281,7 +2293,7 @@ def infer(self, obs): return answer with pimm.World() as world: - harness = Harness(ServedPolicy(_HangingInfer([], hang_sec)), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(ServedPolicy(_HangingInfer([], hang_sec))), make_embodiment()) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) @@ -2306,8 +2318,8 @@ def infer(self, obs): @pytest.mark.timeout(20.0) def test_the_session_is_closed_only_once_its_call_has_left_it(): - """``RemoteSession.close`` shuts the websocket the function in flight is talking over, and ``Session`` - asks for no thread safety, so the session is retired with its runtime and closed after it.""" + """Ending an episode shuts the websocket the function in flight is talking over, so the episode is + retired with its runtime and released after it.""" inside_at_close = [] class _HangingInfer(_FakeInferenceSession): @@ -2326,7 +2338,7 @@ def close(self): inside_at_close.append(self.inside) with pimm.World() as world: - harness = Harness(ServedPolicy(_HangingInfer()), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(ServedPolicy(_HangingInfer())), make_embodiment()) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) @@ -2354,13 +2366,17 @@ def test_a_rescheduled_trajectory_clears_the_channels_it_omits(world): """A trajectory naming only one channel replaces the whole schedule: the omitted channel stops being played rather than draining the previous trajectory's tail.""" - class _GripThenArm(Session): + class _GripThenArmPolicy(Policy): """First a two-channel chunk, then an arm-only one that must silence the gripper.""" def __init__(self): self._calls = 0 - def __call__(self, obs, time_ns): + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} + + def _infer(self, obs): self._calls += 1 pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) command = CartesianPosition(pose=pose) @@ -2371,11 +2387,7 @@ def __call__(self, obs, time_ns): ] return [{keys.ROBOT_COMMAND: command, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10)] - class _GripThenArmPolicy(Policy): - def new_session(self, context=None, rt=None): - return _GripThenArm() - - harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(_GripThenArmPolicy()), make_embodiment()) grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) @@ -2401,7 +2413,7 @@ def new_session(self, context=None, rt=None): @pytest.mark.timeout(3.0) def test_manual_commands_are_emitted_as_plain_values(world): """An operator's command bypasses the schedule: it is the command, not a plan to play.""" - harness = Harness(StubPolicy(), make_embodiment()) + harness = Harness(ChunkPlayer().wrap(StubPolicy()), make_embodiment()) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) @@ -2423,7 +2435,7 @@ def test_finishing_discards_a_call_that_is_still_in_flight(world): """Finishing while the model is still inside its call throws that answer away: the trajectory it carries never reaches the devices.""" policy = RemoteStubPolicy(wall_sec=1.0, chunk=slow_chunk()) - harness = Harness(ChunkedSchedule().wrap(policy), make_embodiment(simulated=True)) + harness = Harness(ChunkPlayer().wrap(policy), make_embodiment(simulated=True)) cmd_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 6d55f6f62..520ff030c 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -1,5 +1,7 @@ -"""Unit tests for Layer composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" +"""Unit tests for Layer composition, ChunkPlayer, TemporalStack, and the policy-pipeline algebra.""" +from collections.abc import Mapping +from contextlib import contextmanager from typing import Any import numpy as np @@ -11,7 +13,7 @@ from positronic.geom import Rotation, Transform3D from positronic.policy import spec from positronic.policy.action import AbsoluteJointsAction, AbsolutePositionAction, IKJointsAction, JointDeltaAction -from positronic.policy.base import Layer, Policy, Session +from positronic.policy.base import INFER, Answer, Fn, Layer, Policy, Runtime, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -23,62 +25,118 @@ RestrictImageSize, SetControlMode, ) -from positronic.policy.layers import ChunkedSchedule, StopOnFault, TemporalStack +from positronic.policy.layers import ChunkPlayer, StopOnFault, TemporalStack from positronic.policy.observation import ObservationCodec +from positronic.policy.tests.conftest import Done -class _ConstSession(Session): +class _Const: + """Answers the same chunk to every call.""" + def __init__(self, actions): self._actions = actions self.call_count = 0 - def __call__(self, obs, time_ns): + def __call__(self, obs): self.call_count += 1 return self._actions class _ConstPolicy(Policy): + """Answers the same chunk to every call, and counts the calls of the episode it last opened.""" + def __init__(self, actions): self._actions = actions - self._session: _ConstSession | None = None + self.infer = _Const(actions) + + @contextmanager + def episode(self, context=None): + self.infer = _Const(self._actions) + yield {INFER: self.infer} + + +class Pending(Answer): + """A call the test answers by hand, so a player can be watched while it waits.""" + + def __init__(self, value: Any = None, failure: BaseException | None = None): + self._value = value + self._failure = failure + self._done = False + + def answer(self) -> None: + self._done = True - def new_session(self, context=None, rt=None): - self._session = _ConstSession(self._actions) - return self._session + def done(self) -> bool: + return self._done + + def result(self) -> Any: + if self._failure is not None: + raise self._failure + return self._value + + +class _FnRuntime(Runtime): + """Serves one ``INFER``, exactly as the test wrote it.""" + + def __init__(self, infer: Fn): + self._fns: Mapping[str, Fn] = {INFER: infer} + + @property + def fns(self) -> Mapping[str, Fn]: + return self._fns + + +def player(infer: Fn) -> Session: + """A ``ChunkPlayer`` session over ``infer``, as the framework builds one.""" + return ChunkPlayer().wrap(Policy()).new_session(_FnRuntime(infer)) def _obs(now_sec=0.0, status=RobotStatus.AVAILABLE): return {keys.OBS_TIME_NS: int(now_sec * 1e9), keys.ROBOT_STATUS: status} +class _CommandSession(Session): + """Answers a fixed command mapping, and asks for its next call one period later.""" + + POLL_SEC = 0.001 + + def __init__(self, commands): + self._commands = commands + self.call_count = 0 + + def __call__(self, obs, time_ns): + self.call_count += 1 + return self._commands, time_ns + int(self.POLL_SEC * 1e9) + + class TestStopOnFault: @pytest.mark.parametrize('unavailable', [RobotStatus.ERROR, RobotStatus.BUSY]) def test_an_unavailable_arm_stops_what_is_executing(self, unavailable): - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session(_obs(0.0, unavailable), 0) == [] + assert session(_obs(0.0, unavailable), 0) == ({}, int(StopOnFault.POLL_SEC * 1e9)) assert inner.call_count == 0, 'the model was asked about an arm that is not tracking it' def test_an_available_arm_reaches_the_model(self): - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session(_obs(0.0, RobotStatus.AVAILABLE), 0) is not None + assert session(_obs(0.0, RobotStatus.AVAILABLE), 0) == ({'v': 1}, int(_CommandSession.POLL_SEC * 1e9)) assert inner.call_count == 1 def test_an_observation_with_no_arm_status_reaches_the_model(self): """A probe replaying a recording has no arm to stop for.""" - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session({keys.OBS_TIME_NS: 0}, 0) is not None + assert session({keys.OBS_TIME_NS: 0}, 0) == ({'v': 1}, int(_CommandSession.POLL_SEC * 1e9)) assert inner.call_count == 1 def test_either_arm_of_a_bimanual_rig_stops_the_pair(self): """Whichever arm is unavailable stops the pair, and the status counts as its number: a server-side stack reads it off a wire with no enum to carry.""" - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) obs = { keys.OBS_TIME_NS: 0, @@ -86,150 +144,220 @@ def test_either_arm_of_a_bimanual_rig_stops_the_pair(self): f'{keys.ROBOT_STATE}.right.status': int(RobotStatus.ERROR), } - assert session(obs, 0) == [] + assert session(obs, 0) == ({}, int(StopOnFault.POLL_SEC * 1e9)) assert inner.call_count == 0 def test_the_status_a_recording_carries_for_a_taken_arm_stops_the_policy(self): """The numbers are the contract between a rig and a server: 1 is an arm its driver has taken.""" - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) assert (RobotStatus.AVAILABLE, RobotStatus.BUSY, RobotStatus.ERROR) == (0, 1, 3) - assert session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 1}, 0) == [] + assert session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 1}, 0) == ({}, int(StopOnFault.POLL_SEC * 1e9)) assert inner.call_count == 0 def test_the_status_published_for_a_travelling_arm_reaches_the_model(self): """The wire protocol publishes 2 for an arm on its way to a setpoint, which is one taking commands.""" - inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 2}, 0) is not None + assert session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 2}, 0) == ( + {'v': 1}, + int(_CommandSession.POLL_SEC * 1e9), + ) assert inner.call_count == 1 def test_a_status_no_arm_answers_to_raises(self): """A number outside ``RobotStatus`` is the rig and the server disagreeing about the protocol, which is not something to drive an arm through.""" - session = StopOnFault().make_session(_ConstSession([])) + session = StopOnFault().make_session(_CommandSession({})) with pytest.raises(ValueError): session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 99}, 0) - def test_recovery_plans_afresh_instead_of_resuming(self): - """The stop resets the scheduler below it, so the first observation from an available arm infers - again rather than waiting out the chunk stamped before.""" + def test_recovery_plans_afresh_instead_of_resuming(self, open_session): + """The stop drops the chunk the player holds, so the first observation from an available arm plans + again rather than playing out the chunk stamped before.""" inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 1.0}]) - session = (StopOnFault() | ChunkedSchedule()).wrap(inner).new_session() + session = open_session((StopOnFault() | ChunkPlayer()).wrap(inner)) - assert session(_obs(0.0), 0) is not None # a chunk that runs until 1.0 - assert session(_obs(0.2, RobotStatus.ERROR), int(0.2e9)) == [] - assert session(_obs(0.3), int(0.3e9)) is not None + assert session(_obs(0.0), 0) == ({'v': 1}, int(1e9)) # a chunk that runs until 1.0 + assert session(_obs(0.2, RobotStatus.ERROR), int(0.2e9)) == ({}, int(0.2e9) + int(StopOnFault.POLL_SEC * 1e9)) + assert session(_obs(0.3), int(0.3e9)) == ({'v': 1}, int(1.3e9)) + assert inner.infer.call_count == 2 -class _ScriptedSession(Session): - """Answers each of ``script`` in turn — ``None`` where a session has nothing to place yet.""" +class _Scripted: + """Answers each of ``script`` in turn, each entry a handle the player holds.""" def __init__(self, script): self._script = list(script) self.call_count = 0 - def __call__(self, obs, time_ns): + def __call__(self, obs): self.call_count += 1 return self._script.pop(0) -class TestChunkedSchedule: - def test_an_inner_with_no_answer_yet_is_asked_again(self): - """A session that waits for a served function answers ``None``, which is no trajectory. The layer - passes the ``None`` on and asks again on the next observation.""" - inner = _ScriptedSession([None, None, [{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]]) - session = ChunkedSchedule().make_session(inner) +class TestChunkPlayer: + def test_a_call_that_has_not_answered_leaves_the_player_holding(self): + """The player asks one time and holds the handle, so a round it waits through costs no second call.""" + pending = Pending([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _Scripted([pending]) + session = player(inner) - assert session(_obs(0.0), int(1e9)) is None - assert session(_obs(0.1), int(1e9)) is None - assert session(_obs(0.2), int(1e9)) == [{'v': 1, keys.ACTION_TIMESTAMP: 1.0}] - assert inner.call_count == 3 + assert session(_obs(0.0), int(1e9)) == ({}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + assert session(_obs(0.1), int(1e9)) == ({}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + pending.answer() + assert session(_obs(0.2), int(1e9)) == ({'v': 1}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + assert inner.call_count == 1 + + def test_a_cancel_drops_the_chunk_of_the_call_in_flight(self): + """A cancelled player drops the chunk it waited for, because that chunk describes a world the cancel + says has gone, and it asks for a new one.""" + pending = Pending([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _Scripted([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) + session = player(inner) + poll_ns = int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9) + + assert session(_obs(), int(1e9)) == ({}, poll_ns) + pending.answer() + session.cancel() + + assert session(_obs(), int(1e9)) == ({}, poll_ns), 'the cancelled chunk was read and thrown away' + assert session(_obs(), int(1e9)) == ({'v': 2}, poll_ns) + assert inner.call_count == 2 + + def test_a_cancelled_call_still_raises_what_it_failed_with(self): + """A dropped chunk drops no failure. The player reads a cancelled call, so a stalled server raises + to the caller that asked for the episode.""" + pending = Pending(failure=TimeoutError('server stalled')) + session = player(_Scripted([pending])) - def test_first_call_runs_inference(self): - # Relative timestamps: trajectory of duration 0.5s - inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) - policy = ChunkedSchedule().wrap(inner) - session = policy.new_session() - result = session(_obs(), int(1e9)) - assert result is not None - assert len(result) == 2 - # Timestamps stamped to absolute by ChunkedSchedule. - assert result[0][keys.ACTION_TIMESTAMP] == 1.0 - assert result[1][keys.ACTION_TIMESTAMP] == 1.5 - - def test_returns_none_while_trajectory_active(self): - # The trajectory starts at the call time 1.0 and ends at 1.0+0.5=1.5. - inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) - policy = ChunkedSchedule().wrap(inner) - session = policy.new_session() session(_obs(), int(1e9)) - assert session(_obs(), int(1.2e9)) is None - assert session(_obs(), int(1.4e9)) is None + pending.answer() + session.cancel() - def test_re_infers_after_trajectory_consumed(self): - inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session() - session(_obs(1.0), int(1e9)) # trajectory ends at 1.5 - assert session(_obs(1.3), int(1.3e9)) is None - result = session(_obs(1.6), int(1.6e9)) - assert result is not None - assert inner._session.call_count == 2 - - def test_single_action_refires_immediately_after(self): - """Single action at ts=0 → trajectory_end is the call's time → next tick re-infers.""" - policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) - session = policy.new_session() - session(_obs(1.0), int(1e9)) - result = session(_obs(1.01), int(1.01e9)) - assert result is not None - - def test_expiry_is_judged_at_the_observation_instant(self): - """Whether the trajectory has run out is a question about the observation, not about the call's time.""" + with pytest.raises(TimeoutError, match='server stalled'): + session(_obs(), int(1e9)) + + def test_a_cancel_ends_with_the_call_it_was_made_against(self): + """A cancel ends with the call it was made against, even when that call fails. A caller that catches + the failure and keeps the session gets the next chunk.""" + pending = Pending(failure=TimeoutError('server stalled')) + inner = _Scripted([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) + session = player(inner) + + session(_obs(), int(1e9)) + pending.answer() + session.cancel() + with pytest.raises(TimeoutError, match='server stalled'): + session(_obs(), int(1e9)) + + assert session(_obs(), int(1e9)) == ({'v': 2}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + + def test_a_chunk_plays_one_waypoint_at_a_time(self, open_session): + """The player anchors the chunk on the call that receives it and asks for a call at each waypoint.""" inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session() - session(_obs(1.0), int(2e9)) # anchored at the call's time 2.0, so the trajectory ends at 2.5 - assert session(_obs(2.4), int(2e9)) is None - assert session(_obs(2.6), int(2e9)) is not None + session = open_session(ChunkPlayer().wrap(inner)) + + assert session(_obs(), int(1e9)) == ({'v': 1}, int(1.5e9)) + assert session(_obs(), int(1.2e9)) == ({}, int(1.5e9)) + assert inner.infer.call_count == 1 + + def test_waypoints_due_together_keep_every_channel_they_name(self, open_session): + """A round that reaches two waypoints commands both channels, not only the ones the later names.""" + inner = _ConstPolicy([ + {'arm': 1, 'grip': 0.5, keys.ACTION_TIMESTAMP: 0.0}, + {'grip': 0.9, keys.ACTION_TIMESTAMP: 0.01}, + {'arm': 2, keys.ACTION_TIMESTAMP: 0.02}, + {keys.ACTION_TIMESTAMP: 0.03}, + ]) + session = open_session(ChunkPlayer().wrap(inner)) + + assert session(_obs(), int(1e9)) == ({'arm': 1, 'grip': 0.5}, int(1.01e9)) + assert session(_obs(), int(1.025e9)) == ({'grip': 0.9, 'arm': 2}, int(1.03e9)) + + def test_a_new_chunk_supersedes_the_one_it_replaces(self): + """The call that drains a chunk loads the next one, and a channel only the drained chunk named + commands nothing: the driver holds what it last took until the new chunk names that channel.""" + inner = _Scripted([ + Done([{'arm': 1, keys.ACTION_TIMESTAMP: 0.0}, {'grip': 0.9, keys.ACTION_TIMESTAMP: 0.5}]), + Done([{'arm': 2, keys.ACTION_TIMESTAMP: 0.2}]), + ]) + session = player(inner) + + assert session(_obs(), int(1e9)) == ({'arm': 1}, int(1.5e9)) + assert session(_obs(), int(1.5e9)) == ({}, int(1.7e9)) + + def test_a_channel_with_several_waypoints_due_keeps_the_last(self, open_session): + """A round that finds more than one waypoint due commands the latest of them.""" + inner = _ConstPolicy([{'v': i, keys.ACTION_TIMESTAMP: i * 0.1} for i in range(4)]) + session = open_session(ChunkPlayer().wrap(inner)) + + assert session(_obs(), int(1e9)) == ({'v': 0}, int(1.1e9)) + assert session(_obs(), int(1.25e9)) == ({'v': 2}, int(1.3e9)) + + def test_the_call_that_drains_the_chunk_asks_for_the_next_one(self, open_session): + """Draining and re-querying happen in one call, so the chunk after it starts where this one ended.""" + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {keys.ACTION_TIMESTAMP: 0.5}]) + session = open_session(ChunkPlayer().wrap(inner)) + + assert session(_obs(), int(1e9)) == ({'v': 1}, int(1.5e9)) + assert session(_obs(), int(1.5e9)) == ({'v': 1}, int(2.0e9)) + assert inner.infer.call_count == 2 + + def test_a_single_action_is_drained_by_the_call_that_loads_it(self, open_session): + """A chunk of one action at ts=0 is drained by the call that loads it, so the next call re-queries.""" + session = open_session(ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]))) + + assert session(_obs(1.0), int(1e9)) == ({'v': 1}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + assert session(_obs(1.01), int(1.01e9)) == ({'v': 1}, int(1.01e9) + int(ChunkPlayer.POLL_SEC * 1e9)) + + def test_a_waypoint_naming_no_channel_commands_nothing(self, open_session): + """The codecs close a chunk with a timestamp-only sentinel; it states where the chunk ends.""" + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {keys.ACTION_TIMESTAMP: 0.5}]) + session = open_session(ChunkPlayer().wrap(inner)) + + assert session(_obs(), int(1e9)) == ({'v': 1}, int(1.5e9)) + assert session(_obs(), int(1.4e9)) == ({}, int(1.5e9)) + + def test_a_chunk_timed_against_another_clock_is_refused(self, open_session): + """A chunk that reaches the player already anchored would place its waypoints decades out.""" + session = open_session(ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 1.77e18 / 1e9}]))) + + with pytest.raises(ValueError, match='clock of its own'): + session(_obs(), int(1e9)) class TestPipelineComposition: """Test | operator across Layer and Codec types.""" - def test_layer_pipe_layer(self): - pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() + def test_layer_pipe_layer(self, open_session): + pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkPlayer() assert isinstance(pipeline, Layer) - policy = pipeline.wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) - session = policy.new_session() - result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}, int(1e9)) - assert result is not None - assert result[0]['v'] == 1 + session = open_session(pipeline.wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]))) + assert session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}, int(1e9)) == ( + {'v': 1}, + int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9), + ) def test_codec_pipe_layer(self): - codec = ActionTimestamp(fps=10.0) - pipeline = codec | ChunkedSchedule() + """A codec wraps the work the player plays, so composing one above a player is refused.""" + pipeline = ActionTimestamp(fps=10.0) | ChunkPlayer() assert isinstance(pipeline, Layer) - policy = pipeline.wrap(_ConstPolicy([{'action': 'test', keys.ACTION_TIMESTAMP: 0.0}])) - session = policy.new_session() - result = session(_obs(), int(1e9)) - assert result is not None + with pytest.raises(AssertionError, match='under the layer that plays the chunk'): + pipeline.wrap(_ConstPolicy([{'action': 'test'}])) - def test_full_pipeline(self): + def test_full_pipeline(self, open_session): codec = ActionTimestamp(fps=10.0) - pipeline = ChunkedSchedule() | codec + pipeline = ChunkPlayer() | codec assert isinstance(pipeline, Layer) - # 5 raw actions → codec stamps relative 0.0, 0.1, 0.2, 0.3, 0.4 - # → ChunkedSchedule shifts to 1.0, 1.1, 1.2, 1.3, 1.4 (call time 1.0). - policy = pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)])) - session = policy.new_session() - result = session(_obs(), int(1e9)) - assert result is not None - assert result[0][keys.ACTION_TIMESTAMP] == 1.0 - # Second call within trajectory window returns None (ChunkedSchedule). - assert session(_obs(), int(1.2e9)) is None + # 5 raw actions → codec stamps relative 0.0, 0.1, 0.2, 0.3, 0.4 and closes the chunk at 0.5 + # → ChunkPlayer plays them from the call's time 1.0. + session = open_session(pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)]))) + assert session(_obs(), int(1e9)) == ({'action': 'a0'}, int(1.1e9)) + assert session(_obs(), int(1.2e9)) == ({'action': 'a2'}, int(1.3e9)) def test_codec_and_stays_codec_only(self): """& only works between codecs, not layers.""" @@ -267,23 +395,21 @@ def test_parallel_frame_codecs_keep_the_frame_they_share(self): ) -class _CaptureSession(Session): +class _CapturePolicy(Policy): + """Records every observation its inference is given.""" + def __init__(self): self.seen = [] - def __call__(self, obs, time_ns): + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} + + def _infer(self, obs): self.seen.append(obs) return [] -class _CapturePolicy(Policy): - def __init__(self): - self.session = _CaptureSession() - - def new_session(self, context=None, rt=None): - return self.session - - def _stack_obs(now_sec, value): return {keys.OBS_TIME_NS: int(now_sec * 1e9), 'v': np.array([value])} @@ -328,40 +454,40 @@ def test_every_arm_channel_is_stamped(self): class TestTemporalStack: OFFSETS = (-0.2, -0.1, 0.0) - def test_pad_start_repeats_oldest(self): + def test_pad_start_repeats_oldest(self, open_session): inner = _CapturePolicy() - session = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS).wrap(inner).new_session() + session = open_session((TemporalStack(keys=('v',), offsets_sec=self.OFFSETS) | ChunkPlayer()).wrap(inner)) session(_stack_obs(0.0, 1.0), 0) - stack = inner.session.seen[0]['v'] + stack = inner.seen[0]['v'] assert stack.shape == (3, 1) assert (stack == 1.0).all() - def test_no_pad_start_grows_from_one(self): + def test_no_pad_start_grows_from_one(self, open_session): inner = _CapturePolicy() - layer = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS, pad_start=False) - session = layer.wrap(inner).new_session() + layer = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS, pad_start=False) | ChunkPlayer() + session = open_session(layer.wrap(inner)) session(_stack_obs(0.0, 1.0), 0) - assert inner.session.seen[0]['v'].shape == (1, 1) + assert inner.seen[0]['v'].shape == (1, 1) session(_stack_obs(0.1, 2.0), int(0.1e9)) - assert inner.session.seen[1]['v'].shape == (2, 1) - assert inner.session.seen[1]['v'][:, 0].tolist() == [1.0, 2.0] + assert inner.seen[1]['v'].shape == (2, 1) + assert inner.seen[1]['v'][:, 0].tolist() == [1.0, 2.0] session(_stack_obs(0.2, 3.0), int(0.2e9)) - assert inner.session.seen[2]['v'].shape == (3, 1) - assert inner.session.seen[2]['v'][:, 0].tolist() == [1.0, 2.0, 3.0] + assert inner.seen[2]['v'].shape == (3, 1) + assert inner.seen[2]['v'][:, 0].tolist() == [1.0, 2.0, 3.0] - def test_no_pad_start_full_window_matches_padded(self): + def test_no_pad_start_full_window_matches_padded(self, open_session): offsets = self.OFFSETS stacks = {} for pad_start in (True, False): inner = _CapturePolicy() - layer = TemporalStack(keys=('v',), offsets_sec=offsets, pad_start=pad_start) - session = layer.wrap(inner).new_session() + layer = TemporalStack(keys=('v',), offsets_sec=offsets, pad_start=pad_start) | ChunkPlayer() + session = open_session(layer.wrap(inner)) for i in range(4): session(_stack_obs(0.1 * i, float(i)), round(0.1 * i * 1e9)) - stacks[pad_start] = inner.session.seen[-1]['v'] + stacks[pad_start] = inner.seen[-1]['v'] assert stacks[True].shape == stacks[False].shape == (3, 1) assert (stacks[True] == stacks[False]).all() @@ -371,7 +497,7 @@ class TestPipelineSpec: def test_split_on_marker(self): stack = TemporalStack(keys=('v',), offsets_sec=(0.0,)) - sched = ChunkedSchedule() + sched = ChunkPlayer() codec = ActionTimestamp(fps=10.0) local, border, rem = spec.split(stack | sched | spec.remote | codec) assert local is not None and local._layers() == (stack, sched) @@ -380,12 +506,12 @@ def test_split_on_marker(self): def test_split_empty_halves(self): assert spec.split(spec.remote) == (None, spec.remote, None) - local, _, rem = spec.split(ChunkedSchedule() | spec.remote) - assert rem is None and isinstance(local, ChunkedSchedule) + local, _, rem = spec.split(ChunkPlayer() | spec.remote) + assert rem is None and isinstance(local, ChunkPlayer) def test_split_requires_exactly_one_marker(self): with pytest.raises(ValueError, match='exactly one'): - spec.split(ChunkedSchedule() | ChunkedSchedule()) + spec.split(ChunkPlayer() | ChunkPlayer()) with pytest.raises(ValueError, match='exactly one'): spec.split(spec.remote | spec.remote) @@ -395,7 +521,7 @@ def test_split_recomposes_codec_half_as_codec(self): def test_border_carries_the_wire_settings(self): """``remote`` is the plain border; calling it describes the wire without changing the split.""" - border = spec.split(ChunkedSchedule() | spec.remote(compress_images=True) | ActionTimestamp(fps=10.0))[1] + border = spec.split(ChunkPlayer() | spec.remote(compress_images=True) | ActionTimestamp(fps=10.0))[1] assert border.compress_images is True assert spec.remote.compress_images is False @@ -404,7 +530,7 @@ def test_marker_cannot_be_applied(self): spec.remote.wrap(_ConstPolicy([])) def test_spec_round_trip(self): - stack = TemporalStack(keys=('a', 'b'), offsets_sec=(-0.5, 0.0), pad_start=False) | ChunkedSchedule() + stack = TemporalStack(keys=('a', 'b'), offsets_sec=(-0.5, 0.0), pad_start=False) | ChunkPlayer() rebuilt = spec.from_spec(stack.to_spec()) assert rebuilt is not None and rebuilt.to_spec() == stack.to_spec() @@ -412,12 +538,12 @@ def test_codec_spec_round_trip(self): obs = ObservationCodec( state={'observation.state': {'grip': 1}}, images={'left': (keys.WRIST_IMAGE, (224, 224))} ) - local = ChunkedSchedule() | ActionTimestamp(fps=10.0) | (obs & AbsolutePositionAction('pose', 'grip')) + local = ChunkPlayer() | ActionTimestamp(fps=10.0) | (obs & AbsolutePositionAction('pose', 'grip')) rebuilt = spec.from_spec(local.to_spec()) assert rebuilt is not None and rebuilt.to_spec() == local.to_spec() def test_leaf_without_args_omits_args_key(self): - assert ChunkedSchedule().to_spec() == {'name': 'chunked_schedule'} + assert ChunkPlayer().to_spec() == {'name': 'chunk_player'} def test_par_topology_round_trips(self, monkeypatch): class _WireCodec(Codec): @@ -440,13 +566,13 @@ def to_spec(self): def test_par_of_non_codecs_is_rejected(self): with pytest.raises(TypeError): - spec.from_spec({'par': [{'name': 'chunked_schedule'}, {'name': 'chunked_schedule'}]}) + spec.from_spec({'par': [{'name': 'chunk_player'}, {'name': 'chunk_player'}]}) def test_empty_declaration_builds_nothing(self): assert spec.from_spec({'seq': []}) is None def test_unknown_name_lists_vocabulary(self): - with pytest.raises(ValueError, match='chunked_schedule'): + with pytest.raises(ValueError, match='chunk_player'): spec.from_spec({'name': 'not_a_layer'}) def test_unknown_arg_fails(self): @@ -461,7 +587,7 @@ def test_the_table_publishes_these_exact_wire_names(self): """The strings a deployed server already declares its local stack with. Spelled out here rather than read off ``WIRE_NAME``, so renaming an attribute cannot quietly rename the wire.""" instances = { - 'chunked_schedule': ChunkedSchedule(), + 'chunk_player': ChunkPlayer(), 'stop_on_fault': StopOnFault(), 'temporal_stack': TemporalStack(('v',), (0.0,)), 'action_timestamp': ActionTimestamp(fps=10.0), @@ -498,7 +624,7 @@ class TestPipe: def test_layer_chain_terminates_into_pipe(self): stack = TemporalStack(keys=('v',), offsets_sec=(0.0,)) - sched = ChunkedSchedule() + sched = ChunkPlayer() codec = ActionTimestamp(fps=10.0) source = spec.PolicySource(_ConstPolicy([])) pipeline = stack | sched | spec.remote | codec | source @@ -518,7 +644,7 @@ def test_bare_marker_terminates_into_pipe(self): assert pipeline.components == (spec.remote,) def test_split_pipe(self): - sched = ChunkedSchedule() + sched = ChunkPlayer() codec = ActionTimestamp(fps=10.0) local, border, rem = spec.split(sched | spec.remote | codec | spec.PolicySource(_ConstPolicy([]))) assert local is sched @@ -527,7 +653,7 @@ def test_split_pipe(self): def test_split_pipe_requires_exactly_one_marker(self): with pytest.raises(ValueError, match='exactly one'): - spec.split(ChunkedSchedule() | spec.PolicySource(_ConstPolicy([]))) + spec.split(ChunkPlayer() | spec.PolicySource(_ConstPolicy([]))) def test_pipe_refuses_a_frame_declared_on_both_sides_of_the_wire(self): """Rig-side and server-side conversion are alternatives; running both puts poses at the product.""" @@ -537,30 +663,26 @@ def test_pipe_refuses_a_frame_declared_on_both_sides_of_the_wire(self): _ = chain | spec.PolicySource(_ConstPolicy([])) def test_pipe_composes_no_further(self): - pipeline: Any = ChunkedSchedule() | spec.remote | spec.PolicySource(_ConstPolicy([])) + pipeline: Any = ChunkPlayer() | spec.remote | spec.PolicySource(_ConstPolicy([])) with pytest.raises(TypeError): _ = pipeline | ActionTimestamp(fps=10.0) with pytest.raises(TypeError): - _ = ChunkedSchedule() | pipeline + _ = ChunkPlayer() | pipeline with pytest.raises(TypeError): _ = pipeline | spec.PolicySource(_ConstPolicy([])) - def test_inline_full_pipe(self): + def test_inline_full_pipe(self, open_session): inner = _ConstPolicy([{'action': f'a{i}'} for i in range(5)]) - policy = spec.inline(ChunkedSchedule() | spec.remote | ActionTimestamp(fps=10.0) | spec.PolicySource(inner)) + policy = spec.inline(ChunkPlayer() | spec.remote | ActionTimestamp(fps=10.0) | spec.PolicySource(inner)) assert isinstance(policy, Policy) - session = policy.new_session() - result = session(_obs(), int(1e9)) - assert result is not None - assert result[0][keys.ACTION_TIMESTAMP] == 1.0 - assert session(_obs(), int(1.2e9)) is None + session = open_session(policy) + assert session(_obs(), int(1e9)) == ({'action': 'a0'}, int(1.1e9)) + assert session(_obs(), int(1.2e9)) == ({'action': 'a2'}, int(1.3e9)) - def test_inline_tolerates_marker_less_pipe(self): + def test_inline_tolerates_marker_less_pipe(self, open_session): inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) - policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) - session = policy.new_session() - result = session(_obs(), int(1e9)) - assert result is not None and result[0][keys.ACTION_TIMESTAMP] == 1.0 + session = open_session(spec.inline(ChunkPlayer() | spec.PolicySource(inner))) + assert session(_obs(), int(1e9)) == ({'v': 1}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) def test_inline_bare_source_pipe_is_the_loaded_policy(self): inner = _ConstPolicy([]) diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index ef09c5056..6e3ffcf2c 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager + import numpy as np import pytest @@ -8,7 +10,7 @@ from positronic.dataset.tests.utils import DummySignal from positronic.geom import Rotation from positronic.policy.action import AbsoluteJointsAction, AbsolutePositionAction -from positronic.policy.base import Policy, Session +from positronic.policy.base import INFER, Policy from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -115,25 +117,15 @@ def test_absolute_joints_action_encode_decode(): assert np.isclose(target_grip, g[0]) -class _FixedSession(Session): - def __init__(self, result): - self._result = result - - def __call__(self, obs, time_ns): - return self._result - - -class _ChunkPolicy(Policy): - def __init__(self, actions: list[dict]): - self._actions = actions +class _FixedPolicy(Policy): + """Answers the same chunk to every call.""" - def new_session(self, context=None, rt=None): - return _FixedSession(list(self._actions)) + def __init__(self, chunk): + self._chunk = chunk - -class _SinglePolicy(Policy): - def new_session(self, context=None, rt=None): - return _FixedSession({'v': 42}) + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: self._chunk} class _PassthroughCodec(Codec): @@ -145,9 +137,9 @@ def encode(self, data): return data -class _MetaPolicy(Policy): - def new_session(self, context=None, rt=None): - return _FixedSession({}) +class _MetaPolicy(_FixedPolicy): + def __init__(self): + super().__init__({}) @property def meta(self): @@ -157,12 +149,19 @@ def meta(self): _T0_OBS = {obs_keys.OBS_TIME_NS: 0} +def _chunk(policy: Policy, obs=_T0_OBS): + """The chunk ``policy``'s inference answers for ``obs``.""" + with policy.episode() as fns: + return fns[INFER](obs) + + def test_action_horizon_sec_truncates_chunk(): actions = [{'v': i} for i in range(10)] # action_horizon_sec=0.1s at action_fps=30 -> 3 actions codec = ActionTiming(fps=30.0, horizon_sec=0.1) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) + assert isinstance(result, list) assert [r['v'] for r in result if 'v' in r] == [0, 1, 2] assert result[-1] == {'timestamp': pytest.approx(0.1)} # horizon sentinel @@ -170,8 +169,8 @@ def test_action_horizon_sec_truncates_chunk(): def test_action_horizon_sec_none_returns_full_chunk(): actions = [{'v': i} for i in range(5)] codec = ActionTiming(fps=30.0) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) assert len(result) == 6 # 5 actions + timestamp sentinel @@ -179,17 +178,17 @@ def test_action_horizon_sec_larger_than_chunk(): actions = [{'v': i} for i in range(3)] # action_horizon_sec=10s at action_fps=10 -> 100 actions max, but only 3 available codec = ActionTiming(fps=10.0, horizon_sec=10.0) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) assert len(result) == 4 # 3 actions + timestamp sentinel (nothing truncated) def test_timestamps_embedded_in_actions(): actions = [{'v': i} for i in range(4)] codec = ActionTiming(fps=10.0) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) - assert len(result) == 5 # 4 actions + timestamp sentinel + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) + assert isinstance(result, list) and len(result) == 5 # 4 actions + timestamp sentinel for i, action in enumerate(result): assert action['timestamp'] == pytest.approx(i * 0.1) @@ -198,9 +197,9 @@ def test_action_horizon_sec_seconds_truncates(): actions = [{'v': i} for i in range(100)] # 0.1s at 30fps -> 3 actions codec = ActionTiming(fps=30.0, horizon_sec=0.1) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) - assert len(result) == 4 # 3 actions + horizon sentinel + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) + assert isinstance(result, list) and len(result) == 4 # 3 actions + horizon sentinel dt = 1.0 / 30.0 for i, action in enumerate(result): assert action['timestamp'] == pytest.approx(i * dt) @@ -249,17 +248,17 @@ def test_action_horizon_meta(): def test_action_timestamp_and_horizon_compose(): actions = [{'v': i} for i in range(10)] codec = ActionHorizon(0.3) | ActionTimestamp(fps=10.0) - policy = codec.wrap(_ChunkPolicy(actions)) - result = policy.new_session()(_T0_OBS, 0) - assert len(result) == 4 # 3 actions + horizon sentinel + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) + assert isinstance(result, list) and len(result) == 4 # 3 actions + horizon sentinel assert [r['v'] for r in result if 'v' in r] == [0, 1, 2] assert result[-1] == {'timestamp': pytest.approx(0.3)} # horizon sentinel def test_single_action_has_zero_timestamp(): codec = ActionTiming(fps=15.0) - policy = codec.wrap(_SinglePolicy()) - result = policy.new_session()(_T0_OBS, 0) + policy = codec.wrap(_FixedPolicy({'v': 42})) + result = _chunk(policy) assert isinstance(result, dict) assert result['timestamp'] == 0.0 assert result['v'] == 42 diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index 63df5b1fb..7ebaecbe3 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -1,73 +1,64 @@ +from contextlib import contextmanager + import numpy as np from positronic import keys from positronic.drivers.roboarm import command from positronic.drivers.roboarm.command import CartesianPosition from positronic.geom import Rotation, Transform3D -from positronic.policy.base import Policy, Session +from positronic.policy.base import INFER, Policy +from positronic.policy.layers import ChunkPlayer from positronic.policy.recording import ( + _TIMELINE_VALUES, Recorder, _build_blueprint, _command_field_arrays, + _CommandTapSession, _flat_wire, _squeeze_batch, _stack_numeric, ) -class _TrackingSession(Session): - def __init__(self, actions, meta): - self._actions = actions - self._meta = meta - - def __call__(self, obs, time_ns): - return list(self._actions) - - @property - def meta(self): - return self._meta +def _chunk(policy: Policy, obs, times: int = 1): + """The chunk ``policy``'s inference answers for ``obs``, over one episode of ``times`` calls.""" + with policy.episode() as fns: + return [fns[INFER](obs) for _ in range(times)][-1] class _TrackingPolicy(Policy): - """Policy that returns a fixed action chunk and tracks session creation.""" + """Answers a fixed action chunk, and counts the episodes it opened.""" def __init__(self, actions: list[dict] | None = None): - self._actions = actions or [{'action': np.array([1.0, 2.0], dtype=np.float32), 'timestamp': 0.0}] - self.session_count = 0 + default = [{'action': np.array([1.0, 2.0], dtype=np.float32), 'timestamp': 0.0}] + self._actions = default if actions is None else actions + self.episodes = 0 - def new_session(self, context=None, rt=None): - self.session_count += 1 - return _TrackingSession(self._actions, {'policy_key': 'policy_value'}) + @contextmanager + def episode(self, context=None): + self.episodes += 1 + yield {INFER: lambda obs: list(self._actions)} @property def meta(self): return {'policy_key': 'policy_value'} -class _CapturingSession(Session): - """Innermost session that snapshots the Recorder's carried timeline state when called.""" +class _CapturingPolicy(Policy): + """Snapshots the timeline values its inference is called under.""" def __init__(self, rec, actions): self._rec = rec self._actions = actions self.seen_timeline_values = None - self.seen_depth = None - - def __call__(self, obs, time_ns): - self.seen_timeline_values = dict(self._rec._timeline_values) - self.seen_depth = self._rec._depth - return list(self._actions) + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} -class _CapturingPolicy(Policy): - def __init__(self, rec, actions): - self._rec = rec - self._actions = actions - self.last_session = None - - def new_session(self, context=None, rt=None): - self.last_session = _CapturingSession(self._rec, self._actions) - return self.last_session + def _infer(self, obs): + self.seen_timeline_values = dict(_TIMELINE_VALUES.get() or {}) + return list(self._actions) def test_squeeze_batch(): @@ -116,7 +107,7 @@ def test_a_cartesian_chunk_records_the_mode_beside_its_trajectory(tmp_path): actions = [{keys.ROBOT_COMMAND: CartesianPosition(pose=pose, mode=mode), 'timestamp': 0.0}] rec = Recorder(tmp_path) - rec.tap('t').wrap(_TrackingPolicy(actions)).new_session()({'x': 1.0, keys.WALL_TIME_NS: 1}, 0) + _chunk(rec.chunk_tap('t').wrap(_TrackingPolicy(actions)), {'x': 1.0, keys.WALL_TIME_NS: 1}) assert any('mode.impedance.kq' in path for path in rec._series_paths), rec._series_paths assert not any(path.endswith('cartesian_pos/pose') for path in rec._series_paths), 'the pose is the trajectory' @@ -151,31 +142,32 @@ def test_a_chunk_that_switches_law_records_both(): def test_single_tap_file_per_episode(tmp_path): rec = Recorder(tmp_path) - policy = rec.tap('raw').wrap(_TrackingPolicy()) + policy = rec.chunk_tap('raw').wrap(_TrackingPolicy()) for _ in range(3): - session = policy.new_session() - session.close() + with policy.episode(): + pass assert len(list(tmp_path.glob('*.rrd'))) == 3 def test_tap_delegates_inner_call(tmp_path): actions = [{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.1}] - policy = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(actions)) - session = policy.new_session() - result = session({'x': 1.0, keys.WALL_TIME_NS: 1_000_000}, 0) - assert result == actions + policy = Recorder(tmp_path).chunk_tap('t').wrap(_TrackingPolicy(actions)) + assert _chunk(policy, {'x': 1.0, keys.WALL_TIME_NS: 1_000_000}) == actions -def test_tap_meta_passthrough(tmp_path): - """A tap contributes no meta; inner meta passes through.""" - policy = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy()) +def test_tap_names_its_recording_once_the_episode_opens(tmp_path): + """Before any episode there is no file to name, and the inner meta passes through either way.""" + policy = Recorder(tmp_path).chunk_tap('t').wrap(_TrackingPolicy()) assert policy.meta == {'policy_key': 'policy_value'} + with policy.episode(): + assert policy.meta['recording.rrd'].endswith('.rrd') + def test_obs_log_filtering_uses_pure_tap_names(tmp_path): rec = Recorder(tmp_path) - session = rec.tap('cam').wrap(_TrackingPolicy([{'v': 1.0, 'timestamp': 0.0}])).new_session() - session( + _chunk( + rec.chunk_tap('cam').wrap(_TrackingPolicy([{'v': 1.0, 'timestamp': 0.0}])), { keys.WALL_TIME_NS: 1_000_000, keys.TASK: 'pick up the cube', @@ -184,7 +176,6 @@ def test_obs_log_filtering_uses_pure_tap_names(tmp_path): 'joints_list': [0.1, 0.2, 0.3], keys.GRIP: 0.5, }, - 0, ) assert 'cam/camera' in rec._image_paths @@ -205,46 +196,59 @@ def test_logs_command_chunk_without_mutating(tmp_path): {keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': 0.5, 'timestamp': 0.0}, {keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': 0.6, 'timestamp': 0.1}, ] - session = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(actions)).new_session() - result = session({'x': 1.0, keys.WALL_TIME_NS: 1}, 0) + policy = Recorder(tmp_path).chunk_tap('t').wrap(_TrackingPolicy(actions)) + result = _chunk(policy, {'x': 1.0, keys.WALL_TIME_NS: 1}) + assert isinstance(result, list) assert result[0][keys.ROBOT_COMMAND] is actions[0][keys.ROBOT_COMMAND] # unchanged on return -def test_handles_none_actions(tmp_path): - class _NoneSession(Session): - def __call__(self, obs, time_ns): - return None +def test_handles_an_empty_chunk(tmp_path): + rec = Recorder(tmp_path) + policy = rec.chunk_tap('t').wrap(_TrackingPolicy([])) - class _NonePolicy(Policy): - def new_session(self, context=None, rt=None): - return _NoneSession() + assert _chunk(policy, {'x': 1.0}, times=2) == [] + assert rec._series_paths == [] - session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() - assert session({'x': 1.0}, 0) is None - assert session._step == 1 +def test_a_two_action_chunk_is_not_read_as_a_command_pair(tmp_path): + """A two-action chunk has the shape of the ``(commands, resume_at_ns)`` pair a session above the player + answers, and a tap under the player must plot both actions.""" + chunk = [{'v': 1.0, 'timestamp': 0.0}, {'v': 2.0, 'timestamp': 0.5}] + policy = Recorder(tmp_path).chunk_tap('t').wrap(_TrackingPolicy(chunk)) -def test_two_taps_share_one_file_per_episode(tmp_path): + assert _chunk(policy, {'x': 1.0}) == chunk + + +def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path, open_session): + """Above a ``ChunkPlayer`` a session answers commands, so the tap plots a point per round.""" + actions = [{'v': 1.0, 'timestamp': 0.0}, {'v': 2.0, 'timestamp': 0.5}] rec = Recorder(tmp_path) - policy = (rec.tap('raw') | rec.tap('server')).wrap(_TrackingPolicy()) + session = open_session((rec.tap('raw') | ChunkPlayer()).wrap(_TrackingPolicy(actions))) + assert isinstance(session, _CommandTapSession) - session = policy.new_session() - assert len(list(tmp_path.glob('*.rrd'))) == 1 - assert rec._live == 2 + assert session({'x': 1.0, keys.WALL_TIME_NS: 1}, 0) == ({'v': 1.0}, int(0.5e9)) + assert session({'x': 1.0, keys.WALL_TIME_NS: 2}, int(0.25e9)) == ({}, int(0.5e9)) + assert 'raw/series/v' in rec._series_paths + + +def test_two_taps_share_one_file_per_episode(tmp_path): + rec = Recorder(tmp_path) + policy = (rec.chunk_tap('raw') | rec.chunk_tap('server')).wrap(_TrackingPolicy()) - session.close() + with policy.episode(): + assert len(list(tmp_path.glob('*.rrd'))) == 1 + assert rec._live == 2 assert rec._live == 0 - policy.new_session() - assert len(list(tmp_path.glob('*.rrd'))) == 2 + with policy.episode(): + assert len(list(tmp_path.glob('*.rrd'))) == 2 def test_two_taps_log_both_seams(tmp_path): rec = Recorder(tmp_path) actions = [{'v': 1.0, 'timestamp': 0.0}] - policy = (rec.tap('raw') | rec.tap('server')).wrap(_TrackingPolicy(actions)) - session = policy.new_session() - session({'camera': np.zeros((4, 4, 3), dtype=np.uint8), keys.WALL_TIME_NS: 1}, 0) + policy = (rec.chunk_tap('raw') | rec.chunk_tap('server')).wrap(_TrackingPolicy(actions)) + _chunk(policy, {'camera': np.zeros((4, 4, 3), dtype=np.uint8), keys.WALL_TIME_NS: 1}) assert 'raw/camera' in rec._image_paths assert 'server/camera' in rec._image_paths @@ -254,27 +258,23 @@ def test_two_taps_log_both_seams(tmp_path): def test_timeline_values_captured_once_and_carried(tmp_path): rec = Recorder(tmp_path) inner = _CapturingPolicy(rec, [{'v': 1.0, 'timestamp': 0.0}]) - policy = (rec.tap('raw') | rec.tap('server')).wrap(inner) - session = policy.new_session() + policy = (rec.chunk_tap('raw') | rec.chunk_tap('server')).wrap(inner) - session({keys.WALL_TIME_NS: 111, keys.OBS_TIME_NS: 222, 'x': 1.0}, 0) + _chunk(policy, {keys.WALL_TIME_NS: 111, keys.OBS_TIME_NS: 222, 'x': 1.0}) - # Both taps entered before the inner session ran, and both share the values - # captured once from the raw obs at the outermost tap. - assert inner.last_session.seen_depth == 2 - assert inner.last_session.seen_timeline_values == {'wall_time': 111, 'obs_time': 222} + # Both taps entered before the inference ran, and both share the values captured once from the raw obs + # at the outermost tap. + assert inner.seen_timeline_values == {'wall_time': 111, 'obs_time': 222} # Per-inference context is cleared once the outermost tap returns. - assert rec._timeline_values == {} - assert rec._depth == 0 + assert _TIMELINE_VALUES.get() is None def test_partial_timelines_only_set_present_keys(tmp_path): rec = Recorder(tmp_path) inner = _CapturingPolicy(rec, [{'v': 1.0, 'timestamp': 0.0}]) - session = rec.tap('raw').wrap(inner).new_session() - session({keys.WALL_TIME_NS: 555, 'x': 1.0}, 0) # no obs_time_ns - assert inner.last_session.seen_timeline_values == {'wall_time': 555} + _chunk(rec.chunk_tap('raw').wrap(inner), {keys.WALL_TIME_NS: 555, 'x': 1.0}) # no obs_time_ns + assert inner.seen_timeline_values == {'wall_time': 555} def test_concurrent_recorders_write_separate_files(tmp_path): @@ -282,8 +282,8 @@ def test_concurrent_recorders_write_separate_files(tmp_path): stream or collide on filenames.""" rec_a = Recorder(tmp_path) rec_b = Recorder(tmp_path) - rec_a.tap('inference').wrap(_TrackingPolicy()).new_session() - rec_b.tap('inference').wrap(_TrackingPolicy()).new_session() + rec_a.chunk_tap('inference').wrap(_TrackingPolicy()).episode().__enter__() + rec_b.chunk_tap('inference').wrap(_TrackingPolicy()).episode().__enter__() assert rec_a._stream is not rec_b._stream assert len(list(tmp_path.glob('*.rrd'))) == 2 diff --git a/positronic/probe.py b/positronic/probe.py index 8e161ee94..284e4240b 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -1,7 +1,7 @@ """Replay one recorded observation through a live inference endpoint and save an ``.rrd``. Point this at a recorded episode and a moment in it; the observation at that moment is -sent to a remote policy endpoint and the returned action chunk is written to a rerun +sent to a remote policy endpoint and the commands it plays back are written to a rerun recording, with the predicted end-effector trajectory overlaid on the robot's actual pose at that moment. Open the ``.rrd`` to see whether the predicted chunk descends toward the object or rises away. @@ -33,8 +33,8 @@ from positronic import keys from positronic.dataset.dataset import Dataset from positronic.drivers.roboarm.command import CartesianPosition, JointDelta -from positronic.policy import Policy, Recorder, is_action -from positronic.policy.executor import blocking +from positronic.policy import Policy, Recorder +from positronic.policy.executor import Executor # Tap name; the recorder logs each obs/action entity under ``{_TAP}/{key}`` (see recording.py). _TAP = 'raw' @@ -68,35 +68,63 @@ def _meta_doc(name: str, meta: dict) -> str: return f'## {name}\n\n{rows}' -def _is_cartesian_chunk(actions: list[dict] | None) -> bool: - """Whether every action carries a Cartesian end-effector command (so a 3D trajectory exists).""" - return bool(actions) and all(isinstance(a.get(keys.ROBOT_COMMAND), CartesianPosition) for a in actions) +def _play(session, obs: dict, runtime: Executor) -> list[tuple[int, dict]]: + """Every command the session emits for ``obs``, from the endpoint's answer to the end of that chunk. + + The observation is one frozen frame, so the walk moves the clock rather than the world. The first call + asks the endpoint and the wait is what that round trip takes. The session anchors the chunk on the call + that receives it and asks for a call at each waypoint, so the walk follows the instants it names. An + endpoint that commands nothing gives an empty list. + """ + session(obs, time.time_ns()) + # The endpoint's own ``infer_timeout`` bounds the wait, and a round trip that fails raises out of the + # call below rather than here. + runtime.wait() + played: list[tuple[int, dict]] = [] + now_ns = time.time_ns() + while True: + commands, resume_at_ns = session(obs, now_ns) + if commands: + played.append((now_ns, dict(commands))) + if runtime.owes_an_answer: # the chunk has run out, and the session is asking the endpoint again + return played + now_ns = resume_at_ns -def _log_commands(actions: list[dict], wall_ns: int, inf_ns: int) -> None: - """Log the chunk's per-step fields as one named time-series on the obs's live timelines. +def _is_cartesian_chunk(played: list[tuple[int, dict]]) -> bool: + """Whether every command is a Cartesian end-effector pose (so a 3D trajectory exists).""" + return bool(played) and all(isinstance(c.get(keys.ROBOT_COMMAND), CartesianPosition) for _, c in played) - Plots EE pose fields for a Cartesian chunk or joint velocities for a DROID chunk. The tap - already logs this on ``action_time``, but a rerun time-series view plots only the active - timeline, and the images live on ``wall_time`` — so to read the chunk on the same timeline as - the scene we re-stamp each waypoint on ``wall_time`` / ``obs_time`` (offset by its horizon), - with a relative ``chunk_time`` axis alongside. + +def _log_commands(played: list[tuple[int, dict]], obs: dict, wall_ns: int, inf_ns: int) -> None: + """Log the commanded fields as one named time-series on the obs's live timelines, and the predicted + end-effector path in 3D beside the pose the robot was actually at. + + Plots EE pose fields for a Cartesian chunk or joint velocities for a DROID chunk. A rerun time-series + view plots only the active timeline, and the images live on ``wall_time`` — so each command is stamped + on ``wall_time`` / ``obs_time`` offset by its horizon, with a relative ``chunk_time`` axis alongside. """ - if _is_cartesian_chunk(actions): - commands = [a[keys.ROBOT_COMMAND] for a in actions] + if not played: + return + commands = [c for _, c in played] + if _is_cartesian_chunk(played): + poses = [c[keys.ROBOT_COMMAND].pose for c in commands] labels = ['tx', 'ty', 'tz', 'qw', 'qx', 'qy', 'qz'] - rows = [[*c.pose.translation, *c.pose.rotation.as_quat] for c in commands] - elif all(isinstance(a.get(keys.ROBOT_COMMAND), JointDelta) for a in actions): - deltas = [a[keys.ROBOT_COMMAND].velocities for a in actions] + rows = [[*p.translation, *p.rotation.as_quat] for p in poses] + path = np.array([p.translation for p in poses], dtype=np.float64) + rr.log('trajectory/path', rr.LineStrips3D([path], radii=0.0012, colors=[120, 120, 120]), static=True) + actual = np.asarray(obs[keys.EE_POSE], dtype=np.float64).reshape(-1)[:3] + rr.log('trajectory/actual', rr.Points3D([actual], radii=0.006, colors=[245, 245, 245]), static=True) + elif all(isinstance(c.get(keys.ROBOT_COMMAND), JointDelta) for c in commands): + deltas = [c[keys.ROBOT_COMMAND].velocities for c in commands] labels = [f'dq{i}' for i in range(len(deltas[0]))] rows = [list(d) for d in deltas] else: return - horizon = np.array([float(a.get(keys.ACTION_TIMESTAMP, i)) for i, a in enumerate(actions)]) - horizon -= horizon[0] - if all(keys.TARGET_GRIP in a for a in actions): + horizon = np.array([(at_ns - played[0][0]) / 1e9 for at_ns, _ in played]) + if all(keys.TARGET_GRIP in c for c in commands): labels.append(keys.TARGET_GRIP) - rows = [row + [a[keys.TARGET_GRIP]] for row, a in zip(rows, actions, strict=True)] + rows = [row + [c[keys.TARGET_GRIP]] for row, c in zip(rows, commands, strict=True)] data = np.array(rows, float) rr.log('commands', rr.SeriesLines(names=labels), static=True) @@ -114,7 +142,7 @@ def _blueprint(image_keys: list[str], has_trajectory: bool) -> rrb.Blueprint: images = [rrb.Spatial2DView(origin=f'{_TAP}/{key}', name=key) for key in image_keys] top = rrb.Horizontal(*images, rrb.TextDocumentView(origin='meta', name='server')) commands = rrb.TimeSeriesView(origin='commands', name='commands') - trajectory = rrb.Spatial3DView(origin=f'{_TAP}/robot_command/trajectory', name='trajectory') + trajectory = rrb.Spatial3DView(origin='trajectory', name='trajectory') bottom = rrb.Horizontal(trajectory, commands) if has_trajectory else commands return rrb.Blueprint(rrb.Vertical(top, bottom)) @@ -135,26 +163,30 @@ def main( image_keys = [k for k in obs if k.startswith(keys.IMAGE_PREFIX)] rec = Recorder(pos3.sync(output_dir)) - tapped = rec.tap(_TAP).wrap(blocking(policy)) - session = tapped.new_session({keys.TASK: task} if task else None) - meta = dict(session.meta) - name = label or _recording_name(meta) - try: - actions = session(obs, now_ns) - if actions is not None: - actions = [a for a in actions if is_action(a)] # drop the codec's keyless validity sentinel - n = 0 if actions is None else len(actions) - print(f'episode {episode} @ {at:.3f}s (ts={ts}) [{name}]: {n} action(s); rrd -> {output_dir}') - with rec.stream: - rec.stream.send_recording_name(name) - rr.log('meta', rr.TextDocument(_meta_doc(name, meta), media_type=rr.MediaType.MARKDOWN), static=True) - if actions: - _log_commands(actions, now_ns, ts) - # Sent here, not at Recorder construction: the layout depends on the chunk type, which is - # only known after inference — a velocity chunk drops the 3D trajectory view. - rr.send_blueprint(_blueprint(image_keys, _is_cartesian_chunk(actions))) - finally: - session.close() + tapped = rec.tap(_TAP).wrap(policy) + with tapped.episode({keys.TASK: task} if task else None) as fns: + runtime = Executor(fns) + session = tapped.new_session(runtime) + meta = dict(tapped.meta) + name = label or _recording_name(meta) + try: + played = _play(session, obs, runtime) + finally: + # The call that drained the chunk asked the endpoint again; closing the runtime waits that + # answer out. It goes before the episode, whose socket the call still holds. + runtime.close() + session.close() + + print(f'episode {episode} @ {at:.3f}s (ts={ts}) [{name}]: {len(played)} command(s); rrd -> {output_dir}') + stream = rec.stream + assert stream is not None, 'the tap opened the recording when the session was made' + with stream: + stream.send_recording_name(name) + rr.log('meta', rr.TextDocument(_meta_doc(name, meta), media_type=rr.MediaType.MARKDOWN), static=True) + _log_commands(played, obs, now_ns, ts) + # Sent here, not at Recorder construction: the layout depends on the chunk type, which is + # only known after inference — a velocity chunk drops the 3D trajectory view. + rr.send_blueprint(_blueprint(image_keys, _is_cartesian_chunk(played))) @pos3.with_mirror() diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index de08094fd..03bef768d 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -20,9 +20,9 @@ from positronic.dataset.local_dataset import LocalDataset from positronic.drivers.roboarm import command as roboarm_command from positronic.eval import Task -from positronic.policy import Policy, Session +from positronic.policy import INFER, Policy from positronic.policy.codec import ActionTimestamp -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.tests.test_harness import StubPolicy from positronic.simulator.env_server.adapter import EnvAdapter, _in_env_control_frame, _wire_command from positronic.simulator.env_server.client import _CLOSE_ACK_TIMEOUT, EnvConnection @@ -330,7 +330,7 @@ def test_remote_eval_runs_to_timeout_without_done(env_server, tmp_path): trial = number_trials(replace(next(iter(ev.tasks())), timeout_sec=0.1), [{keys.EVAL_SEED: 100}])[0] policy = StubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0) main( - policy=ChunkedSchedule().wrap(policy), + policy=ChunkPlayer().wrap(policy), evals=[replace(ev, tasks=partial(iter, [trial]))], output_dir=str(tmp_path), ) @@ -369,19 +369,14 @@ def __init__(self, command: roboarm_command.CommandType, chunk_len: int): self.chunk_len = chunk_len self.chunks = 0 - def new_session(self, context=None, rt=None): - return _JointposChunkSession(self) + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} - -class _JointposChunkSession(Session): - def __init__(self, policy: _JointposChunks): - self._policy = policy - - def __call__(self, obs, time_ns): - self._policy.chunks += 1 + def _infer(self, obs): + self.chunks += 1 return [ - {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.chunks * 100.0 + i} - for i in range(self._policy.chunk_len) + {keys.ROBOT_COMMAND: self.command, 'target_grip': self.chunks * 100.0 + i} for i in range(self.chunk_len) ] @@ -389,8 +384,8 @@ def __call__(self, obs, time_ns): def test_full_chunk_executes_between_replans(env_server, tmp_path): """The recording proves the contract the DROID jointpos codec makes with RoboLab's client: every action of every chunk lands on the wire — including the final one, which ``ActionTimestamp``'s validity - sentinel gives a full period before ``ChunkedSchedule`` re-infers — and replans arrive exactly - ``chunk_len`` control periods apart.""" + sentinel gives a full period before ``ChunkPlayer`` re-infers — and replans arrive one control period + after that, which is what the answer to the re-infer takes to reach the loop.""" host, port = env_server probe = make_mujoco_env([]) control_dt = probe.reset(0)['control_dt'] @@ -398,7 +393,7 @@ def test_full_chunk_executes_between_replans(env_server, tmp_path): chunk_len = 5 raw = _JointposChunks(roboarm_command.JointPosition(np.zeros(7)), chunk_len) - policy = (ChunkedSchedule() | ActionTimestamp(fps=1.0 / control_dt)).wrap(raw) + policy = (ChunkPlayer() | ActionTimestamp(fps=1.0 / control_dt)).wrap(raw) with pos3.mirror(): ev = remote_stack_cubes_eval(host, port, camera_dict=CAMERAS) trial = number_trials(replace(next(iter(ev.tasks())), timeout_sec=20 * control_dt), [{keys.EVAL_SEED: 100}])[0] @@ -413,7 +408,8 @@ def test_full_chunk_executes_between_replans(env_server, tmp_path): assert values[: len(expected)] == expected starts = [ts for v, ts in executed if v % 100 == 0] - period_ns = chunk_len * control_dt * 1e9 + # The chunk's own span, plus the round the player spends reading the answer to the call that drained it. + period_ns = (chunk_len + 1) * control_dt * 1e9 for earlier, later in zip(starts, starts[1:], strict=False): assert later - earlier == pytest.approx(period_ns, abs=period_ns / (2 * chunk_len)) diff --git a/positronic/tests/test_inference_integration.py b/positronic/tests/test_inference_integration.py index f94fc508e..90d73cd11 100644 --- a/positronic/tests/test_inference_integration.py +++ b/positronic/tests/test_inference_integration.py @@ -21,7 +21,7 @@ from positronic.drivers.roboarm import command as roboarm_command from positronic.drivers.roboarm.models import bundled_panda_model from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Eval, Observation, Task -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.tests.test_harness import RemoteStubPolicy, StubPolicy from positronic.simulator.env_server import telemetry as env_telemetry from positronic.simulator.mujoco.sim import MujocoSim @@ -82,7 +82,7 @@ def close(self): ) trials = number_trials(next(iter(ev.tasks())), [{keys.EVAL_SEED: 100 + i} for i in range(2)]) main( - policy=ChunkedSchedule().wrap(policy), + policy=ChunkPlayer().wrap(policy), evals=[replace(ev, tasks=partial(iter, trials))], output_dir=str(tmp_path), ) @@ -218,9 +218,9 @@ def test_every_trial_records_from_its_own_reset(tmp_path): trials = number_trials(next(iter(ev.tasks())), [{keys.EVAL_SEED: i} for i in range(2)]) with pos3.mirror(): main( - policy=StubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0), + policy=ChunkPlayer().wrap(StubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0)), evals=[replace(ev, tasks=partial(iter, trials))], - # the degenerate obs is not Franka-shaped, so run the policy unwrapped + # the degenerate obs is not Franka-shaped, so nothing but the player wraps the policy output_dir=str(tmp_path), ) @@ -241,7 +241,7 @@ def test_timing_writes_telemetry_sidecars(tmp_path): trials = number_trials(next(iter(ev.tasks())), [{}, {}]) with pos3.mirror(): main( - policy=ChunkedSchedule().wrap( + policy=ChunkPlayer().wrap( RemoteStubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0) ), evals=[replace(ev, tasks=partial(iter, trials))], @@ -289,7 +289,7 @@ def test_countdown_terminates_on_done_records_payload(tmp_path): trial = number_trials(next(iter(ev.tasks())), [{keys.EVAL_SEED: 100}])[0] with pos3.mirror(): main( - policy=StubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0), + policy=ChunkPlayer().wrap(StubPolicy(command=roboarm_command.JointPosition(np.zeros(7)), target_grip=0.0)), evals=[replace(ev, tasks=partial(iter, [trial]))], output_dir=str(tmp_path), ) diff --git a/positronic/tests/testing_coutils.py b/positronic/tests/testing_coutils.py index e9aad7f04..81048d972 100644 --- a/positronic/tests/testing_coutils.py +++ b/positronic/tests/testing_coutils.py @@ -123,12 +123,14 @@ class IdleSession(Session): The recording lands on its policy's ``observations`` list. """ + POLL_SEC = 0.01 + def __init__(self, policy): self._policy = policy def __call__(self, obs, time_ns): self._policy.observations.append(obs) - return [] + return {}, time_ns + int(self.POLL_SEC * 1e9) @property def meta(self): diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index 38e960c60..a0274aa11 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -6,6 +6,8 @@ import subprocess import uuid from collections.abc import Callable +from contextlib import contextmanager +from functools import partial from pathlib import Path from typing import Any @@ -20,7 +22,7 @@ from positronic import keys from positronic.offboard.server import serve from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready -from positronic.policy import Codec, Layer, Policy, Session +from positronic.policy import INFER, Codec, Layer, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints @@ -256,40 +258,37 @@ def stop(self): self.process = None -class _DreamZeroSession(Session): - def __init__(self, client: RoboarenaClient, session_id: str): - self._client = client - self._session_id = session_id - - def __call__(self, obs, time_ns): - obs = dict(obs) - obs[roboarena.SESSION_ID] = self._session_id - action_array = np.asarray(self._client.infer(obs)) - - # Response is (N, 8) — 7 joints + 1 gripper - if action_array.ndim == 1: - return [{'action': action_array}] - return [{'action': action_array[i]} for i in range(action_array.shape[0])] - - def close(self): - try: - self._client.reset(session_id=self._session_id) - except (OSError, TimeoutError, ConnectionClosed): - logger.info('DreamZero session reset skipped: backend connection already gone') - finally: - self._client.close() +def _infer(client: RoboarenaClient, session_id: str, obs): + """One model call: an observation in, an action chunk out.""" + obs = dict(obs) + obs[roboarena.SESSION_ID] = session_id + action_array = np.asarray(client.infer(obs)) + # Response is (N, 8) — 7 joints + 1 gripper + if action_array.ndim == 1: + return [{'action': action_array}] + return [{'action': action_array[i]} for i in range(action_array.shape[0])] class DreamZeroPolicy(Policy): - """Owns the DreamZero subprocess; every session talks to it over its own roboarena connection.""" + """Owns the DreamZero subprocess; every episode talks to it over its own roboarena connection.""" def __init__(self, sp: DreamZeroSubprocess): self._subprocess = sp - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): client = RoboarenaClient(port=self._subprocess.roboarena_port) client.connect() - return _DreamZeroSession(client, str(uuid.uuid4())) + session_id = str(uuid.uuid4()) + try: + yield {INFER: partial(_infer, client, session_id)} + finally: + try: + client.reset(session_id=session_id) + except (OSError, TimeoutError, ConnectionClosed): + logger.info('DreamZero session reset skipped: backend connection already gone') + finally: + client.close() def close(self): self._subprocess.stop() diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index c8a24932f..204724497 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -3,6 +3,8 @@ import os import subprocess from collections.abc import Callable +from contextlib import contextmanager +from functools import partial from pathlib import Path from typing import Any @@ -17,9 +19,9 @@ from positronic.offboard.client import DEFAULT_INFER_TIMEOUT from positronic.offboard.server import serve from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup -from positronic.policy import Policy, Session +from positronic.policy import INFER, Policy from positronic.policy.codec import RestrictImageSize -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints from positronic.vendors import gr00t @@ -200,17 +202,14 @@ def stop(self): ########################################################################################### -class _Gr00tSession(Session): - def __init__(self, client: PolicyClient): - self._client = client - - def __call__(self, obs, time_ns): - action_response, _info = self._client.get_action(obs) - action = {k: v[0] for k, v in action_response.items()} - lengths = {len(v) for v in action.values()} - assert len(lengths) == 1, f'All values in action must have the same length, got {lengths}' - time_horizon = lengths.pop() - return [{k: v[i] for k, v in action.items()} for i in range(time_horizon)] +def _infer(client: PolicyClient, obs): + """One model call: an observation in, an action chunk out.""" + action_response, _info = client.get_action(obs) + action = {k: v[0] for k, v in action_response.items()} + lengths = {len(v) for v in action.values()} + assert len(lengths) == 1, f'All values in action must have the same length, got {lengths}' + time_horizon = lengths.pop() + return [{k: v[i] for k, v in action.items()} for i in range(time_horizon)] class Gr00tPolicy(Policy): @@ -220,9 +219,10 @@ def __init__(self, groot: Gr00tSubprocess, checkpoint_path: str): self._groot = groot self._checkpoint_path = checkpoint_path - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): self._groot.client.reset() - return _Gr00tSession(self._groot.client) + yield {INFER: partial(_infer, self._groot.client)} @property def meta(self): @@ -354,7 +354,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=codecs.ee_quat, source=gr00t_source) def pipeline(codec, source): - return StopOnFault() | ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source + return StopOnFault() | ChunkPlayer() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source # Each entry pairs the codec with the matching GR00T modality config; they must agree with training. diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index a518b437c..bacfb8407 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager +from functools import partial from typing import Any import numpy as np @@ -7,7 +9,7 @@ from lerobot.policies.factory import get_policy_class, make_pre_post_processors from positronic import keys -from positronic.policy import Policy, Session +from positronic.policy import INFER, Policy from positronic.policy.observation import TASK_FIELD @@ -30,44 +32,33 @@ def _detect_device() -> str: return 'cpu' -class _LerobotSession(Session): - def __init__(self, policy, preprocessor, postprocessor, device: str, meta: dict[str, Any]): - self._policy = policy - self._preprocessor = preprocessor - self._postprocessor = postprocessor - self._device = device - self._meta = meta - - def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: - obs_int = {} - for key, val in obs.items(): - if key == TASK_FIELD: - obs_int[key] = val - elif isinstance(val, np.ndarray): - if key.startswith('observation.images.'): - val = torch.from_numpy(np.transpose(val, (2, 0, 1)).copy()).float() / 255.0 - else: - val = torch.from_numpy(val).float() - obs_int[key] = val +def _infer(policy, preprocessor, postprocessor, obs: dict[str, Any]) -> list[dict[str, Any]]: + """One model call: an observation in, an action chunk out.""" + obs_int = {} + for key, val in obs.items(): + if key == TASK_FIELD: + obs_int[key] = val + elif isinstance(val, np.ndarray): + if key.startswith('observation.images.'): + val = torch.from_numpy(np.transpose(val, (2, 0, 1)).copy()).float() / 255.0 else: - obs_int[key] = torch.as_tensor(val) - - if self._preprocessor is not None: - obs_int = self._preprocessor(obs_int) + val = torch.from_numpy(val).float() + obs_int[key] = val + else: + obs_int[key] = torch.as_tensor(val) - action = self._policy.select_action(obs_int) + if preprocessor is not None: + obs_int = preprocessor(obs_int) - if self._postprocessor is not None: - action = self._postprocessor(action) + action = policy.select_action(obs_int) - action = action.cpu().numpy().squeeze(0) - if action.ndim == 1: - return [{'action': action}] - return [{'action': a} for a in action] + if postprocessor is not None: + action = postprocessor(action) - @property - def meta(self) -> dict[str, Any]: - return self._meta + action = action.cpu().numpy().squeeze(0) + if action.ndim == 1: + return [{'action': action}] + return [{'action': a} for a in action] def warm_observation(config: PreTrainedConfig) -> dict[str, Any]: @@ -102,9 +93,10 @@ def config(self) -> PreTrainedConfig: """The checkpoint's own declaration of what this policy takes.""" return self._policy.config - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): self._policy.reset() - return _LerobotSession(self._policy, self._preprocessor, self._postprocessor, self._device, self._meta) + yield {INFER: partial(_infer, self._policy, self._preprocessor, self._postprocessor)} @property def meta(self) -> dict[str, Any]: diff --git a/positronic/vendors/lerobot/server.py b/positronic/vendors/lerobot/server.py index 706e9806d..3cae0a5c0 100644 --- a/positronic/vendors/lerobot/server.py +++ b/positronic/vendors/lerobot/server.py @@ -12,7 +12,7 @@ from positronic.offboard.server_utils import run_with_progress, warmup from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, Pipeline, remote from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.vendors.lerobot import codecs as lerobot_codecs @@ -61,7 +61,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=lerobot_codecs.ee, source=lerobot_source) def pipeline(codec: Codec, source: ModelSource) -> Pipeline: - return StopOnFault() | ChunkedSchedule() | RestrictImageSize(512, 512) | remote | codec | source + return StopOnFault() | ChunkPlayer() | RestrictImageSize(512, 512) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/lerobot/tests/test_server.py b/positronic/vendors/lerobot/tests/test_server.py index 6cd5e4562..161839c85 100644 --- a/positronic/vendors/lerobot/tests/test_server.py +++ b/positronic/vendors/lerobot/tests/test_server.py @@ -6,7 +6,7 @@ from positronic.offboard.protocol import deserialise from positronic.offboard.server import PolicyServer -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.spec import remote pytest.importorskip('lerobot', minversion='0.4') @@ -44,7 +44,7 @@ async def close(self, **kwargs): async def test_lerobot_server_uses_configured_checkpoint(monkeypatch): monkeypatch.setattr('positronic.utils.checkpoints.list_checkpoints', lambda _path: ['42']) - server = PolicyServer(ChunkedSchedule() | remote | LerobotSource('s3://bucket/exp', checkpoint='42')) + server = PolicyServer(ChunkPlayer() | remote | LerobotSource('s3://bucket/exp', checkpoint='42')) requested = {} @@ -74,7 +74,7 @@ async def fake_get_policy(checkpoint_id: str, websocket=None): async def test_lerobot_server_reports_missing_checkpoint(monkeypatch): monkeypatch.setattr('positronic.utils.checkpoints.list_checkpoints', lambda _path: ['41']) - server = PolicyServer(ChunkedSchedule() | remote | LerobotSource('s3://bucket/exp', checkpoint='42')) + server = PolicyServer(ChunkPlayer() | remote | LerobotSource('s3://bucket/exp', checkpoint='42')) server._manager.get_policy = AsyncMock() with pytest.raises(ValueError, match=r"Configured checkpoint not found: 42. Available: \['41'\]"): @@ -88,7 +88,7 @@ async def test_lerobot_server_reports_unknown_checkpoint_id(monkeypatch): monkeypatch.setattr('positronic.utils.checkpoints.list_checkpoints', lambda _path: ['41']) monkeypatch.setattr('positronic.utils.checkpoints.get_latest_checkpoint', lambda _path: '41') - server = PolicyServer(ChunkedSchedule() | remote | LerobotSource('s3://bucket/exp')) + server = PolicyServer(ChunkPlayer() | remote | LerobotSource('s3://bucket/exp')) server._manager.get_policy = AsyncMock(return_value=MagicMock()) server._manager.release_session = AsyncMock() await server._startup() diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 8b6ac0667..beaa3e7c8 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Mapping +from contextlib import contextmanager from functools import partial from typing import Any @@ -14,9 +14,8 @@ from positronic import keys from positronic.cfg import codecs -from positronic.policy import Codec, Policy, Session -from positronic.policy.base import Answer, Runtime -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy import INFER, Codec, Policy +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.observation import TASK_FIELD from positronic.policy.spec import PolicySource, inline from positronic.utils.checkpoints import resolve_checkpoint @@ -81,54 +80,15 @@ def _infer(policy: PreTrainedPolicy, device: str, obs: dict[str, Any]) -> list[d class LerobotPolicy(Policy): - _INFER = 'infer' - - class _Session(Session): - """Per-episode session that gives the model call to the runtime, and answers the chunk on a later call.""" - - def __init__(self, rt: Runtime, meta: dict[str, Any]): - self._rt = rt - self._meta = meta - self._answer: Answer | None = None - self._cancelled = False - - def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]] | None: - if self._answer is None: - self._answer = self._rt.fns[LerobotPolicy._INFER](obs) - return None - if not self._answer.done(): - return None - answer, cancelled = self._answer, self._cancelled - # The answer and the flag are cleared before the read, because ``result`` raises what the model - # call raised. A cancel then ends with the answer it was made against, and never drops the next - # chunk. - self._answer, self._cancelled = None, False - result = answer.result() - return None if cancelled else result - - def cancel(self): - # The cancel says the world the chunk applies to has gone. The session still reads the model call - # for its failure, and drops the chunk that comes with it. - self._cancelled = self._answer is not None - - @property - def meta(self) -> dict[str, Any]: - return self._meta - def __init__(self, policy: PreTrainedPolicy, device: str | None = None, extra_meta: dict[str, Any] | None = None): self._device = device or _detect_device() self._policy = policy.to(self._device) self._meta = extra_meta or {} - def new_session(self, context=None, rt=None): - if rt is None: - raise ValueError('A lerobot session runs its model on a runtime: pass rt to new_session.') + @contextmanager + def episode(self, context=None): self._policy.reset() - return LerobotPolicy._Session(rt, self._meta) - - @property - def functions(self) -> Mapping[str, Callable[..., Any]]: - return {self._INFER: partial(_infer, self._policy, self._device)} + yield {INFER: partial(_infer, self._policy, self._device)} @property def meta(self) -> dict[str, Any]: @@ -161,4 +121,4 @@ def act(checkpoints_dir: str, checkpoint: str | None, n_action_steps: int | None ) def act_absolute(base: Policy, codec: Codec): """ACT with the absolute-position codec, composed in-process.""" - return inline(StopOnFault() | ChunkedSchedule() | codec | PolicySource(base)) + return inline(StopOnFault() | ChunkPlayer() | codec | PolicySource(base)) diff --git a/positronic/vendors/lerobot_0_3_3/server.py b/positronic/vendors/lerobot_0_3_3/server.py index d5907aa3a..c8dca0f2e 100644 --- a/positronic/vendors/lerobot_0_3_3/server.py +++ b/positronic/vendors/lerobot_0_3_3/server.py @@ -15,7 +15,7 @@ from positronic.offboard.server_utils import run_with_progress, warmup from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.vendors.lerobot_0_3_3 import codecs as lerobot_codecs @@ -84,7 +84,7 @@ def meta(self, model_id: str) -> dict[str, Any]: # so none has a transform to declare. @cfn.config(codec=lerobot_codecs.ee, source=lerobot_source) def pipeline(codec: Codec, source: ModelSource): - return StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) | remote | codec | source + return StopOnFault() | ChunkPlayer() | RestrictImageSize(224, 224) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/lerobot_0_3_3/tests/test_server.py b/positronic/vendors/lerobot_0_3_3/tests/test_server.py index 67e3744bf..d47bf971b 100644 --- a/positronic/vendors/lerobot_0_3_3/tests/test_server.py +++ b/positronic/vendors/lerobot_0_3_3/tests/test_server.py @@ -5,7 +5,7 @@ from starlette.datastructures import QueryParams from positronic.offboard.protocol import deserialise -from positronic.policy.layers import ChunkedSchedule +from positronic.policy.layers import ChunkPlayer from positronic.policy.spec import remote pytest.importorskip('torch') @@ -74,7 +74,7 @@ def _make_server(checkpoint: str | None) -> PolicyServer: source = lerobot_server.LerobotSource( policy_factory=lambda _checkpoint: MagicMock(), checkpoints_dir='s3://bucket/exp', checkpoint=checkpoint ) - return PolicyServer(ChunkedSchedule() | remote | source) + return PolicyServer(ChunkPlayer() | remote | source) @pytest.mark.asyncio diff --git a/positronic/vendors/molmoact2/README.md b/positronic/vendors/molmoact2/README.md index 84b51e730..9f9516e8e 100644 --- a/positronic/vendors/molmoact2/README.md +++ b/positronic/vendors/molmoact2/README.md @@ -76,5 +76,5 @@ ships one, `droid` (source: [`codecs.py`](./codecs.py)): command (no IK at runtime). - **Observation**: 3 cameras (2 exterior + 1 wrist) + 8-D state + language prompt. - **Inference**: `norm_tag='franka_droid'`, continuous action mode; the model emits a 15-step action chunk at - 15 Hz, executed in full by the client's declared `ChunkedSchedule`. + 15 Hz, played in full by the client's declared `ChunkPlayer`. - **Wire protocol**: Positronic's standard WebSocket protocol — see [Connect Your Model](../../../docs/connect-your-model.md). diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index d74599a7f..ba4c4648c 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager +from functools import partial from typing import Any import numpy as np @@ -5,7 +7,7 @@ from transformers import AutoModelForImageTextToText, AutoProcessor from positronic import keys -from positronic.policy import Policy, Session +from positronic.policy import INFER, Policy from positronic.vendors import molmoact2 # The three views and the 8-D ``[joint_positions(7), grip(1)]`` state of the DROID action space this vendor @@ -24,36 +26,24 @@ def warm_observation() -> dict[str, Any]: } -class _MolmoAct2Session(Session): - def __init__(self, model, processor, norm_tag: str, num_steps: int, meta: dict[str, Any]): - self._model = model - self._processor = processor - self._norm_tag = norm_tag - self._num_steps = num_steps - self._meta = meta - - def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: - # predict_action is decorated @torch.no_grad() and manages its own precision: the model loads - # in bfloat16 and runs bf16 throughout (its autocast path only guards fp32 inputs), so an - # external torch.inference_mode() / torch.autocast wrap or a detach() would all be redundant. - out = self._model.predict_action( - processor=self._processor, - images=obs[molmoact2.IMAGES], - task=obs.get(molmoact2.TASK, ''), - state=np.asarray(obs[molmoact2.STATE], dtype=np.float32), - norm_tag=self._norm_tag, - inference_action_mode='continuous', - enable_depth_reasoning=False, - num_steps=self._num_steps, - normalize_language=True, - enable_cuda_graph=False, - ) - actions = out.actions[0].float().cpu().numpy() - return [{'action': action} for action in actions] - - @property - def meta(self) -> dict[str, Any]: - return self._meta +def _infer(model, processor, norm_tag: str, num_steps: int, obs: dict[str, Any]) -> list[dict[str, Any]]: + """One model call: an observation in, an action chunk out.""" + # predict_action is decorated @torch.no_grad() and manages its own precision: the model loads + # in bfloat16 and runs bf16 throughout (its autocast path only guards fp32 inputs), so an + # external torch.inference_mode() / torch.autocast wrap or a detach() would all be redundant. + out = model.predict_action( + processor=processor, + images=obs[molmoact2.IMAGES], + task=obs.get(molmoact2.TASK, ''), + state=np.asarray(obs[molmoact2.STATE], dtype=np.float32), + norm_tag=norm_tag, + inference_action_mode='continuous', + enable_depth_reasoning=False, + num_steps=num_steps, + normalize_language=True, + enable_cuda_graph=False, + ) + return [{'action': action} for action in out.actions[0].float().cpu().numpy()] class MolmoAct2Policy(Policy): @@ -66,8 +56,9 @@ def __init__(self, model_id: str, *, device_map: str = 'auto', norm_tag: str = ' self._num_steps = num_steps self._meta = {keys.TYPE: 'molmoact2', 'norm_tag': norm_tag} - def new_session(self, context=None, rt=None) -> Session: - return _MolmoAct2Session(self._model, self._processor, self._norm_tag, self._num_steps, self._meta) + @contextmanager + def episode(self, context=None): + yield {INFER: partial(_infer, self._model, self._processor, self._norm_tag, self._num_steps)} @property def meta(self) -> dict[str, Any]: diff --git a/positronic/vendors/molmoact2/server.py b/positronic/vendors/molmoact2/server.py index 37ab71e82..4b826089d 100644 --- a/positronic/vendors/molmoact2/server.py +++ b/positronic/vendors/molmoact2/server.py @@ -9,7 +9,7 @@ from positronic.offboard.server_utils import warmup from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.vendors.molmoact2 import codecs as molmoact2_codecs from positronic.vendors.molmoact2.policy import MolmoAct2Policy, warm_observation @@ -60,7 +60,7 @@ def meta(self, model_id: str) -> dict[str, Any]: @cfn.config(codec=molmoact2_codecs.droid, source=molmoact2_source) def pipeline(codec: Codec, source: ModelSource): - return StopOnFault() | ChunkedSchedule() | RestrictImageSize() | remote | codec | source + return StopOnFault() | ChunkPlayer() | RestrictImageSize() | remote | codec | source droid = pipeline diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 9b72b2193..9c6b1dce6 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -3,6 +3,8 @@ import socket import subprocess from collections.abc import Callable +from contextlib import contextmanager +from functools import partial from pathlib import Path from typing import Any @@ -14,9 +16,9 @@ from positronic import geom, keys from positronic.offboard.server import serve from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup -from positronic.policy import Codec, Policy, Session +from positronic.policy import INFER, Codec, Policy from positronic.policy.codec import ChangeEEFrame, RestrictImageSize -from positronic.policy.layers import ChunkedSchedule, StopOnFault +from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import get_latest_checkpoint, list_checkpoints from positronic.vendors import openpi @@ -124,14 +126,9 @@ def stop(self): ########################################################################################### -class _OpenpiSession(Session): - def __init__(self, client: WebsocketClientPolicy): - self._client = client - - def __call__(self, obs, time_ns): - response = self._client.infer(obs) - actions = response['actions'] - return [{'action': a} for a in actions] +def _infer(client: WebsocketClientPolicy, obs): + """One model call: an observation in, an action chunk out.""" + return [{'action': a} for a in client.infer(obs)['actions']] class OpenpiPolicy(Policy): @@ -140,10 +137,11 @@ class OpenpiPolicy(Policy): def __init__(self, subproc: OpenpiSubprocess): self._subproc = subproc - def new_session(self, context=None, rt=None): + @contextmanager + def episode(self, context=None): client = self._subproc.client client.reset() - return _OpenpiSession(client) + yield {INFER: partial(_infer, client)} def close(self): self._subproc.stop() @@ -258,11 +256,12 @@ def pipeline(codec: Codec, source: ModelSource, ee_frame: geom.Transform3D | Non ``ee_frame`` places the end-effector frame this checkpoint's poses live in relative to ``DEFAULT_FRAME`` (``models.DROID_EE_FRAME``); ``None`` for a checkpoint trained in ``default``, or one speaking joints. """ - local = StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) + local = StopOnFault() | ChunkPlayer() if ee_frame is not None: - # Outermost, so everything downstream — the wire, the server's codec — sees poses already in ``ee_frame``. - local = ChangeEEFrame(ee_frame) | local - return local | remote | codec | source + # Under the player, which answers commands rather than the chunk this converts, and above the wire, + # so everything downstream — the wire, the server's codec — sees poses already in ``ee_frame``. + local = local | ChangeEEFrame(ee_frame) + return local | RestrictImageSize(224, 224) | remote | codec | source # These bind no checkpoint, so they state no frame: whoever binds one passes ``--pipeline.ee_frame`` with it.