From 6eba390b037d2db48a43bfdcb633aaea5af14a62 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 26 Aug 2026 19:50:04 +0200 Subject: [PATCH 01/10] Play a chunk in a layer and command the harness one round at a time `Session.__call__` answers `(commands, resume_at_ns)`: the commands to emit now, and the instant the session wants its next call. `ChunkedSchedule` becomes `ChunkPlayer`: it holds the anchored chunk, drops the waypoints already passed, answers the ones due, and asks the sessions below for a new chunk in the call that drains the last one. The harness sheds its waypoint deques, `_play`, `_reschedule` and `_assert_anchored`. It emits what the session answers and paces the next round by `resume_at_ns`, still capped at the poll period. Sessions under the player keep answering chunks, and `base.py` names that contract `ChunkSession` until PR 9 makes the chunk a served function's answer. `openpi` moves `ChangeEEFrame` under the player, where a codec belongs. `episode.steps` counts the control rounds the harness ran. The golden pipeline does not move. --- .basedpyright/baseline.json | 146 ----------- ARCHITECTURE.md | 24 +- docs/connect-your-model.md | 20 +- docs/inference.md | 2 +- positronic/cfg/layers.py | 6 +- positronic/cli/eval/run.py | 8 +- positronic/offboard/README.md | 11 +- positronic/offboard/server.py | 2 +- positronic/offboard/tests/conftest.py | 27 +- .../offboard/tests/test_remote_policy.py | 26 +- positronic/offboard/tests/test_server.py | 76 +++--- positronic/policy/__init__.py | 14 +- positronic/policy/base.py | 118 ++++++--- positronic/policy/codec.py | 13 +- positronic/policy/executor.py | 18 +- positronic/policy/harness.py | 102 +++----- positronic/policy/layers.py | 136 ++++++---- positronic/policy/recording.py | 45 +++- positronic/policy/remote.py | 6 +- positronic/policy/spec.py | 8 +- positronic/policy/tests/test_executor.py | 16 +- .../policy/tests/test_golden_pipeline.py | 30 ++- positronic/policy/tests/test_harness.py | 240 +++++++++--------- positronic/policy/tests/test_layers.py | 238 ++++++++--------- positronic/policy/tests/test_policy_io.py | 17 +- positronic/policy/tests/test_recording.py | 11 +- positronic/probe.py | 100 +++++--- .../env_server/tests/test_remote_env.py | 12 +- .../tests/test_inference_integration.py | 12 +- positronic/tests/testing_coutils.py | 2 +- positronic/vendors/gr00t/server.py | 4 +- positronic/vendors/lerobot/server.py | 4 +- .../vendors/lerobot/tests/test_server.py | 8 +- positronic/vendors/lerobot_0_3_3/policy.py | 4 +- positronic/vendors/lerobot_0_3_3/server.py | 4 +- .../lerobot_0_3_3/tests/test_server.py | 4 +- positronic/vendors/molmoact2/README.md | 2 +- positronic/vendors/molmoact2/server.py | 4 +- positronic/vendors/openpi/server.py | 11 +- 39 files changed, 774 insertions(+), 757 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 043f8f5b3..1eae3e9a2 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5032,16 +5032,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,30 +5065,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": { @@ -5115,46 +5081,6 @@ "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": { @@ -5171,62 +5097,6 @@ "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": { - "startColumn": 11, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 11, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportCallIssue", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 11, - "endColumn": 22, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -5341,22 +5211,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 17, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 19, - "endColumn": 24, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 621a40dab..fc9c7a2c4 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,21 @@ 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. The harness aims at that instant and may call earlier. **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..3ee4a3e02 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 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 it once per timestep with the latest observation and your clock reading, and it answers the actions to run. A policy served behind the `remote` marker answers a chunk, so its sessions are `ChunkSession`s; the rig's `ChunkPlayer` holds that chunk and emits each waypoint at its own time. Per-episode state (history, the chunk in flight) lives in the session — so one `Policy` can serve several robots at once, each with its own session. **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,12 +172,12 @@ 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 positronic.policy import ChunkSession, 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 MySession(ChunkSession): def __init__(self, model): self._model = model @@ -205,20 +205,20 @@ class MyPolicy(Policy): 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. -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 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: ```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. @@ -257,7 +257,7 @@ import time from positronic.offboard.protocol import serialise, deserialise -session = policy.new_session() # one Session per episode/connection +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) 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..d364e0bb9 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -20,7 +20,7 @@ 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.executor import Executor from positronic.policy.harness import Harness from positronic.simulator.env_server.telemetry import ATTR_RUN_ID, ENV_RUN_ID, ENV_TELEMETRY_DIR @@ -194,7 +194,11 @@ def main(policy, *, evals: list[Eval], output_dir: str | Path | None = None, tim # 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() + warmup = Executor(policy.functions) + try: + policy.new_session(None, warmup).close() + finally: + warmup.close() 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..f71818573 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -174,7 +174,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() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 5a371ff75..baee1ee54 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 typing import Any from unittest.mock import MagicMock import pytest import uvicorn from positronic.offboard.server import PolicyServer -from positronic.policy import Policy, Session +from positronic.policy import AnySession, ChunkSession, 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,11 +58,11 @@ async def _run(): @pytest.fixture -def open_session() -> Generator[Callable[..., tuple[Session, Executor]], None, None]: +def open_session() -> Generator[Callable[..., tuple[AnySession, Executor]], None, None]: """Opens a policy's session against a runtime that serves its functions, as the harness does.""" runtimes: list[Executor] = [] - def make(policy: Policy) -> tuple[Session, Executor]: + def make(policy: Policy) -> tuple[AnySession, Executor]: runtimes.append(Executor(policy.functions)) return policy.new_session(None, runtimes[-1]), runtimes[-1] @@ -74,7 +75,7 @@ def make(policy: Policy) -> tuple[Session, Executor]: ANSWER_SEC = 5.0 -def round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> list[dict] | None: +def round_trip(session: ChunkSession, rt: Executor, obs, time_ns: int = 0) -> list[dict] | dict | None: """What ``session`` answers for ``obs``, over the two calls one round trip takes. Both calls get the same ``time_ns``, so a chunk comes back anchored at the value the test passed. @@ -85,6 +86,18 @@ def round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> list[di return session(obs, time_ns) +def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> Mapping[str, Any]: + """What ``session`` commands for ``obs``, over the two calls one round trip takes. + + For a session topped by a ``ChunkPlayer``: both calls get 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) == ({}, time_ns), 'a round-trip was already in flight' + rt.wait(ANSWER_SEC) + assert not rt.in_flight, 'the round-trip never came back' + return session(obs, time_ns)[0] + + def _make_mock_policy(action, meta): """Create a mock policy with session-based API.""" session = MagicMock() @@ -146,7 +159,7 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple 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 @@ -154,5 +167,5 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple 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)) + host, port, _server = start_server(ChunkPlayer() | remote | _DictSource(mock_policy_registry)) return host, port, mock_policy_registry diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index c1a385296..b1a5a5911 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -12,17 +12,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.offboard.tests.conftest import ANSWER_SEC, played_round_trip, round_trip from positronic.policy import 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): @@ -496,7 +496,7 @@ def test_declared_stack_built_at_session_open(open_session): 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(): @@ -531,19 +531,19 @@ 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}') @@ -553,10 +553,10 @@ def test_remote_policy_lifecycle(inference_server, mock_policy, open_session): 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() diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index e5ef71f01..d6e4a290a 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -17,11 +17,11 @@ 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.offboard.tests.conftest import ANSWER_SEC +from positronic.policy import ChunkSession, Codec, Policy, RemotePolicy from positronic.policy.base import Runtime 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,7 +59,7 @@ 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'} @@ -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 @@ -221,12 +221,12 @@ def test_a_backend_that_cannot_answer_its_warmup_raises_and_still_ends_its_sessi 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() @@ -241,7 +241,7 @@ def test_pipeline_with_no_rig_side_half_refused_at_startup(make_mock_policy): _INFER = 'infer' -class _ScriptedSession(Session): +class _ScriptedSession(ChunkSession): def __init__(self, rt: Runtime): self._rt = rt self._answer = None @@ -259,7 +259,7 @@ def __call__(self, obs, time_ns): 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: + def new_session(self, context=None, rt=None) -> ChunkSession: assert rt is not None return _ScriptedSession(rt) @@ -272,33 +272,43 @@ 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, walked by the instants it asks to be called at.""" + at_ns = int(100e9) + assert session({keys.OBS_TIME_NS: 0}, at_ns) == ({}, at_ns), 'a round trip was already in flight' + runtime.wait(ANSWER_SEC) + played = [] + for _ in range(4): + commands, resume_at_ns = session({keys.OBS_TIME_NS: 0}, at_ns) + played.append((at_ns, commands)) + if resume_at_ns <= at_ns: + break + at_ns = resume_at_ns + 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 +345,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 +440,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 +515,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..ee7b6361b 100644 --- a/positronic/policy/__init__.py +++ b/positronic/policy/__init__.py @@ -1,4 +1,13 @@ -from .base import DelegatingPolicy, DelegatingSession, Layer, Policy, Session +from .base import ( + AnySession, + ChunkSession, + DelegatingChunkSession, + DelegatingPolicy, + DelegatingSession, + Layer, + Policy, + Session, +) from .codec import ActionHorizon, ActionTimestamp, ActionTiming, Codec, is_action from .recording import Recorder from .remote import RemotePolicy @@ -6,8 +15,11 @@ __all__ = [ 'Policy', 'Session', + 'ChunkSession', + 'AnySession', 'DelegatingPolicy', 'DelegatingSession', + 'DelegatingChunkSession', 'Layer', 'RemotePolicy', 'Codec', diff --git a/positronic/policy/base.py b/positronic/policy/base.py index db88af250..a1f7ce69c 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -43,36 +43,13 @@ def fns(self) -> Mapping[str, Fn]: """The policy's functions, under the names it declared them by.""" -class Session(ABC): - """Per-episode inference session. Created by ``Policy.new_session()``. +class _SessionBase: + """Per-episode state a policy opens: what the framework holds, cancels and closes. - 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:: - - session = policy.new_session(context) - trajectory = session(obs, time_ns) - - **Plain-data contract**: sessions accept and return only plain data - (dicts, lists, numpy arrays, scalars). No tensors or custom objects. - - **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. + ``Session`` and ``ChunkSession`` add the call. Both are opened and ended the same way, so the + lifecycle lives here. """ - @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``. - - ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. - """ - @property def meta(self) -> dict[str, Any]: """Session metadata (may include policy meta + per-session info).""" @@ -80,7 +57,7 @@ def meta(self) -> dict[str, Any]: def cancel(self): """Drop any in-flight trajectory state. Layers that buffer/schedule a - trajectory (e.g. ``ChunkedSchedule``) should reset so the next call + trajectory (e.g. ``ChunkPlayer``) should reset so the next call triggers a fresh inference. Override and propagate via ``super().cancel()``. """ return None @@ -90,8 +67,57 @@ def close(self): return None +class Session(_SessionBase, ABC): + """Per-episode control session. Created by ``Policy.new_session()``. + + Sessions hold per-episode state (the chunk being played, latency tracking) and are what a caller + controls a robot through. Call the session like a function:: + + session = policy.new_session(context) + commands, resume_at_ns = session(obs, time_ns) + + **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) -> 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.functions``. + + ``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 on the clock of ``time_ns``. The caller aims at it and may call earlier. + """ + + +class ChunkSession(_SessionBase, ABC): + """Session that answers action chunks: what a ``ChunkPlayer`` holds and plays. + + A call answers a list of actions, each naming command channels and carrying its own + ``keys.ACTION_TIMESTAMP`` in seconds from the call. ``None`` means "no new chunk" — what a session + answers while the work it asked for is still in flight. An empty list drops what is playing. One + action may come back bare, and the player wraps it. + + TODO(#661): the chunk becomes the answer of a served function, and this class goes with the + session-serving wire. + """ + + @abstractmethod + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | dict[str, Any] | None: + """The chunk to play from now, or ``None`` while the work is in flight. + + ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. + """ + + +# What a layer wraps: a session on either side of a ``ChunkPlayer``. +# TODO(#661): the chunk side goes, and this alias with it. +AnySession = Session | ChunkSession + + 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 @@ -110,16 +136,40 @@ def close(self): self._inner.close() +class DelegatingChunkSession(ChunkSession): + """Chunk session that delegates all methods to an inner ``ChunkSession``. + + The body repeats ``DelegatingSession`` because the two answer different calls, and each one's ``_inner`` + is the side it forwards to. TODO(#661): it goes with ``ChunkSession``. + """ + + def __init__(self, inner: ChunkSession): + self._inner = inner + + 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() + + def close(self): + self._inner.close() + + class Policy(ABC): """Factory for inference sessions. A Policy holds shared resources (model weights, connections) and creates - per-episode ``Session`` instances. One Policy can serve multiple robots + per-episode session instances. One Policy can serve multiple robots by creating independent sessions. """ @abstractmethod - def new_session(self, context: dict[str, Any] | None = None, rt: Runtime | None = None) -> Session: + def new_session(self, context: dict[str, Any] | None = None, rt: Runtime | None = None) -> AnySession: """Create a new inference session for an episode. Args: @@ -174,7 +224,7 @@ 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 @@ -186,8 +236,8 @@ def wrap(self, policy: Policy) -> Policy: """Apply this layer to a policy. Default: wrap every session it creates via ``make_session``.""" return _LayerPolicy(policy, self) - def make_session(self, inner: Session) -> Session: - """Make this layer's session around ``inner``.""" + def make_session(self, inner: Any) -> AnySession: + """Make this layer's session around ``inner``, on the side of a ``ChunkPlayer`` the layer sits on.""" raise NotImplementedError('Override make_session or wrap') # The name this layer travels under, set by every deliverable subclass. ``WIRE_LAYERS`` is keyed by diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index c701dc24c..844f31173 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -24,7 +24,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, ChunkSession, DelegatingChunkSession, Layer, _ComposedLayer from positronic.utils import merge_dicts _QUAT = geom.Rotation.Representation.QUAT @@ -81,7 +81,7 @@ def training_encoder(self) -> EpisodeTransform: def meta(self) -> dict: return {} - def make_session(self, inner: Session): + def make_session(self, inner: ChunkSession) -> ChunkSession: return _CodecSession(inner, self) @final @@ -100,16 +100,19 @@ def __and__(self, other): return NotImplemented -class _CodecSession(DelegatingSession): +class _CodecSession(DelegatingChunkSession): """Session wrapped with a codec: encodes observations, decodes actions.""" - def __init__(self, inner: Session, codec: 'Codec'): + def __init__(self, inner: ChunkSession, 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) + # A codec decodes a chunk, so it belongs under the player. Above one it would be handed commands + # paired with a resume time, and each codec would drop or corrupt them in its own way. + assert not isinstance(action, tuple), 'compose a Codec under the ChunkPlayer, not above it' if action is None: return None return self._codec.decode(action) @@ -447,7 +450,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..08f014a9d 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -7,12 +7,13 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor from functools import partial -from typing import Any +from typing import Any, cast from positronic.policy.base import ( Answer, + ChunkSession, + DelegatingChunkSession, DelegatingPolicy, - DelegatingSession, Fn, NotAnswered, Policy, @@ -115,12 +116,12 @@ def close(self) -> None: class _BlockingPolicy(DelegatingPolicy): - class _Session(DelegatingSession): - def __init__(self, inner: Session, rt: Executor): + class _Session(DelegatingChunkSession): + def __init__(self, inner: ChunkSession, rt: Executor): super().__init__(inner) self._rt = rt - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | None: + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | 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: @@ -132,11 +133,14 @@ def close(self): self._rt.close() self._inner.close() - def new_session(self, context=None, rt=None) -> Session: + def new_session(self, context=None, rt=None) -> ChunkSession: 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) + inner = self._inner.new_session(context, own) + # Waiting out the work is a question about a chunk, so nothing above a ``ChunkPlayer`` blocks. + assert not isinstance(inner, Session), f'{type(inner).__name__} answers commands, not a chunk' + return _BlockingPolicy._Session(cast(ChunkSession, inner), own) except BaseException: own.close() raise diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 9aabb524e..8d2c5a90f 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,6 +1,5 @@ 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 @@ -12,16 +11,12 @@ from positronic.dataset.serializers import expand_suffixed from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Task -from positronic.policy.base import Policy +from positronic.policy.base import Policy, Session 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. +# How long a real-time round may last when the session asks for nothing 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 @@ -35,7 +30,12 @@ def __init__(self, policy: Policy, context: dict[str, Any], charges_wall_time: b self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() self._runtime = Executor(policy.functions) try: - self._session = policy.new_session(context, self._runtime) + session = policy.new_session(context, self._runtime) + # A stack that answers chunks has no player in it, so nothing turns them into commands. + assert isinstance(session, Session), ( + f'{type(session).__name__} answers a chunk; a policy played here needs a ChunkPlayer in its stack' + ) + self._session = session except BaseException: self._runtime.close() raise @@ -53,7 +53,7 @@ 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: @@ -109,6 +109,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 +145,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 +170,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) @@ -239,20 +240,21 @@ def _ready(self, should_stop: pimm.SignalReceiver, args: dict[str, Any]) -> Gene 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.""" + each observation instantly. Real: sleep to the moment the session asked for, capped at the poll + period, so a waypoint is emitted at its own time and a round rarely finds more than one due. A + session that names a moment already reached asks for the next round, one poll period away. + """ 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)) + due_in_sec = (self._resume_at_ns - clock.now_ns()) / 1e9 if self._resume_at_ns is not None else 0.0 + return pimm.Sleep(min(POLL_PERIOD_SEC, due_in_sec) if due_in_sec > 0 else POLL_PERIOD_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 +264,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 @@ -376,49 +377,11 @@ 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 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 +435,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..a420f26af 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`` +converts a chunk into the commands a caller runs; the layers around 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 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 ChunkSession, DelegatingSession, Layer, 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,10 +42,10 @@ 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' @@ -50,58 +55,97 @@ def __call__(self, obs, time_ns): if _arms_available(obs): return self._inner(obs, time_ns) self.cancel() - return [] + return {}, time_ns - 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 the sessions below answer, 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. + Owns relative→absolute time conversion: the sessions below (codecs, models) emit timestamps relative to + the call; this layer anchors them to the call's ``time_ns``. Each call answers the waypoints that have + come due, and asks to be called again at the next one. The call that drains the last waypoint asks the + sessions below for a new chunk, and plays whatever of it is already due. """ - WIRE_NAME = 'chunked_schedule' + WIRE_NAME = 'chunk_player' - class _Session(DelegatingSession): - """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" + class _Session(Session): + """Plays the chunk it holds; re-queries in the call that drains it. - def __init__(self, inner: Session): - super().__init__(inner) - self._trajectory_end: float | None = None + The one session with a different contract on each side, so it forwards rather than delegates. + """ + + def __init__(self, inner: ChunkSession): + self._inner = inner + self._waypoints: deque[tuple[int, dict[str, Any]]] = deque() 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 + commands = self._due(time_ns) + if not self._waypoints: + chunk = self._inner(obs, time_ns) + if chunk is not None: + self._load(chunk, time_ns) + commands = {**commands, **self._due(time_ns)} + # A spent chunk asks for the next call at once: the work for the chunk after it is in flight. + return commands, self._waypoints[0][0] if self._waypoints else time_ns + + def _due(self, time_ns: int) -> dict[str, Any]: + """The waypoints reached at ``time_ns``, taken off the chunk. + + A channel with several waypoints due keeps the last: exact for an absolute setpoint, lossy for + a relative one. + """ + commands: dict[str, Any] = {} + while self._waypoints and self._waypoints[0][0] <= time_ns: + commands.update(self._waypoints.popleft()[1]) + return commands + + 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 session 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] + anchor = time_ns / 1e9 + 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 sessions below are timing actions against a clock of their own' + ) + # The single explicit seconds->ns seam: the sessions below time actions in float seconds, and + # every pimm channel is in ns. A waypoint naming no channel — the codecs' end-of-chunk sentinel — + # commands nothing and states where the chunk ends. + self._waypoints = deque( + ( + int((anchor + action.get(keys.ACTION_TIMESTAMP, 0.0)) * 1e9), + {name: value for name, value in action.items() if name != keys.ACTION_TIMESTAMP}, + ) + for action in chunk + ) + + @property + def meta(self): + return self._inner.meta def cancel(self): - self._trajectory_end = None - super().cancel() + self._waypoints.clear() + self._inner.cancel() + + def close(self): + self._inner.close() - def make_session(self, inner: Session): - return ChunkedSchedule._Session(inner) + def make_session(self, inner: ChunkSession) -> Session: + return ChunkPlayer._Session(inner) def to_spec(self): return {'name': self.WIRE_NAME} @@ -156,7 +200,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 +240,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..30ca1c5aa 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -60,7 +60,7 @@ 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 AnySession, Layer, Session from positronic.policy.codec import is_action from positronic.utils.rerun_compat import log_numeric_series, set_timeline_sequence, set_timeline_time @@ -361,16 +361,20 @@ def __init__(self, rec: Recorder, name: str): self._rec = rec self._name = name - def make_session(self, inner: Session) -> Session: + def make_session(self, inner: AnySession) -> AnySession: stream = self._rec._open_stream() return _RecordingTapSession(inner, self._rec, self._name, stream) -class _RecordingTapSession(DelegatingSession): - """Logs the observation flowing down and the action chunk flowing up at one point.""" +class _RecordingTapSession(Session): + """Logs the observation flowing down and the actions flowing up at one point. - def __init__(self, inner: Session, rec: Recorder, name: str, stream: rr.RecordingStream): - super().__init__(inner) + A tap sits on either side of a ``ChunkPlayer``, so it forwards whatever its inner answers and reads the + side off that. TODO(#661): the framework records each boundary it carries, and a tap stops guessing. + """ + + def __init__(self, inner: AnySession, rec: Recorder, name: str, stream: rr.RecordingStream): + self._inner = inner self._rec = rec self._name = name self._stream = stream @@ -385,7 +389,7 @@ def _set_timelines(self) -> None: set_timeline_time(timeline, value) set_timeline_sequence('step', self._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,7 +404,7 @@ 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], obs: Mapping[str, Any]) -> 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)] @@ -450,6 +454,18 @@ def _log_action_chunk(self, prefix: str, actions: list[dict], obs: dict) -> None _log_action_series(f'{prefix}/series/{suffix}', field_arr, group_h, base_ns, names=None) self._rec._series_paths.append(f'{prefix}/series/{suffix}') + @staticmethod + def _actions(answer) -> list[dict]: + """The actions in what a session answered: a chunk under a ``ChunkPlayer``, one command above one. + + Above the player the answer holds the commands of one round, so a plot reads a point per round + rather than a chunk per inference. + """ + if isinstance(answer, tuple): + commands, _ = answer + return [dict(commands)] if commands else [] + return answer or [] + def _send_blueprint(self) -> None: rec = self._rec bp = rec._blueprint or _build_blueprint( @@ -458,7 +474,7 @@ def _send_blueprint(self) -> None: if bp is not None: rr.send_blueprint(bp) - def __call__(self, obs, time_ns): + def __call__(self, obs, time_ns) -> Any: rec = self._rec outermost = rec._depth == 0 if outermost: @@ -469,9 +485,9 @@ def __call__(self, obs, time_ns): self._set_timelines() self._log(self._name, obs) - actions = self._inner(obs, time_ns) + answer = self._inner(obs, time_ns) - if actions is not None: + if actions := self._actions(answer): with self._stream: self._set_timelines() # TODO(#661): the chunk is logged against the observation and the timelines of this @@ -484,12 +500,15 @@ def __call__(self, obs, time_ns): with self._stream: self._send_blueprint() self._step += 1 - return actions + return answer finally: rec._depth -= 1 if rec._depth == 0: rec._timeline_values = {} + def cancel(self): + self._inner.cancel() + def close(self): - super().close() + self._inner.close() self._rec._release_stream() diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index a051fded6..8a20d39e6 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -10,7 +10,7 @@ from positronic.utils import flatten_dict from positronic.utils.serialization import encode_jpeg -from .base import Answer, Layer, Policy, Runtime, Session +from .base import Answer, AnySession, ChunkSession, Layer, Policy, Runtime from .recording import Recorder from .spec import from_spec @@ -53,7 +53,7 @@ def round_trip( telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns()) -class RemoteSession(Session): +class RemoteSession(ChunkSession): """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 @@ -197,7 +197,7 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, rt=None) -> Session: + def new_session(self, context=None, rt=None) -> AnySession: return self._policy().new_session(context, rt) @property 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/test_executor.py b/positronic/policy/tests/test_executor.py index 7e5ff780b..378c97ad3 100644 --- a/positronic/policy/tests/test_executor.py +++ b/positronic/policy/tests/test_executor.py @@ -10,7 +10,7 @@ import pytest -from positronic.policy.base import Answer, DelegatingSession, Layer, NotAnswered, Policy, Session +from positronic.policy.base import Answer, ChunkSession, DelegatingChunkSession, Layer, NotAnswered, Policy from positronic.policy.executor import Executor, blocking # How long a test waits for the worker threads before calling the call lost. @@ -237,7 +237,7 @@ def test_close_stays_quiet_about_a_call_that_answered(serve, caplog): class _PlainPolicy(Policy): """Serves nothing: its session answers inside the call that asked.""" - class _Session(Session): + class _Session(ChunkSession): def __init__(self): self.calls = 0 @@ -248,7 +248,7 @@ def __call__(self, obs, time_ns): def __init__(self): self.session = _PlainPolicy._Session() - def new_session(self, context=None, rt=None) -> Session: + def new_session(self, context=None, rt=None) -> ChunkSession: return self.session @@ -258,7 +258,7 @@ def new_session(self, context=None, rt=None) -> Session: class _EchoPolicy(Policy): """Serves ``echo``, and makes sessions that take ``rounds`` calls of it to answer.""" - class _Session(Session): + class _Session(ChunkSession): def __init__(self, rt, rounds: int): self._rt = rt self._answer = None @@ -280,7 +280,7 @@ def __init__(self, rounds: int): self._rounds = rounds self.session: _EchoPolicy._Session - def new_session(self, context=None, rt=None) -> Session: + def new_session(self, context=None, rt=None) -> ChunkSession: assert rt is not None self.session = _EchoPolicy._Session(rt, self._rounds) return self.session @@ -296,8 +296,8 @@ class _CountingLayer(Layer): def __init__(self): self.calls = 0 - class _Session(DelegatingSession): - def __init__(self, inner: Session, layer: '_CountingLayer'): + class _Session(DelegatingChunkSession): + def __init__(self, inner: ChunkSession, layer: '_CountingLayer'): super().__init__(inner) self._layer = layer @@ -314,7 +314,7 @@ def opened(): """Opens the sessions a test asks for, and closes every one at teardown.""" sessions = [] - def make(policy: Policy) -> Session: + def make(policy: Policy) -> ChunkSession: sessions.append(policy.new_session()) return sessions[-1] diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index 1c4a36cfb..e2a754e5e 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -41,10 +41,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 ChunkSession, DelegatingChunkSession, DelegatingPolicy, Policy 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,7 +64,7 @@ CAPTURED_SIGNALS = (keys.EE_POSE, keys.JOINTS, keys.GRIP) -class _ScriptedSession(Session): +class _ScriptedSession(ChunkSession): def __call__(self, obs, time_ns): current = np.asarray(obs[keys.EE_POSE][:3], dtype=np.float32) delta = TARGET_POS - current @@ -80,7 +80,7 @@ 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. Codec stamps/truncates; ``ChunkedSchedule`` anchors and the harness plays. + clock, no images. The codec stamps and truncates; ``ChunkPlayer`` anchors the chunk and plays it. """ def new_session(self, context=None, rt=None): @@ -88,19 +88,19 @@ def new_session(self, context=None, rt=None): 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.""" + """A fixed inference latency in world time: every chunk is stamped for, and reaches the player at, + ``latency_sec`` after the call that produced it, whatever the machine took. Sits under the player, so + nothing below it sees an observation while a chunk is held. A skip or a stop passes through at once.""" def __init__(self, inner: Policy, latency_sec: float): super().__init__(inner) self._latency_ns = round(latency_sec * 1e9) - class _Session(DelegatingSession): - def __init__(self, inner: Session, latency_ns: int): + class _Session(DelegatingChunkSession): + def __init__(self, inner: ChunkSession, latency_ns: int): super().__init__(inner) self._latency_ns = latency_ns - self._held: list | None = None + self._held: list | dict | None = None self._release_at_ns = 0 def __call__(self, obs, time_ns): @@ -109,7 +109,7 @@ def __call__(self, obs, time_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. + # The chunk is held for ``latency_ns``, so the sessions below time it from the release instant. result = self._inner(obs, time_ns + self._latency_ns) if not result: return result @@ -120,8 +120,10 @@ 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) + def new_session(self, context=None, rt=None) -> ChunkSession: + inner = self._inner.new_session(context, rt) + assert isinstance(inner, ChunkSession) + return _SimulatedLatency._Session(inner, self._latency_ns) class FakeRobot(pimm.ControlSystem): @@ -208,7 +210,7 @@ def _run_pipeline(tmp_path: Path) -> dict: simulated=True, ) harness = Harness( - _SimulatedLatency((StopOnFault() | ChunkedSchedule()).wrap(policy), INFERENCE_LATENCY_S), embodiment + (StopOnFault() | ChunkPlayer()).wrap(_SimulatedLatency(policy, INFERENCE_LATENCY_S)), 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..94ad27065 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,10 +22,10 @@ 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 ChunkSession, DelegatingSession, 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.layers import ChunkPlayer, StopOnFault from positronic.policy.remote import INFER, RemoteSession, round_trip from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -79,7 +80,7 @@ def make_embodiment( ) -class _SpySession(Session): +class _SpySession(ChunkSession): def __init__(self, policy): self._policy = policy @@ -105,7 +106,7 @@ def new_session(self, context=None, rt=None): return _SpySession(self) -class _StubSession(Session): +class _StubSession(ChunkSession): def __init__(self, policy): self._policy = policy self._meta = dict(policy._meta) @@ -144,13 +145,13 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, rt=None) -> Session: + def new_session(self, context=None, rt=None) -> ChunkSession: self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) -class _ChunkSession(Session): +class _ChunkSession(ChunkSession): def __init__(self, policy): self._policy = policy @@ -356,7 +357,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 +418,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 +448,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 +475,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 +517,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 +564,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,7 +593,7 @@ 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): + class _StaticMetaSession(ChunkSession): def __init__(self, command): self._command = command @@ -611,7 +612,7 @@ def new_session(self, context=None, rt=None): 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 +634,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 +655,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 +672,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,7 +690,7 @@ def test_an_uncharged_wait_ends_when_the_world_comes_down(world): never_answers = threading.Event() class _HangingPolicy(Policy): - class _Session(Session): + class _Session(ChunkSession): def __init__(self, rt): self._rt = rt @@ -705,7 +706,7 @@ def new_session(self, context=None, rt=None): 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,7 +720,7 @@ 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): + class _RaisingSession(ChunkSession): def __call__(self, obs, time_ns): raise RuntimeError('inference boom') @@ -727,7 +728,7 @@ class _RaisingPolicy(Policy): def new_session(self, context=None, rt=None): return _RaisingSession() - 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 +749,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 +767,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 +784,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 +810,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 +881,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 +925,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 +946,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 +972,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 +994,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 +1030,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 +1047,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 @@ -1086,7 +1087,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 +1105,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 +1126,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 +1144,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 +1163,15 @@ 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): + class _BoomSession(ChunkSession): + def __call__(self, obs, time_ns): raise RuntimeError('inference boom') class _BoomPolicy(Policy): - def new_session(self, context=None, now=None, rt=None): + def new_session(self, context=None, rt=None): return _BoomSession() - 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 +1193,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 +1222,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 +1253,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 +1288,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 +1310,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 +1346,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)) @@ -1399,7 +1400,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 +1424,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 +1450,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 +1479,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,7 +1513,7 @@ 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): + class _EmptyChunkSession(ChunkSession): def __call__(self, obs, time_ns): return [] @@ -1520,7 +1521,7 @@ class EmptyChunkPolicy(Policy): def new_session(self, context=None, rt=None): return _EmptyChunkSession() - 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 +1549,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 +1576,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]) @@ -1606,7 +1607,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 +1678,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 +1733,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. @@ -1776,7 +1777,7 @@ def test_timing_spans_recorded_with_taxonomy(world, tmp_path): 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()) + 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 +1811,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 +1824,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 +1858,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 +1900,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,61 +1931,56 @@ 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) - - -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) - - -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) - - -@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.""" +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()) - 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) + 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(AssertionError, match='needs a ChunkPlayer'): + 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): + def __init__(self, inner: ChunkSession): + self._inner = inner + self._waypoints: deque[tuple[int, dict[str, Any]]] = deque() + self._replan_at_ns: int | 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] + return commands, min(due, default=time_ns) + + def _replan(self, obs, time_ns: int) -> None: + chunk = self._inner(obs, time_ns) + if not isinstance(chunk, list): # the function it asked has still to answer + return 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 - - def make_session(self, inner: Session): + 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: ChunkSession): return _ReplanEarly._Session(inner) @@ -2035,7 +2033,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 +2044,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 +2055,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' @@ -2069,7 +2067,7 @@ class _ObservedTicks(Layer): def __init__(self): self.seen: list[float] = [] - class _Session(DelegatingSession): + class _Session(DelegatingSession, Session): def __init__(self, inner: Session, seen: list[float]): super().__init__(inner) self._seen = seen @@ -2093,7 +2091,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 +2121,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 +2165,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 +2195,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 +2237,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 +2279,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()) @@ -2326,7 +2324,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,7 +2352,7 @@ 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 _GripThenArm(ChunkSession): """First a two-channel chunk, then an arm-only one that must silence the gripper.""" def __init__(self): @@ -2375,7 +2373,7 @@ 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 +2399,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 +2421,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..bc7c9d7ab 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -1,4 +1,4 @@ -"""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 typing import Any @@ -11,7 +11,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 ChunkSession, Layer, Policy, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -23,11 +23,11 @@ RestrictImageSize, SetControlMode, ) -from positronic.policy.layers import ChunkedSchedule, StopOnFault, TemporalStack +from positronic.policy.layers import ChunkPlayer, StopOnFault, TemporalStack from positronic.policy.observation import ObservationCodec -class _ConstSession(Session): +class _ConstSession(ChunkSession): def __init__(self, actions): self._actions = actions self.call_count = 0 @@ -40,9 +40,9 @@ def __call__(self, obs, time_ns): class _ConstPolicy(Policy): def __init__(self, actions): self._actions = actions - self._session: _ConstSession | None = None + self._session = _ConstSession(actions) - def new_session(self, context=None, rt=None): + def new_session(self, context=None, rt=None) -> _ConstSession: self._session = _ConstSession(self._actions) return self._session @@ -51,34 +51,46 @@ 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 at once, as a layer above a player does.""" + + 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 + + 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) == ({}, 0) 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}, 0) 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}, 0) 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,46 +98,47 @@ 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) == ({}, 0) 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) == ({}, 0) 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}, 0) 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.""" + """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 = (StopOnFault() | ChunkPlayer()).wrap(inner).new_session() - 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)) + assert session(_obs(0.3), int(0.3e9)) == ({'v': 1}, int(1.3e9)) + assert inner._session.call_count == 2 -class _ScriptedSession(Session): +class _ScriptedSession(ChunkSession): """Answers each of ``script`` in turn — ``None`` where a session has nothing to place yet.""" def __init__(self, script): @@ -137,99 +150,95 @@ def __call__(self, obs, time_ns): return self._script.pop(0) -class TestChunkedSchedule: +class TestChunkPlayer: 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.""" + """A session that waits for a served function answers ``None``, which is no chunk. The player + commands nothing, asks for the next call at once, and asks the session again on it.""" inner = _ScriptedSession([None, None, [{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]]) - session = ChunkedSchedule().make_session(inner) + session = ChunkPlayer().make_session(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 session(_obs(0.0), int(1e9)) == ({}, int(1e9)) + assert session(_obs(0.1), int(1e9)) == ({}, int(1e9)) + assert session(_obs(0.2), int(1e9)) == ({'v': 1}, int(1e9)) assert inner.call_count == 3 - 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. + def test_a_chunk_plays_one_waypoint_at_a_time(self): + """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}]) - 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 + session = ChunkPlayer().wrap(inner).new_session() - 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 session(_obs(), int(1e9)) == ({'v': 1}, int(1.5e9)) + assert session(_obs(), int(1.2e9)) == ({}, int(1.5e9)) + assert inner._session.call_count == 1 + + def test_a_channel_with_several_waypoints_due_keeps_the_last(self): + """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 = ChunkPlayer().wrap(inner).new_session() + + 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): + """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 = ChunkPlayer().wrap(inner).new_session() + + assert session(_obs(), int(1e9)) == ({'v': 1}, int(1.5e9)) + assert session(_obs(), int(1.5e9)) == ({'v': 1}, int(2.0e9)) 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_a_single_action_asks_for_the_next_chunk_at_once(self): + """A chunk of one action at ts=0 is drained by the call that loads it, so the next call re-queries.""" + session = ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])).new_session() - 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.""" - 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 + assert session(_obs(1.0), int(1e9)) == ({'v': 1}, int(1e9)) + assert session(_obs(1.01), int(1.01e9)) == ({'v': 1}, int(1.01e9)) + + def test_a_waypoint_naming_no_channel_commands_nothing(self): + """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 = ChunkPlayer().make_session(inner.new_session()) + + 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): + """A chunk that reaches the player already anchored would place its waypoints decades out.""" + session = ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 1.77e18 / 1e9}])).new_session() + + with pytest.raises(ValueError, match='clock of their 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() + 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 + assert session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}, int(1e9)) == ({'v': 1}, int(1e9)) def test_codec_pipe_layer(self): - codec = ActionTimestamp(fps=10.0) - pipeline = codec | ChunkedSchedule() + """A codec composes above a player, and refuses the commands it would silently corrupt.""" + 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 + session = pipeline.wrap(_ConstPolicy([{'action': 'test'}])).new_session() + with pytest.raises(AssertionError, match='under the ChunkPlayer'): + session(_obs(), int(1e9)) def test_full_pipeline(self): 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). + # 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. 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 + 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,7 +276,7 @@ def test_parallel_frame_codecs_keep_the_frame_they_share(self): ) -class _CaptureSession(Session): +class _CaptureSession(ChunkSession): def __init__(self): self.seen = [] @@ -371,7 +380,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 +389,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 +404,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 +413,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 +421,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 +449,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 +470,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 +507,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 +527,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 +536,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 +546,27 @@ 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): 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 + 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): inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) - policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) + policy = spec.inline(ChunkPlayer() | 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 + assert session(_obs(), int(1e9)) == ({'v': 1}, int(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..2f4822c62 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -8,7 +8,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 ChunkSession, Policy from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -115,7 +115,7 @@ def test_absolute_joints_action_encode_decode(): assert np.isclose(target_grip, g[0]) -class _FixedSession(Session): +class _FixedSession(ChunkSession): def __init__(self, result): self._result = result @@ -127,12 +127,12 @@ class _ChunkPolicy(Policy): def __init__(self, actions: list[dict]): self._actions = actions - def new_session(self, context=None, rt=None): + def new_session(self, context=None, rt=None) -> ChunkSession: return _FixedSession(list(self._actions)) class _SinglePolicy(Policy): - def new_session(self, context=None, rt=None): + def new_session(self, context=None, rt=None) -> ChunkSession: return _FixedSession({'v': 42}) @@ -146,7 +146,7 @@ def encode(self, data): class _MetaPolicy(Policy): - def new_session(self, context=None, rt=None): + def new_session(self, context=None, rt=None) -> ChunkSession: return _FixedSession({}) @property @@ -163,6 +163,7 @@ def test_action_horizon_sec_truncates_chunk(): codec = ActionTiming(fps=30.0, horizon_sec=0.1) policy = codec.wrap(_ChunkPolicy(actions)) result = policy.new_session()(_T0_OBS, 0) + 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 @@ -189,7 +190,7 @@ def test_timestamps_embedded_in_actions(): 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 + 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) @@ -200,7 +201,7 @@ def test_action_horizon_sec_seconds_truncates(): 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 + 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) @@ -251,7 +252,7 @@ def test_action_timestamp_and_horizon_compose(): 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 + 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 diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index 63df5b1fb..df4c04148 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -4,18 +4,19 @@ 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 ChunkSession, Policy from positronic.policy.recording import ( Recorder, _build_blueprint, _command_field_arrays, _flat_wire, + _RecordingTapSession, _squeeze_batch, _stack_numeric, ) -class _TrackingSession(Session): +class _TrackingSession(ChunkSession): def __init__(self, actions, meta): self._actions = actions self._meta = meta @@ -44,7 +45,7 @@ def meta(self): return {'policy_key': 'policy_value'} -class _CapturingSession(Session): +class _CapturingSession(ChunkSession): """Innermost session that snapshots the Recorder's carried timeline state when called.""" def __init__(self, rec, actions): @@ -207,11 +208,12 @@ def test_logs_command_chunk_without_mutating(tmp_path): ] session = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(actions)).new_session() result = session({'x': 1.0, keys.WALL_TIME_NS: 1}, 0) + 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): + class _NoneSession(ChunkSession): def __call__(self, obs, time_ns): return None @@ -220,6 +222,7 @@ def new_session(self, context=None, rt=None): return _NoneSession() session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() + assert isinstance(session, _RecordingTapSession) assert session({'x': 1.0}, 0) is None assert session._step == 1 diff --git a/positronic/probe.py b/positronic/probe.py index 8e161ee94..ab3cb891d 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,13 +33,16 @@ 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' # Observation keys the endpoint expects, mirroring the inference harness, plus every image.*. _STATE_KEYS = (keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP) +# How long to wait for the endpoint, and how often to ask the session again while it is out. +_ANSWER_TIMEOUT_SEC = 120.0 +_POLL_PERIOD_SEC = 0.01 def _build_wire_obs(sample: dict, task: str | None, now_ns: int, recorded_ts: int) -> dict: @@ -68,35 +71,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) -> list[tuple[int, dict]]: + """Every command the session emits for ``obs``, from the endpoint's answer to the end of that chunk. - -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. - - 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. + The observation is one frozen frame, so the walk moves the clock rather than the world. The session + anchors the chunk on the call that receives it and asks for a call at each waypoint; a resume time that + has not moved past the call says it holds nothing — the endpoint is still out, or the chunk has run out. + """ + played: list[tuple[int, dict]] = [] + deadline = time.monotonic() + _ANSWER_TIMEOUT_SEC + now_ns = time.time_ns() + while True: + commands, resume_at_ns = session(obs, now_ns) + if commands: + played.append((now_ns, dict(commands))) + if resume_at_ns > now_ns: + now_ns = resume_at_ns + elif played: + return played + elif time.monotonic() >= deadline: + raise TimeoutError(f'the endpoint commanded nothing in {_ANSWER_TIMEOUT_SEC:.0f}s') + else: + time.sleep(_POLL_PERIOD_SEC) + now_ns = time.time_ns() + + +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) + + +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] + 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 +145,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,25 +166,24 @@ 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) + tapped = rec.tap(_TAP).wrap(policy) + runtime = Executor(policy.functions) + session = tapped.new_session({keys.TASK: task} if task else None, runtime) 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}') + played = _play(session, obs) + print(f'episode {episode} @ {at:.3f}s (ts={ts}) [{name}]: {len(played)} command(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) + _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(actions))) + rr.send_blueprint(_blueprint(image_keys, _is_cartesian_chunk(played))) finally: + # The runtime closes first: the call the last waypoint started is still using the session's socket. + runtime.close() session.close() diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index de08094fd..1172ef2ec 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 ChunkSession, 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), ) @@ -373,7 +373,7 @@ def new_session(self, context=None, rt=None): return _JointposChunkSession(self) -class _JointposChunkSession(Session): +class _JointposChunkSession(ChunkSession): def __init__(self, policy: _JointposChunks): self._policy = policy @@ -389,7 +389,7 @@ 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 + sentinel gives a full period before ``ChunkPlayer`` re-infers — and replans arrive exactly ``chunk_len`` control periods apart.""" host, port = env_server probe = make_mujoco_env([]) @@ -398,7 +398,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] 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..d92e7b417 100644 --- a/positronic/tests/testing_coutils.py +++ b/positronic/tests/testing_coutils.py @@ -128,7 +128,7 @@ def __init__(self, policy): def __call__(self, obs, time_ns): self._policy.observations.append(obs) - return [] + return {}, time_ns @property def meta(self): diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index c8a24932f..cd63a8af4 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -19,7 +19,7 @@ from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup from positronic.policy import Policy, Session 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 @@ -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/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..bef059a35 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -16,7 +16,7 @@ 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.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 @@ -161,4 +161,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/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..6b4837630 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -16,7 +16,7 @@ from positronic.offboard.server_utils import run_with_progress, wait_for_subprocess_ready, warmup from positronic.policy import Codec, Policy, Session 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 @@ -258,11 +258,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. From 61292fdc3dc42f09412b6b5e99d8b4e8b4baccde Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 27 Aug 2026 13:38:09 +0200 Subject: [PATCH 02/10] Put every model session on the chunk contract, and fix what a review found A model session answers a chunk, so the six vendor sessions move to `ChunkSession`. They stayed on `Session`, which now means "answers commands", and `blocking()`'s guard refused them: no inference server could load a checkpoint or open a websocket. `basedpyright` excludes `positronic/vendors`, so nothing caught it. A recording tap picks its side when it is made, from the contract its inner answers, so a tap no longer claims to answer commands while it forwards a chunk. It also wraps a bare-dict chunk, which `ChunkSession` allows. `ChunkPlayer` converts the offset rather than the sum, so a waypoint at 0.0 lands on the call that loads it whatever the clock reads. `probe` asks the endpoint once and waits the round trip out, instead of calling the session every 10 ms and logging the observation again each time. It closes the runtime before it reports, so the wait is not a surprise. An episode that refuses to open closes the session it opened. The eval warmup closes its runtime before the session it served. Tests: waypoints due in one round keep every channel they name; the drained chunk's last command survives the call that loads the next; the harness sleeps to the moment the session asked for; a tap above the player logs a command per round. Each one fails against the code it pins. --- .basedpyright/baseline.json | 16 ------- positronic/cli/eval/run.py | 4 +- positronic/policy/executor.py | 4 +- positronic/policy/harness.py | 5 ++- positronic/policy/layers.py | 8 ++-- positronic/policy/recording.py | 42 ++++++++++++------ positronic/policy/remote.py | 2 +- positronic/policy/tests/test_harness.py | 19 +++++++- positronic/policy/tests/test_layers.py | 25 +++++++++++ positronic/policy/tests/test_recording.py | 14 ++++++ positronic/probe.py | 50 +++++++++++----------- positronic/vendors/dreamzero/server.py | 4 +- positronic/vendors/gr00t/server.py | 4 +- positronic/vendors/lerobot/policy.py | 4 +- positronic/vendors/lerobot_0_3_3/policy.py | 4 +- positronic/vendors/molmoact2/policy.py | 6 +-- positronic/vendors/openpi/server.py | 4 +- 17 files changed, 139 insertions(+), 76 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 1eae3e9a2..7d7089635 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5276,22 +5276,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/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index d364e0bb9..63a3479c0 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -196,9 +196,11 @@ def main(policy, *, evals: list[Eval], output_dir: str | Path | None = None, tim # The session runs no inference, but a session that serves its model on a runtime needs one to open. warmup = Executor(policy.functions) try: - policy.new_session(None, warmup).close() + session = policy.new_session(None, warmup) finally: + # The runtime closes first: a call the session started still uses what the session holds. warmup.close() + session.close() output_dir = prepare_output_dir(output_dir) try: diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index 08f014a9d..ab38273bd 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -7,7 +7,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import Future, ThreadPoolExecutor from functools import partial -from typing import Any, cast +from typing import Any from positronic.policy.base import ( Answer, @@ -140,7 +140,7 @@ def new_session(self, context=None, rt=None) -> ChunkSession: inner = self._inner.new_session(context, own) # Waiting out the work is a question about a chunk, so nothing above a ``ChunkPlayer`` blocks. assert not isinstance(inner, Session), f'{type(inner).__name__} answers commands, not a chunk' - return _BlockingPolicy._Session(cast(ChunkSession, inner), own) + return _BlockingPolicy._Session(inner, own) except BaseException: own.close() raise diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 8d2c5a90f..c1777d677 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -29,16 +29,19 @@ def __init__(self, policy: Policy, context: dict[str, Any], charges_wall_time: b # 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) + session = None try: session = policy.new_session(context, self._runtime) # A stack that answers chunks has no player in it, so nothing turns them into commands. assert isinstance(session, Session), ( f'{type(session).__name__} answers a chunk; a policy played here needs a ChunkPlayer in its stack' ) - self._session = session except BaseException: self._runtime.close() + if session is not None: + session.close() raise + self._session = session @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index a420f26af..4e8056ada 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -115,7 +115,6 @@ def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> N """ if isinstance(chunk, dict): chunk = [chunk] - anchor = time_ns / 1e9 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( @@ -123,11 +122,12 @@ def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> N f'the sessions below are timing actions against a clock of their own' ) # The single explicit seconds->ns seam: the sessions below time actions in float seconds, and - # every pimm channel is in ns. A waypoint naming no channel — the codecs' end-of-chunk sentinel — - # commands nothing and states where the chunk ends. + # 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( ( - int((anchor + action.get(keys.ACTION_TIMESTAMP, 0.0)) * 1e9), + time_ns + int(action.get(keys.ACTION_TIMESTAMP, 0.0) * 1e9), {name: value for name, value in action.items() if name != keys.ACTION_TIMESTAMP}, ) for action in chunk diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index 30ca1c5aa..217ac1522 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. @@ -60,7 +60,7 @@ 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 AnySession, Layer, Session +from positronic.policy.base import AnySession, ChunkSession, Layer, Session from positronic.policy.codec import is_action from positronic.utils.rerun_compat import log_numeric_series, set_timeline_sequence, set_timeline_time @@ -363,14 +363,15 @@ def __init__(self, rec: Recorder, name: str): def make_session(self, inner: AnySession) -> AnySession: stream = self._rec._open_stream() - return _RecordingTapSession(inner, self._rec, self._name, stream) + # The side a tap sits on is the side its inner is on, so a chain of taps carries it outwards. + tap = _CommandTapSession if isinstance(inner, Session) else _ChunkTapSession + return tap(inner, self._rec, self._name, stream) -class _RecordingTapSession(Session): +class _RecordingTapSession: """Logs the observation flowing down and the actions flowing up at one point. - A tap sits on either side of a ``ChunkPlayer``, so it forwards whatever its inner answers and reads the - side off that. TODO(#661): the framework records each boundary it carries, and a tap stops guessing. + ``_CommandTapSession`` and ``_ChunkTapSession`` add the call, one per side of a ``ChunkPlayer``. """ def __init__(self, inner: AnySession, rec: Recorder, name: str, stream: rr.RecordingStream): @@ -461,10 +462,13 @@ def _actions(answer) -> list[dict]: Above the player the answer holds the commands of one round, so a plot reads a point per round rather than a chunk per inference. """ - if isinstance(answer, tuple): - commands, _ = answer - return [dict(commands)] if commands else [] - return answer or [] + match answer: + case (commands, _): + return [dict(commands)] if commands else [] + case dict(): + return [answer] + case _: + return answer or [] def _send_blueprint(self) -> None: rec = self._rec @@ -474,7 +478,7 @@ def _send_blueprint(self) -> None: if bp is not None: rr.send_blueprint(bp) - def __call__(self, obs, time_ns) -> Any: + def _tapped(self, obs, time_ns) -> Any: rec = self._rec outermost = rec._depth == 0 if outermost: @@ -512,3 +516,17 @@ def cancel(self): def close(self): self._inner.close() self._rec._release_stream() + + +class _CommandTapSession(_RecordingTapSession, Session): + """A tap above a ``ChunkPlayer``: it logs the command of each round.""" + + def __call__(self, obs, time_ns) -> tuple[Mapping[str, Any], int]: + return self._tapped(obs, time_ns) + + +class _ChunkTapSession(_RecordingTapSession, ChunkSession): + """A tap under a ``ChunkPlayer``: it logs the chunk a session answers.""" + + def __call__(self, obs, time_ns) -> list[dict[str, Any]] | dict[str, Any] | None: + return self._tapped(obs, time_ns) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 8a20d39e6..a4f421803 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -73,7 +73,7 @@ def __init__(self, ws_session: InferenceSession, rt: Runtime, compress_images: b 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__`` + A server answer of one action becomes a 1-element list, which is the form ``ChunkSession.__call__`` returns. """ if self._answer is None: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 94ad27065..fd89f3923 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1931,6 +1931,23 @@ def test_episode_virtual_duration_starts_when_the_rig_is_ready(world, tmp_path): assert virtual_s < draw_s +@pytest.mark.parametrize( + ('resume_in_sec', 'sleep_sec'), + [(0.002, 0.002), (0.05, POLL_PERIOD_SEC), (-0.001, POLL_PERIOD_SEC), (None, POLL_PERIOD_SEC)], +) +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, and no longer than its own poll period. An + instant already reached asks for the next round.""" + 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) + + assert isinstance(command, pimm.Sleep) + assert command.seconds == pytest.approx(sleep_sec, abs=1e-3) + + 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()) @@ -2067,7 +2084,7 @@ class _ObservedTicks(Layer): def __init__(self): self.seen: list[float] = [] - class _Session(DelegatingSession, Session): + class _Session(DelegatingSession): def __init__(self, inner: Session, seen: list[float]): super().__init__(inner) self._seen = seen diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index bc7c9d7ab..f11790215 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -171,6 +171,31 @@ def test_a_chunk_plays_one_waypoint_at_a_time(self): assert session(_obs(), int(1.2e9)) == ({}, int(1.5e9)) assert inner._session.call_count == 1 + def test_waypoints_due_together_keep_every_channel_they_name(self): + """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 = ChunkPlayer().wrap(inner).new_session() + + 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_the_last_waypoint_of_a_chunk_survives_the_call_that_loads_the_next(self): + """The call that drains a chunk loads the next one, and a channel only the drained waypoint names + still reaches the driver.""" + inner = _ScriptedSession([ + [{'arm': 1, keys.ACTION_TIMESTAMP: 0.0}, {'grip': 0.9, keys.ACTION_TIMESTAMP: 0.5}], + [{'arm': 2, keys.ACTION_TIMESTAMP: 0.2}], + ]) + session = ChunkPlayer().make_session(inner) + + assert session(_obs(), int(1e9)) == ({'arm': 1}, int(1.5e9)) + assert session(_obs(), int(1.5e9)) == ({'grip': 0.9}, int(1.7e9)) + def test_a_channel_with_several_waypoints_due_keeps_the_last(self): """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)]) diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index df4c04148..de5af346e 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -5,10 +5,12 @@ from positronic.drivers.roboarm.command import CartesianPosition from positronic.geom import Rotation, Transform3D from positronic.policy.base import ChunkSession, Policy +from positronic.policy.layers import ChunkPlayer from positronic.policy.recording import ( Recorder, _build_blueprint, _command_field_arrays, + _CommandTapSession, _flat_wire, _RecordingTapSession, _squeeze_batch, @@ -227,6 +229,18 @@ def new_session(self, context=None, rt=None): assert session._step == 1 +def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path): + """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) + session = (rec.tap('raw') | ChunkPlayer()).wrap(_TrackingPolicy(actions)).new_session() + assert isinstance(session, _CommandTapSession) + + 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.tap('raw') | rec.tap('server')).wrap(_TrackingPolicy()) diff --git a/positronic/probe.py b/positronic/probe.py index ab3cb891d..3822706bb 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -40,9 +40,8 @@ _TAP = 'raw' # Observation keys the endpoint expects, mirroring the inference harness, plus every image.*. _STATE_KEYS = (keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP) -# How long to wait for the endpoint, and how often to ask the session again while it is out. +# How long to wait for the endpoint before calling the round trip lost. _ANSWER_TIMEOUT_SEC = 120.0 -_POLL_PERIOD_SEC = 0.01 def _build_wire_obs(sample: dict, task: str | None, now_ns: int, recorded_ts: int) -> dict: @@ -71,29 +70,26 @@ def _meta_doc(name: str, meta: dict) -> str: return f'## {name}\n\n{rows}' -def _play(session, obs: dict) -> list[tuple[int, dict]]: +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 session - anchors the chunk on the call that receives it and asks for a call at each waypoint; a resume time that - has not moved past the call says it holds nothing — the endpoint is still out, or the chunk has run out. + 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. """ + session(obs, time.time_ns()) + runtime.wait(_ANSWER_TIMEOUT_SEC) + if runtime.in_flight: + raise TimeoutError(f'the endpoint answered nothing in {_ANSWER_TIMEOUT_SEC:.0f}s') played: list[tuple[int, dict]] = [] - deadline = time.monotonic() + _ANSWER_TIMEOUT_SEC now_ns = time.time_ns() while True: commands, resume_at_ns = session(obs, now_ns) if commands: played.append((now_ns, dict(commands))) - if resume_at_ns > now_ns: - now_ns = resume_at_ns - elif played: + if resume_at_ns <= now_ns: # the chunk has run out, and the session is asking the endpoint again return played - elif time.monotonic() >= deadline: - raise TimeoutError(f'the endpoint commanded nothing in {_ANSWER_TIMEOUT_SEC:.0f}s') - else: - time.sleep(_POLL_PERIOD_SEC) - now_ns = time.time_ns() + now_ns = resume_at_ns def _is_cartesian_chunk(played: list[tuple[int, dict]]) -> bool: @@ -172,20 +168,24 @@ def main( meta = dict(session.meta) name = label or _recording_name(meta) try: - played = _play(session, obs) - print(f'episode {episode} @ {at:.3f}s (ts={ts}) [{name}]: {len(played)} command(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) - _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))) + played = _play(session, obs, runtime) finally: - # The runtime closes first: the call the last waypoint started is still using the session's socket. + # The call that drained the chunk asked the endpoint again; closing the runtime waits that answer + # out. It goes before the session, 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() def _internal_main(): diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index 38e960c60..b2651554c 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -20,7 +20,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 ChunkSession, Codec, Layer, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints @@ -256,7 +256,7 @@ def stop(self): self.process = None -class _DreamZeroSession(Session): +class _DreamZeroSession(ChunkSession): def __init__(self, client: RoboarenaClient, session_id: str): self._client = client self._session_id = session_id diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index cd63a8af4..c9d6837e7 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -17,7 +17,7 @@ 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 ChunkSession, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -200,7 +200,7 @@ def stop(self): ########################################################################################### -class _Gr00tSession(Session): +class _Gr00tSession(ChunkSession): def __init__(self, client: PolicyClient): self._client = client diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index a518b437c..d2efa028b 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -7,7 +7,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 ChunkSession, Policy from positronic.policy.observation import TASK_FIELD @@ -30,7 +30,7 @@ def _detect_device() -> str: return 'cpu' -class _LerobotSession(Session): +class _LerobotSession(ChunkSession): def __init__(self, policy, preprocessor, postprocessor, device: str, meta: dict[str, Any]): self._policy = policy self._preprocessor = preprocessor diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index bef059a35..cbf031017 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -14,7 +14,7 @@ from positronic import keys from positronic.cfg import codecs -from positronic.policy import Codec, Policy, Session +from positronic.policy import ChunkSession, Codec, Policy from positronic.policy.base import Answer, Runtime from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.observation import TASK_FIELD @@ -83,7 +83,7 @@ def _infer(policy: PreTrainedPolicy, device: str, obs: dict[str, Any]) -> list[d class LerobotPolicy(Policy): _INFER = 'infer' - class _Session(Session): + class _Session(ChunkSession): """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]): diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index d74599a7f..9b6d4ff22 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -5,7 +5,7 @@ from transformers import AutoModelForImageTextToText, AutoProcessor from positronic import keys -from positronic.policy import Policy, Session +from positronic.policy import ChunkSession, 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,7 +24,7 @@ def warm_observation() -> dict[str, Any]: } -class _MolmoAct2Session(Session): +class _MolmoAct2Session(ChunkSession): def __init__(self, model, processor, norm_tag: str, num_steps: int, meta: dict[str, Any]): self._model = model self._processor = processor @@ -66,7 +66,7 @@ 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: + def new_session(self, context=None, rt=None) -> ChunkSession: return _MolmoAct2Session(self._model, self._processor, self._norm_tag, self._num_steps, self._meta) @property diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 6b4837630..f15faf0ed 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -14,7 +14,7 @@ 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 ChunkSession, Codec, Policy from positronic.policy.codec import ChangeEEFrame, RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -124,7 +124,7 @@ def stop(self): ########################################################################################### -class _OpenpiSession(Session): +class _OpenpiSession(ChunkSession): def __init__(self, client: WebsocketClientPolicy): self._client = client From 1a4b8dcbc236cf9a8b4ceeefa55ad2dc4adc8522 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 27 Aug 2026 13:49:25 +0200 Subject: [PATCH 03/10] Read a tap's answer by the side it sits on, not by its shape A `match` sequence pattern matches a two-element list, so a chunk of two actions read as a `(commands, resume_at_ns)` pair: the tap plotted the first action as a command and dropped the second. Each tap now states what its own side answers, which is what it knew when it was made. `probe` logs nothing for an endpoint that commands nothing, instead of taking the joint-velocity branch on an empty list and raising `IndexError`. It waits on the round trip without a deadline of its own: the endpoint's `infer_timeout` is the one that bounds it, and a failed round trip raises out of the next call with its own cause. --- positronic/policy/recording.py | 27 ++++++++++------------- positronic/policy/tests/test_recording.py | 12 ++++++++++ positronic/probe.py | 13 ++++++----- 3 files changed, 31 insertions(+), 21 deletions(-) diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index 217ac1522..bcecde58d 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -455,20 +455,9 @@ def _log_action_chunk(self, prefix: str, actions: list[dict], obs: Mapping[str, _log_action_series(f'{prefix}/series/{suffix}', field_arr, group_h, base_ns, names=None) self._rec._series_paths.append(f'{prefix}/series/{suffix}') - @staticmethod - def _actions(answer) -> list[dict]: - """The actions in what a session answered: a chunk under a ``ChunkPlayer``, one command above one. - - Above the player the answer holds the commands of one round, so a plot reads a point per round - rather than a chunk per inference. - """ - match answer: - case (commands, _): - return [dict(commands)] if commands else [] - case dict(): - return [answer] - case _: - return answer or [] + def _actions(self, answer) -> list[dict]: + """The actions the chunk logger plots out of what the inner answered.""" + raise NotImplementedError def _send_blueprint(self) -> None: rec = self._rec @@ -519,14 +508,22 @@ def close(self): class _CommandTapSession(_RecordingTapSession, Session): - """A tap above a ``ChunkPlayer``: it logs the command of each round.""" + """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 __call__(self, obs, time_ns) -> tuple[Mapping[str, Any], int]: return self._tapped(obs, time_ns) + def _actions(self, answer) -> list[dict]: + commands, _ = answer + return [dict(commands)] if commands else [] + class _ChunkTapSession(_RecordingTapSession, ChunkSession): """A tap under a ``ChunkPlayer``: it logs the chunk a session answers.""" def __call__(self, obs, time_ns) -> list[dict[str, Any]] | dict[str, Any] | None: return self._tapped(obs, time_ns) + + def _actions(self, answer) -> list[dict]: + return [dict(answer)] if isinstance(answer, Mapping) else (answer or []) diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index de5af346e..62b9127c0 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -9,6 +9,7 @@ from positronic.policy.recording import ( Recorder, _build_blueprint, + _ChunkTapSession, _command_field_arrays, _CommandTapSession, _flat_wire, @@ -229,6 +230,17 @@ def new_session(self, context=None, rt=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}] + session = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(chunk)).new_session() + assert isinstance(session, _ChunkTapSession) + + assert session._actions(chunk) == chunk + assert session._actions({'v': 1.0}) == [{'v': 1.0}] # one action may come back bare + + def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path): """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}] diff --git a/positronic/probe.py b/positronic/probe.py index 3822706bb..d7426a343 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -40,8 +40,6 @@ _TAP = 'raw' # Observation keys the endpoint expects, mirroring the inference harness, plus every image.*. _STATE_KEYS = (keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP) -# How long to wait for the endpoint before calling the round trip lost. -_ANSWER_TIMEOUT_SEC = 120.0 def _build_wire_obs(sample: dict, task: str | None, now_ns: int, recorded_ts: int) -> dict: @@ -75,12 +73,13 @@ def _play(session, obs: dict, runtime: Executor) -> list[tuple[int, dict]]: 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. + 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()) - runtime.wait(_ANSWER_TIMEOUT_SEC) - if runtime.in_flight: - raise TimeoutError(f'the endpoint answered nothing in {_ANSWER_TIMEOUT_SEC:.0f}s') + # 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: @@ -105,6 +104,8 @@ def _log_commands(played: list[tuple[int, dict]], obs: dict, wall_ns: int, inf_n 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 not played: + return commands = [c for _, c in played] if _is_cartesian_chunk(played): poses = [c[keys.ROBOT_COMMAND].pose for c in commands] From b409a023a41c4553723012c9a0ad651b49349db0 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 27 Aug 2026 13:58:44 +0200 Subject: [PATCH 04/10] Command nothing once the trial passes the instant it stops at The round tests the deadline before it calls the session, and the world can reach the deadline while that call runs. `_reschedule` used to drop a chunk answered past it; `_infer` now holds back the commands instead, and `_run` finishes the trial on the next round. --- positronic/policy/harness.py | 7 +++++-- positronic/policy/tests/test_harness.py | 24 ++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index c1777d677..4b6c32a20 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -382,8 +382,11 @@ def _infer(self, inference: _EpisodeInference, clock: pimm.Clock, should_stop: p return # no function is in flight yet, so this skips no wait commands, self._resume_at_ns = inference(obs) self._telemetry.step() - # 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}) + # 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 is None or clock.now() < self._deadline: + # 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: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index fd89f3923..0545780c6 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -296,6 +296,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() @@ -1948,6 +1951,27 @@ def test_a_real_rig_wakes_at_the_moment_the_session_asked_for(world, resume_in_s assert command.seconds == pytest.approx(sleep_sec, abs=1e-3) +@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 = world.clock.now() + (-1.0 if expired else 1.0) + 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 + + 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()) From 73b80fd87f4f89d35a17a69ea11d471dcdba99e1 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 27 Aug 2026 14:06:00 +0200 Subject: [PATCH 05/10] Close the session `blocking()` opened before it refused it Finding out which contract a session answers takes opening it, and a model session holds a backend connection from that moment. The runtime still closes first, and the session it served follows. --- positronic/policy/executor.py | 5 ++++- positronic/policy/tests/test_executor.py | 26 +++++++++++++++++++++++- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index ab38273bd..ce90a83a3 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -136,14 +136,17 @@ def close(self): def new_session(self, context=None, rt=None) -> ChunkSession: assert rt is None, 'a blocking policy serves its own functions; nothing above it runs them' own = Executor(self._inner.functions) + inner = None try: inner = self._inner.new_session(context, own) # Waiting out the work is a question about a chunk, so nothing above a ``ChunkPlayer`` blocks. assert not isinstance(inner, Session), f'{type(inner).__name__} answers commands, not a chunk' - return _BlockingPolicy._Session(inner, own) except BaseException: own.close() + if inner is not None: + inner.close() raise + return _BlockingPolicy._Session(inner, own) @property def functions(self) -> Mapping[str, Callable[..., Any]]: diff --git a/positronic/policy/tests/test_executor.py b/positronic/policy/tests/test_executor.py index 378c97ad3..422ea261e 100644 --- a/positronic/policy/tests/test_executor.py +++ b/positronic/policy/tests/test_executor.py @@ -10,7 +10,7 @@ import pytest -from positronic.policy.base import Answer, ChunkSession, DelegatingChunkSession, Layer, NotAnswered, Policy +from positronic.policy.base import Answer, ChunkSession, DelegatingChunkSession, Layer, NotAnswered, Policy, Session from positronic.policy.executor import Executor, blocking # How long a test waits for the worker threads before calling the call lost. @@ -326,6 +326,30 @@ def make(policy: Policy) -> ChunkSession: class TestBlocking: """A policy whose sessions answer in the call that asked.""" + def test_a_session_that_answers_commands_is_refused_and_closed(self): + """Nothing above a ``ChunkPlayer`` has work to wait out, and the session opened to find that out + holds a connection of its own.""" + + class _Commanding(Session): + def __init__(self): + self.closed = False + + def __call__(self, obs, time_ns): + return {}, time_ns + + def close(self): + self.closed = True + + session = _Commanding() + + class _Commander(Policy): + def new_session(self, context=None, rt=None): + return session + + with pytest.raises(AssertionError, match='answers commands'): + blocking(_Commander()).new_session() + assert session.closed, 'the session it opened to check was dropped without closing' + def test_a_session_that_answers_in_its_own_call_is_called_one_time(self, opened): policy = _PlainPolicy() From a807baa0c09480141e7c57cf38f8b1afc7244006 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 27 Aug 2026 14:35:49 +0200 Subject: [PATCH 06/10] Read the deadline off the field that carries it `deadline_ns` replaced the float `_deadline` on main while this branch was open, and the guard the round runs after a session call was written against the old name. --- positronic/policy/harness.py | 2 +- positronic/policy/tests/test_harness.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 4b6c32a20..b9fe7ee39 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -384,7 +384,7 @@ def _infer(self, inference: _EpisodeInference, clock: pimm.Clock, should_stop: p self._telemetry.step() # 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 is None or clock.now() < self._deadline: + 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) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 0545780c6..610e6444e 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1964,7 +1964,7 @@ def test_a_command_is_emitted_only_while_the_trial_still_has_budget(world, expir 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 = world.clock.now() + (-1.0 if expired else 1.0) + 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) From 009593bbe8d200363049a6a9180bc228d67f746b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 28 Aug 2026 13:50:24 +0200 Subject: [PATCH 07/10] Let a new chunk supersede the one it replaces The call that drains a chunk also loads the next, and it merged what the drained chunk still held into what the new one commands. A channel only the drained chunk named then took a setpoint from the plan that had just been replaced, while every other channel followed the new one. A driver holds what it last took, so the new chunk alone says what to command. This is what `_reschedule` did, which cleared every channel before it refilled. --- positronic/policy/layers.py | 8 ++------ positronic/policy/tests/test_layers.py | 8 ++++---- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 4e8056ada..99dc7431b 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -91,16 +91,12 @@ def __call__(self, obs, time_ns): chunk = self._inner(obs, time_ns) if chunk is not None: self._load(chunk, time_ns) - commands = {**commands, **self._due(time_ns)} + commands = self._due(time_ns) # A spent chunk asks for the next call at once: the work for the chunk after it is in flight. return commands, self._waypoints[0][0] if self._waypoints else time_ns def _due(self, time_ns: int) -> dict[str, Any]: - """The waypoints reached at ``time_ns``, taken off the chunk. - - A channel with several waypoints due keeps the last: exact for an absolute setpoint, lossy for - a relative one. - """ + """The waypoints reached at ``time_ns``, taken off the chunk; several on a channel keep the last.""" commands: dict[str, Any] = {} while self._waypoints and self._waypoints[0][0] <= time_ns: commands.update(self._waypoints.popleft()[1]) diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index f11790215..5fbee9e9e 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -184,9 +184,9 @@ def test_waypoints_due_together_keep_every_channel_they_name(self): 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_the_last_waypoint_of_a_chunk_survives_the_call_that_loads_the_next(self): - """The call that drains a chunk loads the next one, and a channel only the drained waypoint names - still reaches the driver.""" + 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 = _ScriptedSession([ [{'arm': 1, keys.ACTION_TIMESTAMP: 0.0}, {'grip': 0.9, keys.ACTION_TIMESTAMP: 0.5}], [{'arm': 2, keys.ACTION_TIMESTAMP: 0.2}], @@ -194,7 +194,7 @@ def test_the_last_waypoint_of_a_chunk_survives_the_call_that_loads_the_next(self session = ChunkPlayer().make_session(inner) assert session(_obs(), int(1e9)) == ({'arm': 1}, int(1.5e9)) - assert session(_obs(), int(1.5e9)) == ({'grip': 0.9}, int(1.7e9)) + assert session(_obs(), int(1.5e9)) == ({}, int(1.7e9)) def test_a_channel_with_several_waypoints_due_keeps_the_last(self): """A round that finds more than one waypoint due commands the latest of them.""" From e6ded061730e9700fd0bbcb144800ca4f2a59c7c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 29 Aug 2026 17:54:42 +0200 Subject: [PATCH 08/10] Pace the control loop by the instant the session names The harness sleeps to `resume_at_ns` instead of a fixed poll period, bounded by the episode deadline, a 1 ms floor and a 1 s ceiling. `_EpisodeInference` asserts the instant is in the future, so `ChunkPlayer` and `StopOnFault` each name a poll period of their own while they wait for something they cannot time. --- ARCHITECTURE.md | 3 +- positronic/offboard/tests/conftest.py | 2 +- positronic/offboard/tests/test_server.py | 4 +-- positronic/policy/base.py | 3 +- positronic/policy/executor.py | 2 -- positronic/policy/harness.py | 31 +++++++++------- positronic/policy/layers.py | 46 ++++++++++-------------- positronic/policy/tests/test_harness.py | 35 ++++++++++++------ positronic/policy/tests/test_layers.py | 44 +++++++++++++---------- positronic/probe.py | 2 +- positronic/tests/testing_coutils.py | 4 ++- 11 files changed, 98 insertions(+), 78 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index fc9c7a2c4..ba4aa3894 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -130,7 +130,8 @@ anything it encodes about either side breaks the any-to-any goal. It assembles t 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. The harness aims at that instant and may call earlier. +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/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index baee1ee54..9f450cb55 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -92,7 +92,7 @@ def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> For a session topped by a ``ChunkPlayer``: both calls get 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) == ({}, time_ns), 'a round-trip was already in flight' + assert session(obs, time_ns)[0] == {}, 'a round-trip was already in flight' rt.wait(ANSWER_SEC) assert not rt.in_flight, 'the round-trip never came back' return session(obs, time_ns)[0] diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index d6e4a290a..e0dd18c92 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -282,14 +282,12 @@ def pipeline(): def play_out(session, runtime): """Every waypoint the session plays, walked by the instants it asks to be called at.""" at_ns = int(100e9) - assert session({keys.OBS_TIME_NS: 0}, at_ns) == ({}, at_ns), 'a round trip was already in flight' + assert session({keys.OBS_TIME_NS: 0}, at_ns)[0] == {}, 'a round trip was already in flight' runtime.wait(ANSWER_SEC) played = [] for _ in range(4): commands, resume_at_ns = session({keys.OBS_TIME_NS: 0}, at_ns) played.append((at_ns, commands)) - if resume_at_ns <= at_ns: - break at_ns = resume_at_ns return played diff --git a/positronic/policy/base.py b/positronic/policy/base.py index a1f7ce69c..bb4e9f762 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -87,7 +87,8 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> tuple[Mapping[str, A ``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 on the clock of ``time_ns``. The caller aims at it and may call earlier. + ``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. """ diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index ce90a83a3..e2a869fd3 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -69,8 +69,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) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index b9fe7ee39..6ac1cb2e9 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -15,9 +15,12 @@ from positronic.policy.executor import Executor from positronic.utils import flatten_dict, frozen_view -# How long a real-time round may last when the session asks for nothing 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: @@ -61,7 +64,9 @@ def __call__(self, obs: dict[str, Any]) -> tuple[Mapping[str, Any], int]: # 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.""" @@ -73,7 +78,7 @@ 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. @@ -235,22 +240,22 @@ 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 moment the session asked for, capped at the poll - period, so a waypoint is emitted at its own time and a round rarely finds more than one due. A - session that names a moment already reached asks for the next round, one poll period away. - """ + """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_in_sec = (self._resume_at_ns - clock.now_ns()) / 1e9 if self._resume_at_ns is not None else 0.0 - return pimm.Sleep(min(POLL_PERIOD_SEC, due_in_sec) if due_in_sec > 0 else POLL_PERIOD_SEC) + 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 diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 99dc7431b..ba2939631 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -6,7 +6,7 @@ """ from collections import deque -from typing import Any +from typing import Any, NamedTuple import numpy as np @@ -49,13 +49,14 @@ class StopOnFault(Layer): """ 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 {}, time_ns + return {}, time_ns + int(StopOnFault.POLL_SEC * 1e9) def make_session(self, inner: Session) -> Session: return StopOnFault._Session(inner) @@ -65,42 +66,33 @@ def to_spec(self): class ChunkPlayer(Layer): - """Hold the chunk the sessions below answer, and emit each waypoint at its own time. - - Owns relative→absolute time conversion: the sessions below (codecs, models) emit timestamps relative to - the call; this layer anchors them to the call's ``time_ns``. Each call answers the waypoints that have - come due, and asks to be called again at the next one. The call that drains the last waypoint asks the - sessions below for a new chunk, and plays whatever of it is already due. - """ + """Hold the chunk the sessions below answer, anchor it to the call that received it, and emit each + waypoint at its own time.""" WIRE_NAME = 'chunk_player' + POLL_SEC = 0.01 class _Session(Session): - """Plays the chunk it holds; re-queries in the call that drains it. - - The one session with a different contract on each side, so it forwards rather than delegates. - """ + class _Waypoint(NamedTuple): + cmd: dict[str, Any] + time_ns: int def __init__(self, inner: ChunkSession): self._inner = inner - self._waypoints: deque[tuple[int, dict[str, Any]]] = deque() + self._waypoints: deque[ChunkPlayer._Session._Waypoint] = deque() def __call__(self, obs, time_ns): - commands = self._due(time_ns) - if not self._waypoints: + """Plays the chunk it holds; re-queries in the call that drains it.""" + if not self._waypoints or self._waypoints[-1].time_ns <= time_ns: chunk = self._inner(obs, time_ns) if chunk is not None: self._load(chunk, time_ns) - commands = self._due(time_ns) - # A spent chunk asks for the next call at once: the work for the chunk after it is in flight. - return commands, self._waypoints[0][0] if self._waypoints else time_ns - - def _due(self, time_ns: int) -> dict[str, Any]: - """The waypoints reached at ``time_ns``, taken off the chunk; several on a channel keep the last.""" commands: dict[str, Any] = {} - while self._waypoints and self._waypoints[0][0] <= time_ns: - commands.update(self._waypoints.popleft()[1]) - return commands + 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 _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> None: """Anchor ``chunk`` to ``time_ns`` and hold it. @@ -122,9 +114,9 @@ def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> N # 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( - ( - time_ns + int(action.get(keys.ACTION_TIMESTAMP, 0.0) * 1e9), + 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 ) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 610e6444e..576efa4cd 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -24,12 +24,12 @@ from positronic.offboard.client import InferenceSession from positronic.policy.base import ChunkSession, DelegatingSession, Layer, Policy, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import POLL_PERIOD_SEC, Harness, _EpisodeInference +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 INFER, RemoteSession, 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 @@ -1063,8 +1063,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) @@ -1936,11 +1936,12 @@ def test_episode_virtual_duration_starts_when_the_rig_is_ready(world, tmp_path): @pytest.mark.parametrize( ('resume_in_sec', 'sleep_sec'), - [(0.002, 0.002), (0.05, POLL_PERIOD_SEC), (-0.001, POLL_PERIOD_SEC), (None, POLL_PERIOD_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, and no longer than its own poll period. An - instant already reached asks for the next round.""" + """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) @@ -1948,7 +1949,19 @@ def test_a_real_rig_wakes_at_the_moment_the_session_asked_for(world, resume_in_s command = harness._pace(world.clock) assert isinstance(command, pimm.Sleep) - assert command.seconds == pytest.approx(sleep_sec, abs=1e-3) + assert command.seconds == pytest.approx(sleep_sec, abs=1e-4) + + +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)]) @@ -1993,6 +2006,8 @@ class _ReplanEarly(Layer): """ class _Session(Session): + POLL_SEC = 0.001 + def __init__(self, inner: ChunkSession): self._inner = inner self._waypoints: deque[tuple[int, dict[str, Any]]] = deque() @@ -2005,8 +2020,8 @@ def __call__(self, obs, time_ns): 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] - return commands, min(due, default=time_ns) + 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: chunk = self._inner(obs, time_ns) diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 5fbee9e9e..3a2653b37 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -52,7 +52,9 @@ def _obs(now_sec=0.0, status=RobotStatus.AVAILABLE): class _CommandSession(Session): - """Answers a fixed command mapping and asks for its next call at once, as a layer above a player does.""" + """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 @@ -60,7 +62,7 @@ def __init__(self, commands): def __call__(self, obs, time_ns): self.call_count += 1 - return self._commands, time_ns + return self._commands, time_ns + int(self.POLL_SEC * 1e9) class TestStopOnFault: @@ -69,14 +71,14 @@ def test_an_unavailable_arm_stops_what_is_executing(self, unavailable): inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session(_obs(0.0, unavailable), 0) == ({}, 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 = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session(_obs(0.0, RobotStatus.AVAILABLE), 0) == ({'v': 1}, 0) + 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): @@ -84,7 +86,7 @@ def test_an_observation_with_no_arm_status_reaches_the_model(self): inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session({keys.OBS_TIME_NS: 0}, 0) == ({'v': 1}, 0) + 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): @@ -98,7 +100,7 @@ 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) == ({}, 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): @@ -107,7 +109,7 @@ def test_the_status_a_recording_carries_for_a_taken_arm_stops_the_policy(self): 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) == ({}, 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): @@ -115,7 +117,10 @@ def test_the_status_published_for_a_travelling_arm_reaches_the_model(self): inner = _CommandSession({'v': 1}) session = StopOnFault().make_session(inner) - assert session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 2}, 0) == ({'v': 1}, 0) + 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): @@ -133,7 +138,7 @@ def test_recovery_plans_afresh_instead_of_resuming(self): session = (StopOnFault() | ChunkPlayer()).wrap(inner).new_session() 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)) + 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._session.call_count == 2 @@ -153,13 +158,13 @@ def __call__(self, obs, time_ns): class TestChunkPlayer: 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 chunk. The player - commands nothing, asks for the next call at once, and asks the session again on it.""" + commands nothing, waits its own poll period, and asks the session again.""" inner = _ScriptedSession([None, None, [{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]]) session = ChunkPlayer().make_session(inner) - assert session(_obs(0.0), int(1e9)) == ({}, int(1e9)) - assert session(_obs(0.1), int(1e9)) == ({}, int(1e9)) - assert session(_obs(0.2), int(1e9)) == ({'v': 1}, int(1e9)) + 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)) + assert session(_obs(0.2), int(1e9)) == ({'v': 1}, int(1e9) + int(ChunkPlayer.POLL_SEC * 1e9)) assert inner.call_count == 3 def test_a_chunk_plays_one_waypoint_at_a_time(self): @@ -213,12 +218,12 @@ def test_the_call_that_drains_the_chunk_asks_for_the_next_one(self): assert session(_obs(), int(1.5e9)) == ({'v': 1}, int(2.0e9)) assert inner._session.call_count == 2 - def test_a_single_action_asks_for_the_next_chunk_at_once(self): + def test_a_single_action_is_drained_by_the_call_that_loads_it(self): """A chunk of one action at ts=0 is drained by the call that loads it, so the next call re-queries.""" session = ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])).new_session() - assert session(_obs(1.0), int(1e9)) == ({'v': 1}, int(1e9)) - assert session(_obs(1.01), int(1.01e9)) == ({'v': 1}, int(1.01e9)) + 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): """The codecs close a chunk with a timestamp-only sentinel; it states where the chunk ends.""" @@ -244,7 +249,10 @@ def test_layer_pipe_layer(self): assert isinstance(pipeline, Layer) policy = pipeline.wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session() - assert session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}, int(1e9)) == ({'v': 1}, int(1e9)) + 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): """A codec composes above a player, and refuses the commands it would silently corrupt.""" @@ -591,7 +599,7 @@ def test_inline_tolerates_marker_less_pipe(self): inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) policy = spec.inline(ChunkPlayer() | spec.PolicySource(inner)) session = policy.new_session() - assert session(_obs(), int(1e9)) == ({'v': 1}, int(1e9)) + 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/probe.py b/positronic/probe.py index d7426a343..8ca0b6e42 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -86,7 +86,7 @@ def _play(session, obs: dict, runtime: Executor) -> list[tuple[int, dict]]: commands, resume_at_ns = session(obs, now_ns) if commands: played.append((now_ns, dict(commands))) - if resume_at_ns <= now_ns: # the chunk has run out, and the session is asking the endpoint again + 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 diff --git a/positronic/tests/testing_coutils.py b/positronic/tests/testing_coutils.py index d92e7b417..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 {}, time_ns + return {}, time_ns + int(self.POLL_SEC * 1e9) @property def meta(self): From 34af1a8841e0ddf0f34936f81f0f7d6ffd67b6ce Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 29 Aug 2026 20:41:21 +0200 Subject: [PATCH 09/10] Let the caller own the calls a chunk source makes `ChunkSession.__call__` answers an `Answer` instead of a chunk-or-`None`. The one-call-in-flight state machine leaves `RemoteSession` and `LerobotPolicy._Session`, which each become one line, and moves to `ChunkPlayer`: it holds the handle, decides when to ask and when to read, and could hold more than one. A source that runs the work inside the call answers a `Done`. --- .basedpyright/baseline.json | 24 ----- positronic/offboard/server.py | 8 +- positronic/offboard/server_utils.py | 5 +- positronic/offboard/tests/conftest.py | 23 +++-- .../offboard/tests/test_remote_policy.py | 95 +------------------ positronic/offboard/tests/test_server.py | 24 ++--- positronic/policy/__init__.py | 4 + positronic/policy/base.py | 45 +++++++-- positronic/policy/codec.py | 11 +-- positronic/policy/executor.py | 12 +-- positronic/policy/layers.py | 31 +++++- positronic/policy/recording.py | 94 ++++++++++-------- positronic/policy/remote.py | 37 +------- positronic/policy/tests/test_executor.py | 62 ++++++------ .../policy/tests/test_golden_pipeline.py | 55 ++++++----- positronic/policy/tests/test_harness.py | 34 ++++--- positronic/policy/tests/test_layers.py | 87 ++++++++++++++--- positronic/policy/tests/test_policy_io.py | 25 +++-- positronic/policy/tests/test_recording.py | 36 ++++--- .../env_server/tests/test_remote_env.py | 6 +- positronic/vendors/dreamzero/server.py | 6 +- positronic/vendors/gr00t/server.py | 4 +- positronic/vendors/lerobot/policy.py | 8 +- positronic/vendors/lerobot_0_3_3/policy.py | 26 +---- positronic/vendors/molmoact2/policy.py | 6 +- positronic/vendors/openpi/server.py | 4 +- 26 files changed, 378 insertions(+), 394 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 7d7089635..c9c019d20 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": { @@ -5065,22 +5057,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 21, - "lineCount": 1 - } - }, { "code": "reportCallIssue", "range": { diff --git a/positronic/offboard/server.py b/positronic/offboard/server.py index f71818573..c6e2e25b0 100644 --- a/positronic/offboard/server.py +++ b/positronic/offboard/server.py @@ -18,7 +18,7 @@ from starlette.datastructures import QueryParams from positronic import keys -from positronic.policy import Policy, Recorder +from positronic.policy import ChunkSession, Policy, Recorder from positronic.policy.base import Layer from positronic.policy.executor import blocking from positronic.policy.spec import ModelSource, Pipeline, split @@ -334,7 +334,8 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): session = await asyncio.to_thread(served.new_session) finally: self._infer_lock.release() - assert session is not None + # The wire carries chunks, so a stack that answers commands has nothing this server can send. + assert isinstance(session, ChunkSession), f'{type(session).__name__} answers commands, not a chunk' # Later entries win: per-episode session facts over static ones, the server's own last. meta = { **self.metadata, @@ -358,7 +359,8 @@ async def _serve_session(self, websocket: WebSocket, model_id: str | None): # 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()) + answer = await asyncio.to_thread(session, raw_obs, time.time_ns()) + actions = answer.result() await websocket.send_bytes(serialise({protocol.RESULT: actions})) except Exception as e: logger.error(f'Error processing message: {e}', exc_info=True) diff --git a/positronic/offboard/server_utils.py b/positronic/offboard/server_utils.py index 6e27d2ff4..f55420a24 100644 --- a/positronic/offboard/server_utils.py +++ b/positronic/offboard/server_utils.py @@ -11,7 +11,7 @@ from collections.abc import Callable from typing import Any -from positronic.policy import Policy +from positronic.policy import ChunkSession, Policy from positronic.policy.executor import blocking logger = logging.getLogger(__name__) @@ -43,8 +43,9 @@ 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() + assert isinstance(session, ChunkSession), f'{type(session).__name__} answers commands, not a chunk' try: - run_with_progress(lambda: session(obs, time.time_ns()), 'Running warmup inference', on_progress) + run_with_progress(lambda: session(obs, time.time_ns()).result(), 'Running warmup inference', on_progress) finally: session.close() diff --git a/positronic/offboard/tests/conftest.py b/positronic/offboard/tests/conftest.py index 9f450cb55..23fbc9965 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -10,7 +10,7 @@ import uvicorn from positronic.offboard.server import PolicyServer -from positronic.policy import AnySession, ChunkSession, Policy, Session +from positronic.policy import AnySession, ChunkSession, Done, Policy, Session from positronic.policy.executor import Executor from positronic.policy.layers import ChunkPlayer from positronic.policy.spec import ModelSource, PolicySource, remote @@ -76,23 +76,22 @@ def make(policy: Policy) -> tuple[AnySession, Executor]: def round_trip(session: ChunkSession, rt: Executor, obs, time_ns: int = 0) -> list[dict] | dict | None: - """What ``session`` answers for ``obs``, over the two calls one round trip takes. - - Both calls get the same ``time_ns``, so a chunk comes back anchored at the value the test passed. - """ - assert session(obs, time_ns) is None, 'a round-trip was already in flight' + """What ``session`` answers for ``obs``, once its round trip has come back.""" + answer = session(obs, time_ns) rt.wait(ANSWER_SEC) assert not rt.in_flight, 'the round-trip never came back' - return session(obs, time_ns) + return answer.result() def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> Mapping[str, Any]: - """What ``session`` commands for ``obs``, over the two calls one round trip takes. + """What ``session`` commands for ``obs`` once the call under its ``ChunkPlayer`` has come back. - For a session topped by a ``ChunkPlayer``: both calls get 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. + 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)[0] == {}, '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)[0] @@ -101,7 +100,7 @@ def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> def _make_mock_policy(action, meta): """Create a mock policy with session-based API.""" session = MagicMock() - session.return_value = action + session.return_value = Done(action) session.meta = meta session.close = MagicMock() diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index b1a5a5911..459795b37 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,7 +11,7 @@ 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, played_round_trip, round_trip +from positronic.offboard.tests.conftest import played_round_trip, round_trip from positronic.policy import RemotePolicy from positronic.policy.codec import ActionHorizon from positronic.policy.layers import ChunkPlayer @@ -320,31 +319,6 @@ def test_remote_session_passes_through_none(open_session): 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 - - 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.""" @@ -354,73 +328,6 @@ def test_opening_a_session_without_a_runtime_is_refused(): 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(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): """The ``policy.infer`` span is recorded at the remote inference boundary itself, not by a layer in front of it.""" diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index e0dd18c92..089c1943a 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -244,16 +244,9 @@ def test_pipeline_with_no_rig_side_half_refused_at_startup(make_mock_policy): class _ScriptedSession(ChunkSession): 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() + return self._rt.fns[_INFER](obs) class _ScriptedPolicy(Policy): @@ -280,15 +273,18 @@ def pipeline(): local_session, local_rt = open_session(inline(pipeline())) def play_out(session, runtime): - """Every waypoint the session plays, walked by the instants it asks to be called at.""" + """Every waypoint the session plays, and how long after the chunk loaded each one went out.""" at_ns = int(100e9) - assert session({keys.OBS_TIME_NS: 0}, at_ns)[0] == {}, 'a round trip was already in flight' - runtime.wait(ANSWER_SEC) - played = [] - for _ in range(4): + 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) - played.append((at_ns, commands)) + 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) diff --git a/positronic/policy/__init__.py b/positronic/policy/__init__.py index ee7b6361b..d97c83f72 100644 --- a/positronic/policy/__init__.py +++ b/positronic/policy/__init__.py @@ -1,9 +1,11 @@ from .base import ( + Answer, AnySession, ChunkSession, DelegatingChunkSession, DelegatingPolicy, DelegatingSession, + Done, Layer, Policy, Session, @@ -28,4 +30,6 @@ 'ActionTiming', 'is_action', 'Recorder', + 'Answer', + 'Done', ] diff --git a/positronic/policy/base.py b/positronic/policy/base.py index bb4e9f762..e584b6a89 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -2,7 +2,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Mapping -from typing import Any, ClassVar +from typing import Any, ClassVar, final # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. SEQ = 'seq' @@ -26,6 +26,36 @@ def done(self) -> bool: ... def result(self) -> Any: """What the function returned. Raises what the function raised, or ``NotAnswered`` before it answers.""" + @final + def map(self, fn: Callable[[Any], Any]) -> Answer: + """This answer with ``fn`` applied to what it carries, at the moment it is read.""" + return _Mapped(self, fn) + + +class _Mapped(Answer): + def __init__(self, inner: Answer, fn: Callable[[Any], Any]): + self._inner = inner + self._fn = fn + + def done(self) -> bool: + return self._inner.done() + + def result(self) -> Any: + return self._fn(self._inner.result()) + + +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 + # A call to an ``Fn`` starts the work and returns at once. Fn = Callable[..., Answer] @@ -95,20 +125,21 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> tuple[Mapping[str, A class ChunkSession(_SessionBase, ABC): """Session that answers action chunks: what a ``ChunkPlayer`` holds and plays. - A call answers a list of actions, each naming command channels and carrying its own - ``keys.ACTION_TIMESTAMP`` in seconds from the call. ``None`` means "no new chunk" — what a session - answers while the work it asked for is still in flight. An empty list drops what is playing. One - action may come back bare, and the player wraps it. + A chunk is 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. One action + may come back bare, and the player wraps it. TODO(#661): the chunk becomes the answer of a served function, and this class goes with the session-serving wire. """ @abstractmethod - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | dict[str, Any] | None: - """The chunk to play from now, or ``None`` while the work is in flight. + def __call__(self, obs: Mapping[str, Any], time_ns: int) -> Answer: + """Start the work for one chunk, and answer the caller's handle on it. ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. + The caller decides how many calls it keeps in flight, and reads each handle when it wants the + chunk. A session that runs the work inside the call answers a ``Done``. """ diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 844f31173..51008dfc4 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -24,7 +24,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, ChunkSession, DelegatingChunkSession, Layer, _ComposedLayer +from positronic.policy.base import PAR, SEQ, Answer, ChunkSession, DelegatingChunkSession, Layer, _ComposedLayer from positronic.utils import merge_dicts _QUAT = geom.Rotation.Representation.QUAT @@ -108,14 +108,11 @@ def __init__(self, inner: ChunkSession, codec: 'Codec'): self._codec = codec def __call__(self, obs, time_ns): - encoded = self._codec.encode(obs) - action = self._inner(encoded, time_ns) + answer = self._inner(self._codec.encode(obs), time_ns) # A codec decodes a chunk, so it belongs under the player. Above one it would be handed commands # paired with a resume time, and each codec would drop or corrupt them in its own way. - assert not isinstance(action, tuple), 'compose a Codec under the ChunkPlayer, not above it' - if action is None: - return None - return self._codec.decode(action) + assert isinstance(answer, Answer), 'compose a Codec under the ChunkPlayer, not above it' + return answer.map(self._codec.decode) @property def meta(self): diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index e2a869fd3..fe3a1c0b1 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -119,12 +119,10 @@ def __init__(self, inner: ChunkSession, rt: Executor): super().__init__(inner) self._rt = rt - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] | 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 __call__(self, obs: Mapping[str, Any], time_ns: int) -> Answer: + answer = self._inner(obs, time_ns) + self._rt.wait() + return answer def close(self): # The runtime closes first: a call in flight is still using what the session holds. @@ -152,7 +150,7 @@ def functions(self) -> Mapping[str, Callable[..., Any]]: def blocking(policy: Policy) -> Policy: - """``policy`` with its heavy work waited out: a session answers in the call that asked. + """``policy`` with its heavy work waited out: the handle a call answers is already done. 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 diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index ba2939631..50127557a 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -12,7 +12,7 @@ from positronic import keys from positronic.drivers.roboarm import RobotStatus -from positronic.policy.base import ChunkSession, DelegatingSession, Layer, Session +from positronic.policy.base import Answer, ChunkSession, DelegatingSession, Layer, 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. @@ -80,13 +80,20 @@ class _Waypoint(NamedTuple): def __init__(self, inner: ChunkSession): self._inner = inner self._waypoints: deque[ChunkPlayer._Session._Waypoint] = deque() + # The one call this layer 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): - """Plays the chunk it holds; re-queries in the call that drains it.""" + """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: - chunk = self._inner(obs, time_ns) - if chunk is not None: - self._load(chunk, time_ns) + if self._answer is None: + self._answer = self._inner(obs, time_ns) + 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) @@ -94,6 +101,18 @@ def __call__(self, obs, time_ns): 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 layer 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. @@ -127,6 +146,8 @@ def meta(self): def cancel(self): 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 self._inner.cancel() def close(self): diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index bcecde58d..67c166a4f 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -42,16 +42,16 @@ 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 itertools from collections.abc import Mapping 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,7 +60,7 @@ 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 AnySession, ChunkSession, Layer, Session +from positronic.policy.base import Answer, AnySession, ChunkSession, Layer, Session from positronic.policy.codec import is_action from positronic.utils.rerun_compat import log_numeric_series, set_timeline_sequence, set_timeline_time @@ -368,6 +368,17 @@ def make_session(self, inner: AnySession) -> AnySession: return tap(inner, self._rec, self._name, stream) +class _Frame(NamedTuple): + """What one call is logged against: its observation, the rerun timelines it set, and its step. + + A chunk is logged when the caller reads it, so its frame outlives the call that started the work. + """ + + obs: Mapping[str, Any] + timelines: dict[str, Any] + step: int + + class _RecordingTapSession: """Logs the observation flowing down and the actions flowing up at one point. @@ -385,10 +396,10 @@ def __init__(self, inner: AnySession, rec: Recorder, name: str, stream: rr.Recor def meta(self): return {**self._inner.meta, 'recording.rrd': str(self._rec._rrd_path)} - 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: Mapping[str, Any]) -> None: """Recursively log obs *data* under *prefix*, recording entity paths on the Recorder.""" @@ -405,20 +416,21 @@ def _log(self, prefix: str, data: Mapping[str, Any]) -> None: log_numeric_series(path, num) self._rec._numeric_paths.append(path) - def _log_action_chunk(self, prefix: str, actions: list[dict], obs: Mapping[str, Any]) -> 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: @@ -467,38 +479,37 @@ def _send_blueprint(self) -> None: if bp is not None: rr.send_blueprint(bp) - def _tapped(self, obs, time_ns) -> Any: + def _tapped(self, obs, time_ns) -> tuple[Any, _Frame]: + """What the inner answered, and the frame this call is logged against.""" 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 try: + frame = _Frame(obs, dict(rec._timeline_values), self._step) with self._stream: - self._set_timelines() + self._set_timelines(frame) self._log(self._name, obs) - - answer = self._inner(obs, time_ns) - - if actions := self._actions(answer): - 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 answer + return self._inner(obs, time_ns), frame finally: rec._depth -= 1 if rec._depth == 0: rec._timeline_values = {} + def _log_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 outermost tap.""" + if self._rec._depth == 0 and self._step == 0: + with self._stream: + self._send_blueprint() + self._step += 1 + def cancel(self): self._inner.cancel() @@ -512,18 +523,21 @@ class _CommandTapSession(_RecordingTapSession, Session): rather than a chunk per inference.""" def __call__(self, obs, time_ns) -> tuple[Mapping[str, Any], int]: - return self._tapped(obs, time_ns) - - def _actions(self, answer) -> list[dict]: + answer, frame = self._tapped(obs, time_ns) commands, _ = answer - return [dict(commands)] if commands else [] + self._log_actions([dict(commands)] if commands else [], frame) + self._end_step() + return answer class _ChunkTapSession(_RecordingTapSession, ChunkSession): - """A tap under a ``ChunkPlayer``: it logs the chunk a session answers.""" + """A tap under a ``ChunkPlayer``: it logs the chunk a call brings back, when the caller reads it.""" - def __call__(self, obs, time_ns) -> list[dict[str, Any]] | dict[str, Any] | None: - return self._tapped(obs, time_ns) + def __call__(self, obs, time_ns) -> Answer: + answer, frame = self._tapped(obs, time_ns) + self._end_step() + return answer.map(partial(self._log_chunk, frame)) - def _actions(self, answer) -> list[dict]: - return [dict(answer)] if isinstance(answer, Mapping) else (answer or []) + def _log_chunk(self, frame: _Frame, chunk): + self._log_actions([dict(chunk)] if isinstance(chunk, Mapping) else (chunk or []), frame) + return chunk diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index a4f421803..bb714e4ee 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -56,10 +56,6 @@ def round_trip( class RemoteSession(ChunkSession): """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``). """ @@ -67,43 +63,20 @@ def __init__(self, ws_session: InferenceSession, rt: Runtime, compress_images: b 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. + def __call__(self, obs: cabc.Mapping[str, Any], time_ns: int) -> Answer: + """One round trip to the server, and the caller's handle on the trajectory it brings back. - A server answer of one action becomes a 1-element list, which is the form ``ChunkSession.__call__`` - returns. + A server answer of one action becomes a 1-element list, which is the form a chunk takes. """ - 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 + answer = self._rt.fns[INFER](self._session, obs, self._compress_images) + return answer.map(lambda result: [result] if isinstance(result, dict) else result) @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() diff --git a/positronic/policy/tests/test_executor.py b/positronic/policy/tests/test_executor.py index 422ea261e..3550b7f79 100644 --- a/positronic/policy/tests/test_executor.py +++ b/positronic/policy/tests/test_executor.py @@ -10,7 +10,16 @@ import pytest -from positronic.policy.base import Answer, ChunkSession, DelegatingChunkSession, Layer, NotAnswered, Policy, Session +from positronic.policy.base import ( + Answer, + ChunkSession, + DelegatingChunkSession, + Done, + Layer, + NotAnswered, + Policy, + Session, +) from positronic.policy.executor import Executor, blocking # How long a test waits for the worker threads before calling the call lost. @@ -243,7 +252,7 @@ def __init__(self): def __call__(self, obs, time_ns): self.calls += 1 - return [{'action': obs}] + return Done([{'action': obs}]) def __init__(self): self.session = _PlainPolicy._Session() @@ -256,33 +265,23 @@ def new_session(self, context=None, rt=None) -> ChunkSession: class _EchoPolicy(Policy): - """Serves ``echo``, and makes sessions that take ``rounds`` calls of it to answer.""" + """Serves ``echo``: its session hands the caller the handle on one call of it.""" class _Session(ChunkSession): - def __init__(self, rt, rounds: int): + def __init__(self, rt): 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 + return self._rt.fns[_ECHO](obs) + + def __init__(self): self.session: _EchoPolicy._Session def new_session(self, context=None, rt=None) -> ChunkSession: assert rt is not None - self.session = _EchoPolicy._Session(rt, self._rounds) + self.session = _EchoPolicy._Session(rt) return self.session @property @@ -353,39 +352,36 @@ def new_session(self, context=None, rt=None): 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 opened(blocking(policy))({'x': 1}, 0.0).result() == [{'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_the_handle_a_call_answers_is_already_done(self, opened): + """The wait is what ``blocking`` adds: the caller reads the chunk without asking a second time.""" + policy = _EchoPolicy() - def test_a_session_that_starts_nothing_and_answers_none_is_called_one_time(self, opened): - policy = _EchoPolicy(rounds=0) + answer = opened(blocking(policy))({'x': 1}, 0.0) - assert opened(blocking(policy))({'x': 1}, 0.0) is None + assert answer.done() + assert answer.result() == {'x': 1} 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) + layer, policy = _CountingLayer(), _EchoPolicy() - assert opened(layer.wrap(blocking(policy)))({'x': 1}, 0.0) == {'x': 1} - assert (layer.calls, policy.session.calls) == (1, 3) + assert opened(layer.wrap(blocking(policy)))({'x': 1}, 0.0).result() == {'x': 1} + assert (layer.calls, policy.session.calls) == (1, 1) 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 == {} + assert blocking(_EchoPolicy()).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) + policy = _EchoPolicy() session = blocking(policy).new_session() session.close() diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index e2a754e5e..7f5c16fd6 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -41,7 +41,7 @@ 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 ChunkSession, DelegatingChunkSession, DelegatingPolicy, Policy +from positronic.policy.base import Answer, ChunkSession, DelegatingChunkSession, DelegatingPolicy, Done, Policy from positronic.policy.codec import ActionTiming from positronic.policy.harness import Harness from positronic.policy.layers import ChunkPlayer, StopOnFault @@ -73,7 +73,7 @@ def __call__(self, obs, time_ns): step = current + delta * 0.5 * ((i + 1) / 10.0) pose = Transform3D(translation=step.astype(np.float32), rotation=Rotation.identity) chunk.append({keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': round(0.50 + 0.01 * i, 4)}) - return chunk + return Done(chunk) class ScriptedProportionalPolicy(Policy): @@ -88,42 +88,44 @@ def new_session(self, context=None, rt=None): class _SimulatedLatency(DelegatingPolicy): - """A fixed inference latency in world time: every chunk is stamped for, and reaches the player at, - ``latency_sec`` after the call that produced it, whatever the machine took. Sits under the player, so - nothing below it sees an observation while a chunk is held. A skip or a stop passes through at once.""" + """A fixed inference latency in world time: every chunk is stamped for, and answers at, ``latency_sec`` + after the call that produced it, whatever the machine took. Sits under the player, so nothing below it + sees an observation while a chunk is in flight.""" - def __init__(self, inner: Policy, latency_sec: float): + def __init__(self, inner: Policy, latency_sec: float, clock: pimm.Clock): super().__init__(inner) self._latency_ns = round(latency_sec * 1e9) + self._clock = clock + + class _Held(Answer): + """The chunk 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 + + def done(self) -> bool: + return self._clock.now_ns() >= self._release_at_ns and self._inner.done() + + def result(self): + return self._inner.result() class _Session(DelegatingChunkSession): - def __init__(self, inner: ChunkSession, latency_ns: int): + def __init__(self, inner: ChunkSession, latency_ns: int, clock: pimm.Clock): super().__init__(inner) self._latency_ns = latency_ns - self._held: list | dict | None = None - self._release_at_ns = 0 + self._clock = clock 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 time 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() + # The chunk answers at ``latency_ns``, so the sessions below time it from that instant. + release_at_ns = time_ns + self._latency_ns + return _SimulatedLatency._Held(self._inner(obs, release_at_ns), release_at_ns, self._clock) def new_session(self, context=None, rt=None) -> ChunkSession: inner = self._inner.new_session(context, rt) assert isinstance(inner, ChunkSession) - return _SimulatedLatency._Session(inner, self._latency_ns) + return _SimulatedLatency._Session(inner, self._latency_ns, self._clock) class FakeRobot(pimm.ControlSystem): @@ -210,7 +212,8 @@ def _run_pipeline(tmp_path: Path) -> dict: simulated=True, ) harness = Harness( - (StopOnFault() | ChunkPlayer()).wrap(_SimulatedLatency(policy, INFERENCE_LATENCY_S)), embodiment + (StopOnFault() | ChunkPlayer()).wrap(_SimulatedLatency(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 576efa4cd..1024f6b53 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -22,7 +22,7 @@ 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 ChunkSession, DelegatingSession, Layer, Policy, Session +from positronic.policy.base import Answer, ChunkSession, DelegatingSession, Done, Layer, Policy, Session from positronic.policy.codec import ActionTimestamp from positronic.policy.harness import MAX_ROUND_SEC, MIN_ROUND_SEC, WAIT_PERIOD_SEC, Harness, _EpisodeInference from positronic.policy.layers import ChunkPlayer, StopOnFault @@ -86,7 +86,9 @@ def __init__(self, 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}] + return Done([ + {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.target_grip, 'timestamp': 0.0} + ]) class SpyPolicy(Policy): @@ -114,7 +116,9 @@ def __init__(self, policy): 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}] + return Done([ + {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.target_grip, 'timestamp': 0.0} + ]) @property def meta(self): @@ -158,14 +162,14 @@ def __init__(self, policy): def __call__(self, obs, time_ns): self._policy.counter += 1 dt = 0.005 - return [ + return Done([ { keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.counter * 100.0 + i, 'timestamp': i * dt, } for i in range(10) - ] + ]) class ChunkPolicy(StubPolicy): @@ -601,7 +605,7 @@ 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}] + return Done([{keys.ROBOT_COMMAND: self._command, 'target_grip': 0.0, 'timestamp': 0.0}]) class _StaticMetaPolicy(Policy): def __init__(self): @@ -698,8 +702,7 @@ def __init__(self, rt): self._rt = rt def __call__(self, obs, time_ns): - self._rt.fns[INFER]() - return None + return self._rt.fns[INFER]() def new_session(self, context=None, rt=None): assert rt is not None @@ -1518,7 +1521,7 @@ def test_empty_trajectory_leaves_every_channel_holding(world): class _EmptyChunkSession(ChunkSession): def __call__(self, obs, time_ns): - return [] + return Done([]) class EmptyChunkPolicy(Policy): def new_session(self, context=None, rt=None): @@ -2012,6 +2015,7 @@ def __init__(self, inner: ChunkSession): self._inner = inner 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): commands: dict[str, Any] = {} @@ -2024,9 +2028,11 @@ def __call__(self, obs, time_ns): return commands, min(due, default=time_ns + int(self.POLL_SEC * 1e9)) def _replan(self, obs, time_ns: int) -> None: - chunk = self._inner(obs, time_ns) - if not isinstance(chunk, list): # the function it asked has still to answer + if self._answer is None: + self._answer = self._inner(obs, time_ns) + 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 self._waypoints = deque( (int((anchor + action[keys.ACTION_TIMESTAMP]) * 1e9), {keys.TARGET_GRIP: action[keys.TARGET_GRIP]}) @@ -2419,11 +2425,11 @@ def __call__(self, obs, time_ns): pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) command = CartesianPosition(pose=pose) if self._calls == 1: - return [ + return Done([ {keys.ROBOT_COMMAND: command, keys.TARGET_GRIP: 0.5, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10) - ] - return [{keys.ROBOT_COMMAND: command, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10)] + ]) + return Done([{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): diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 3a2653b37..0117bbf97 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -11,7 +11,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 ChunkSession, Layer, Policy, Session +from positronic.policy.base import Answer, ChunkSession, Done, Layer, Policy, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -27,6 +27,26 @@ from positronic.policy.observation import ObservationCodec +class _Pending(Answer): + """A call the test answers by hand, so a player can be watched while it waits.""" + + def __init__(self, value=None, failure: BaseException | None = None): + self._value = value + self._failure = failure + self._done = False + + def answer(self): + self._done = True + + def done(self): + return self._done + + def result(self): + if self._failure is not None: + raise self._failure + return self._value + + class _ConstSession(ChunkSession): def __init__(self, actions): self._actions = actions @@ -34,7 +54,7 @@ def __init__(self, actions): def __call__(self, obs, time_ns): self.call_count += 1 - return self._actions + return Done(self._actions) class _ConstPolicy(Policy): @@ -144,7 +164,7 @@ def test_recovery_plans_afresh_instead_of_resuming(self): class _ScriptedSession(ChunkSession): - """Answers each of ``script`` in turn — ``None`` where a session has nothing to place yet.""" + """Answers each of ``script`` in turn, each entry a handle the player holds.""" def __init__(self, script): self._script = list(script) @@ -156,16 +176,61 @@ def __call__(self, obs, time_ns): class TestChunkPlayer: - 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 chunk. The player - commands nothing, waits its own poll period, and asks the session again.""" - inner = _ScriptedSession([None, None, [{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]]) + 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 = _ScriptedSession([pending]) session = ChunkPlayer().make_session(inner) 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 == 3 + 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 = _ScriptedSession([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) + session = ChunkPlayer().make_session(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 = ChunkPlayer().make_session(_ScriptedSession([pending])) + + session(_obs(), int(1e9)) + pending.answer() + session.cancel() + + 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 = _ScriptedSession([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) + session = ChunkPlayer().make_session(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): """The player anchors the chunk on the call that receives it and asks for a call at each waypoint.""" @@ -193,8 +258,8 @@ 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 = _ScriptedSession([ - [{'arm': 1, keys.ACTION_TIMESTAMP: 0.0}, {'grip': 0.9, keys.ACTION_TIMESTAMP: 0.5}], - [{'arm': 2, keys.ACTION_TIMESTAMP: 0.2}], + 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 = ChunkPlayer().make_session(inner) @@ -315,7 +380,7 @@ def __init__(self): def __call__(self, obs, time_ns): self.seen.append(obs) - return [] + return Done([]) class _CapturePolicy(Policy): diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index 2f4822c62..d7191fa0a 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -8,7 +8,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 ChunkSession, Policy +from positronic.policy.base import Answer, ChunkSession, Done, Policy from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -120,7 +120,7 @@ def __init__(self, result): self._result = result def __call__(self, obs, time_ns): - return self._result + return Done(self._result) class _ChunkPolicy(Policy): @@ -157,12 +157,19 @@ def meta(self): _T0_OBS = {obs_keys.OBS_TIME_NS: 0} +def _chunk(session, obs=_T0_OBS, time_ns: int = 0): + """The chunk a session under a player answers for ``obs``.""" + answer = session(obs, time_ns) + assert isinstance(answer, Answer), 'the session answers commands, not a chunk' + return answer.result() + + 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) + result = _chunk(policy.new_session()) 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 @@ -172,7 +179,7 @@ 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) + result = _chunk(policy.new_session()) assert len(result) == 6 # 5 actions + timestamp sentinel @@ -181,7 +188,7 @@ def test_action_horizon_sec_larger_than_chunk(): # 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) + result = _chunk(policy.new_session()) assert len(result) == 4 # 3 actions + timestamp sentinel (nothing truncated) @@ -189,7 +196,7 @@ 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) + result = _chunk(policy.new_session()) 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) @@ -200,7 +207,7 @@ def test_action_horizon_sec_seconds_truncates(): # 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) + result = _chunk(policy.new_session()) assert isinstance(result, list) and len(result) == 4 # 3 actions + horizon sentinel dt = 1.0 / 30.0 for i, action in enumerate(result): @@ -251,7 +258,7 @@ 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) + result = _chunk(policy.new_session()) 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 @@ -260,7 +267,7 @@ def test_action_timestamp_and_horizon_compose(): def test_single_action_has_zero_timestamp(): codec = ActionTiming(fps=15.0) policy = codec.wrap(_SinglePolicy()) - result = policy.new_session()(_T0_OBS, 0) + result = _chunk(policy.new_session()) 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 62b9127c0..a6c8e05ad 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -4,7 +4,7 @@ from positronic.drivers.roboarm import command from positronic.drivers.roboarm.command import CartesianPosition from positronic.geom import Rotation, Transform3D -from positronic.policy.base import ChunkSession, Policy +from positronic.policy.base import Answer, ChunkSession, Done, Policy from positronic.policy.layers import ChunkPlayer from positronic.policy.recording import ( Recorder, @@ -19,13 +19,20 @@ ) +def _chunk(session, obs, time_ns: int = 0): + """The chunk a session under a player answers for ``obs``.""" + answer = session(obs, time_ns) + assert isinstance(answer, Answer), 'the session answers commands, not a chunk' + return answer.result() + + class _TrackingSession(ChunkSession): def __init__(self, actions, meta): self._actions = actions self._meta = meta def __call__(self, obs, time_ns): - return list(self._actions) + return Done(list(self._actions)) @property def meta(self): @@ -60,7 +67,7 @@ def __init__(self, rec, actions): 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) + return Done(list(self._actions)) class _CapturingPolicy(Policy): @@ -120,7 +127,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.tap('t').wrap(_TrackingPolicy(actions)).new_session(), {'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' @@ -166,7 +173,7 @@ 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) + result = _chunk(session, {'x': 1.0, keys.WALL_TIME_NS: 1_000_000}) assert result == actions @@ -210,23 +217,23 @@ def test_logs_command_chunk_without_mutating(tmp_path): {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) + result = _chunk(session, {'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(ChunkSession): +def test_handles_an_empty_chunk(tmp_path): + class _EmptySession(ChunkSession): def __call__(self, obs, time_ns): - return None + return Done([]) - class _NonePolicy(Policy): + class _EmptyPolicy(Policy): def new_session(self, context=None, rt=None): - return _NoneSession() + return _EmptySession() - session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() + session = Recorder(tmp_path).tap('t').wrap(_EmptyPolicy()).new_session() assert isinstance(session, _RecordingTapSession) - assert session({'x': 1.0}, 0) is None + assert _chunk(session, {'x': 1.0}) == [] assert session._step == 1 @@ -237,8 +244,7 @@ def test_a_two_action_chunk_is_not_read_as_a_command_pair(tmp_path): session = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(chunk)).new_session() assert isinstance(session, _ChunkTapSession) - assert session._actions(chunk) == chunk - assert session._actions({'v': 1.0}) == [{'v': 1.0}] # one action may come back bare + assert session({'x': 1.0}, 0).result() == chunk def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path): diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 1172ef2ec..12418a62e 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -20,7 +20,7 @@ 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 ChunkSession, Policy +from positronic.policy import ChunkSession, Done, Policy from positronic.policy.codec import ActionTimestamp from positronic.policy.layers import ChunkPlayer from positronic.policy.tests.test_harness import StubPolicy @@ -379,10 +379,10 @@ def __init__(self, policy: _JointposChunks): def __call__(self, obs, time_ns): self._policy.chunks += 1 - return [ + return Done([ {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.chunks * 100.0 + i} for i in range(self._policy.chunk_len) - ] + ]) @pytest.mark.timeout(60.0) diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index b2651554c..a7e922384 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -20,7 +20,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 ChunkSession, Codec, Layer, Policy +from positronic.policy import ChunkSession, Codec, Done, Layer, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote from positronic.utils.checkpoints import list_checkpoints @@ -268,8 +268,8 @@ def __call__(self, obs, time_ns): # 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])] + return Done([{'action': action_array}]) + return Done([{'action': action_array[i]} for i in range(action_array.shape[0])]) def close(self): try: diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index c9d6837e7..467a03dee 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -17,7 +17,7 @@ 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 ChunkSession, Policy +from positronic.policy import ChunkSession, Done, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -210,7 +210,7 @@ def __call__(self, obs, time_ns): 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)] + return Done([{k: v[i] for k, v in action.items()} for i in range(time_horizon)]) class Gr00tPolicy(Policy): diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index d2efa028b..730eefa4e 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -7,7 +7,7 @@ from lerobot.policies.factory import get_policy_class, make_pre_post_processors from positronic import keys -from positronic.policy import ChunkSession, Policy +from positronic.policy import ChunkSession, Done, Policy from positronic.policy.observation import TASK_FIELD @@ -38,7 +38,7 @@ def __init__(self, policy, preprocessor, postprocessor, device: str, meta: dict[ self._device = device self._meta = meta - def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: + def __call__(self, obs: dict[str, Any], time_ns: int) -> Done: obs_int = {} for key, val in obs.items(): if key == TASK_FIELD: @@ -62,8 +62,8 @@ def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: action = action.cpu().numpy().squeeze(0) if action.ndim == 1: - return [{'action': action}] - return [{'action': a} for a in action] + return Done([{'action': action}]) + return Done([{'action': a} for a in action]) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index cbf031017..bea4927f9 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -84,32 +84,14 @@ class LerobotPolicy(Policy): _INFER = 'infer' class _Session(ChunkSession): - """Per-episode session that gives the model call to the runtime, and answers the chunk on a later call.""" + """Per-episode session that gives the model call to the runtime.""" 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 + + def __call__(self, obs: dict[str, Any], time_ns: int) -> Answer: + return self._rt.fns[LerobotPolicy._INFER](obs) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index 9b6d4ff22..c2e4bf5d8 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -5,7 +5,7 @@ from transformers import AutoModelForImageTextToText, AutoProcessor from positronic import keys -from positronic.policy import ChunkSession, Policy +from positronic.policy import ChunkSession, Done, 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 @@ -32,7 +32,7 @@ def __init__(self, model, processor, norm_tag: str, num_steps: int, meta: dict[s self._num_steps = num_steps self._meta = meta - def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: + def __call__(self, obs: dict[str, Any], time_ns: int) -> Done: # 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. @@ -49,7 +49,7 @@ def __call__(self, obs: dict[str, Any], time_ns: int) -> list[dict[str, Any]]: enable_cuda_graph=False, ) actions = out.actions[0].float().cpu().numpy() - return [{'action': action} for action in actions] + return Done([{'action': action} for action in actions]) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index f15faf0ed..74c5ddb9d 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -14,7 +14,7 @@ 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 ChunkSession, Codec, Policy +from positronic.policy import ChunkSession, Codec, Done, Policy from positronic.policy.codec import ChangeEEFrame, RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -131,7 +131,7 @@ def __init__(self, client: WebsocketClientPolicy): def __call__(self, obs, time_ns): response = self._client.infer(obs) actions = response['actions'] - return [{'action': a} for a in actions] + return Done([{'action': a} for a in actions]) class OpenpiPolicy(Policy): From 7390bb2ea59e02e96f7b5b16f7c8dde178a466d6 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sun, 30 Aug 2026 12:16:24 +0200 Subject: [PATCH 10/10] Open a policy's work per episode and play it in one `Session` `Session` and `ChunkSession` answered different types from one method name, so every caller narrowed a union, `DelegatingChunkSession` repeated its twin, and three `isinstance` asserts stood in for a type the code could not state. A policy now declares the work of one episode under `INFER`, and `Policy.episode` opens what that work holds and closes it after. `ChunkPlayer` holds `rt.fns[INFER]` and is the bottom session. `Codec` and the chunk tap become `ChunkLayer`s, which wrap that call rather than a session, so the encode and the JPEG leave the control thread. `ChunkSession`, `AnySession`, `DelegatingChunkSession`, `Done`, `Answer.map`, `RemoteSession`, `blocking`, `Session.meta` and `Policy.functions` go with it. --- .basedpyright/baseline.json | 24 -- docs/connect-your-model.md | 44 ++-- positronic/cli/eval/run.py | 12 +- positronic/offboard/server.py | 42 ++- positronic/offboard/server_utils.py | 12 +- positronic/offboard/tests/conftest.py | 102 +++++--- positronic/offboard/tests/test_offboard.py | 12 +- .../offboard/tests/test_remote_policy.py | 90 +++---- positronic/offboard/tests/test_server.py | 44 +--- positronic/policy/__init__.py | 20 +- positronic/policy/base.py | 241 +++++++----------- positronic/policy/codec.py | 33 +-- positronic/policy/executor.py | 59 +---- positronic/policy/harness.py | 42 ++- positronic/policy/layers.py | 49 ++-- positronic/policy/recording.py | 174 +++++++------ positronic/policy/remote.py | 64 ++--- positronic/policy/tests/conftest.py | 52 ++++ positronic/policy/tests/test_executor.py | 160 +----------- .../policy/tests/test_golden_pipeline.py | 65 ++--- positronic/policy/tests/test_harness.py | 206 ++++++--------- positronic/policy/tests/test_layers.py | 210 ++++++++------- positronic/policy/tests/test_policy_io.py | 71 +++--- positronic/policy/tests/test_recording.py | 167 +++++------- positronic/probe.py | 23 +- .../env_server/tests/test_remote_env.py | 30 +-- positronic/vendors/dreamzero/server.py | 51 ++-- positronic/vendors/gr00t/server.py | 28 +- positronic/vendors/lerobot/policy.py | 64 ++--- positronic/vendors/lerobot_0_3_3/policy.py | 32 +-- positronic/vendors/molmoact2/policy.py | 57 ++--- positronic/vendors/openpi/server.py | 20 +- 32 files changed, 936 insertions(+), 1364 deletions(-) create mode 100644 positronic/policy/tests/conftest.py diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index c9c019d20..89b62aed3 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5186,30 +5186,6 @@ "endColumn": 58, "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": [ diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 3ee4a3e02..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 it once per timestep with the latest observation and your clock reading, and it answers the actions to run. A policy served behind the `remote` marker answers a chunk, so its sessions are `ChunkSession`s; the rig's `ChunkPlayer` holds that chunk and emits each waypoint at its own time. Per-episode state (history, the chunk 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 ChunkSession, Policy +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 ChunkPlayer, StopOnFault -class MySession(ChunkSession): +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,14 +200,6 @@ class MySession(ChunkSession): 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'} @@ -213,9 +213,9 @@ The pipeline reads left to right: everything left of the `remote` marker is the 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 (`ChunkPlayer() | 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() | ChunkPlayer() | RestrictImageSize() | remote | codec | source @@ -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/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index 63a3479c0..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 Executor from positronic.policy.harness import Harness from positronic.simulator.env_server.telemetry import ATTR_RUN_ID, ENV_RUN_ID, ENV_TELEMETRY_DIR @@ -193,14 +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. - warmup = Executor(policy.functions) - try: - session = policy.new_session(None, warmup) - finally: - # The runtime closes first: a call the session started still uses what the session holds. - warmup.close() - 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/server.py b/positronic/offboard/server.py index c6e2e25b0..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 ChunkSession, 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 @@ -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,34 +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() - # The wire carries chunks, so a stack that answers commands has nothing this server can send. - assert isinstance(session, ChunkSession), f'{type(session).__name__} answers commands, not a chunk' - # 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'), @@ -358,9 +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. - answer = await asyncio.to_thread(session, raw_obs, time.time_ns()) - actions = answer.result() + 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) @@ -381,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 f55420a24..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 ChunkSession, Policy -from positronic.policy.executor import blocking +from positronic.policy import INFER, Policy logger = logging.getLogger(__name__) @@ -42,12 +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() - assert isinstance(session, ChunkSession), f'{type(session).__name__} answers commands, not a chunk' - try: - run_with_progress(lambda: session(obs, time.time_ns()).result(), '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 23fbc9965..ff6e59213 100644 --- a/positronic/offboard/tests/conftest.py +++ b/positronic/offboard/tests/conftest.py @@ -3,14 +3,14 @@ import threading import time from collections.abc import Callable, Generator, Mapping +from contextlib import ExitStack, contextmanager from typing import Any -from unittest.mock import MagicMock import pytest import uvicorn from positronic.offboard.server import PolicyServer -from positronic.policy import AnySession, ChunkSession, Done, Policy, Session +from positronic.policy import INFER, Policy, Session from positronic.policy.executor import Executor from positronic.policy.layers import ChunkPlayer from positronic.policy.spec import ModelSource, PolicySource, remote @@ -58,29 +58,35 @@ async def _run(): @pytest.fixture -def open_session() -> Generator[Callable[..., tuple[AnySession, 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) -> tuple[AnySession, Executor]: - runtimes.append(Executor(policy.functions)) - return policy.new_session(None, runtimes[-1]), runtimes[-1] + def make(policy: Policy) -> Mapping[str, Callable[..., Any]]: + return closing.enter_context(policy.episode()) 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 +@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]: + rt = Executor(closing.enter_context(policy.episode())) + closing.callback(rt.close) + session = policy.new_session(rt) + closing.callback(session.close) + return session, rt -def round_trip(session: ChunkSession, rt: Executor, obs, time_ns: int = 0) -> list[dict] | dict | None: - """What ``session`` answers for ``obs``, once its round trip has come back.""" - answer = session(obs, time_ns) - rt.wait(ANSWER_SEC) - assert not rt.in_flight, 'the round-trip never came back' - return answer.result() + yield make + closing.close() + + +# How long a round trip against a local server may take before a test calls it lost. +ANSWER_SEC = 5.0 def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> Mapping[str, Any]: @@ -97,24 +103,39 @@ def played_round_trip(session: Session, rt: Executor, obs, time_ns: int = 0) -> return session(obs, time_ns)[0] -def _make_mock_policy(action, meta): - """Create a mock policy with session-based API.""" - session = MagicMock() - session.return_value = Done(action) - session.meta = meta - session.close = MagicMock() +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 _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): @@ -138,21 +159,20 @@ 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: @@ -164,7 +184,7 @@ def inference_server(start_server: StartServer, mock_policy: MagicMock) -> tuple @pytest.fixture def multi_policy_server( - start_server: StartServer, mock_policy_registry: dict[str, MagicMock] -) -> tuple[str, int, dict[str, MagicMock]]: + 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 459795b37..899026a7d 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -11,8 +11,8 @@ 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 played_round_trip, 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 ChunkPlayer from positronic.policy.remote import _prepare_obs @@ -276,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}, @@ -284,67 +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 + assert open_episode(endpoint)[INFER]({}) is None -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() + with pytest.raises(NotImplementedError, match='ChunkPlayer'): + endpoint.new_session(MagicMock()) -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): @@ -353,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 @@ -361,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] @@ -385,20 +379,20 @@ 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) @@ -409,23 +403,21 @@ def test_declared_stack_built_at_session_open(open_session): 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) @@ -456,7 +448,7 @@ def test_remote_policy_lifecycle(inference_server, mock_policy, open_session): 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' @@ -472,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 089c1943a..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 @@ -18,8 +19,7 @@ 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 ANSWER_SEC -from positronic.policy import ChunkSession, Codec, Policy, RemotePolicy -from positronic.policy.base import Runtime +from positronic.policy import INFER, Codec, Policy, RemotePolicy from positronic.policy.codec import ActionTimestamp from positronic.policy.layers import ChunkPlayer, TemporalStack from positronic.policy.spec import ModelSource, PolicySource, inline, remote @@ -65,7 +65,7 @@ def test_full_inference_cycle(stub_server): 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() @@ -199,24 +199,23 @@ 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): @@ -238,27 +237,12 @@ def test_pipeline_with_no_rig_side_half_refused_at_startup(make_mock_policy): PolicyServer(remote | _StubSource(stub)) -_INFER = 'infer' - - -class _ScriptedSession(ChunkSession): - def __init__(self, rt: Runtime): - self._rt = rt - - def __call__(self, obs, time_ns): - return self._rt.fns[_INFER](obs) - - class _ScriptedPolicy(Policy): - """Deterministic base policy: every session serves the same untimestamped chunk from its runtime.""" - - def new_session(self, context=None, rt=None) -> ChunkSession: - 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): diff --git a/positronic/policy/__init__.py b/positronic/policy/__init__.py index d97c83f72..03acf0b13 100644 --- a/positronic/policy/__init__.py +++ b/positronic/policy/__init__.py @@ -1,28 +1,17 @@ -from .base import ( - Answer, - AnySession, - ChunkSession, - DelegatingChunkSession, - DelegatingPolicy, - DelegatingSession, - Done, - 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', - 'ChunkSession', - 'AnySession', + 'Runtime', 'DelegatingPolicy', 'DelegatingSession', - 'DelegatingChunkSession', 'Layer', + 'ChunkLayer', 'RemotePolicy', 'Codec', 'ActionTimestamp', @@ -31,5 +20,4 @@ 'is_action', 'Recorder', 'Answer', - 'Done', ] diff --git a/positronic/policy/base.py b/positronic/policy/base.py index e584b6a89..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 typing import Any, ClassVar, final +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.""" @@ -26,94 +32,40 @@ def done(self) -> bool: ... def result(self) -> Any: """What the function returned. Raises what the function raised, or ``NotAnswered`` before it answers.""" - @final - def map(self, fn: Callable[[Any], Any]) -> Answer: - """This answer with ``fn`` applied to what it carries, at the moment it is read.""" - return _Mapped(self, fn) - - -class _Mapped(Answer): - def __init__(self, inner: Answer, fn: Callable[[Any], Any]): - self._inner = inner - self._fn = fn - - def done(self) -> bool: - return self._inner.done() - - def result(self) -> Any: - return self._fn(self._inner.result()) - - -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 - # A call to an ``Fn`` starts the work and returns at once. Fn = Callable[..., Answer] 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 _SessionBase: - """Per-episode state a policy opens: what the framework holds, cancels and closes. +class Session(ABC): + """Per-episode control session, created by ``Policy.new_session``. - ``Session`` and ``ChunkSession`` add the call. Both are opened and ended the same way, so the - lifecycle lives here. - """ + 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:: - @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. ``ChunkPlayer``) should reset so the next call - triggers a fresh inference. Override and propagate via ``super().cancel()``. - """ - return None - - def close(self): - """End this session and release per-episode resources.""" - return None - - -class Session(_SessionBase, ABC): - """Per-episode control session. Created by ``Policy.new_session()``. - - Sessions hold per-episode state (the chunk being played, latency tracking) and are what a caller - controls a robot through. Call the session like a function:: - - session = policy.new_session(context) + session = policy.new_session(rt) commands, resume_at_ns = session(obs, time_ns) - **Plain-data contract**: sessions accept and return only plain data - (dicts, lists, numpy arrays, scalars). No tensors or custom objects. + **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) -> 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.functions``. + 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. @@ -121,31 +73,13 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> tuple[Mapping[str, A earlier. A session that waits for work it cannot time names a poll period of its own. """ + def cancel(self): + """Drop what is in flight, so the next call plans afresh. Propagate with ``super().cancel()``.""" + return None -class ChunkSession(_SessionBase, ABC): - """Session that answers action chunks: what a ``ChunkPlayer`` holds and plays. - - A chunk is 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. One action - may come back bare, and the player wraps it. - - TODO(#661): the chunk becomes the answer of a served function, and this class goes with the - session-serving wire. - """ - - @abstractmethod - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> Answer: - """Start the work for one chunk, and answer the caller's handle on it. - - ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. - The caller decides how many calls it keeps in flight, and reads each handle when it wants the - chunk. A session that runs the work inside the call answers a ``Done``. - """ - - -# What a layer wraps: a session on either side of a ``ChunkPlayer``. -# TODO(#661): the chunk side goes, and this alias with it. -AnySession = Session | ChunkSession + def close(self): + """End this session and release what it holds.""" + return None class DelegatingSession(Session): @@ -157,34 +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() - - def close(self): - self._inner.close() - - -class DelegatingChunkSession(ChunkSession): - """Chunk session that delegates all methods to an inner ``ChunkSession``. - - The body repeats ``DelegatingSession`` because the two answer different calls, and each one's ``_inner`` - is the side it forwards to. TODO(#661): it goes with ``ChunkSession``. - """ - - def __init__(self, inner: ChunkSession): - self._inner = inner - - 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() @@ -192,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) -> AnySession: - """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.).""" @@ -230,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): @@ -261,21 +164,26 @@ class Layer: **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: """Apply this layer to a policy. Default: wrap every session it creates via ``make_session``.""" return _LayerPolicy(policy, self) - def make_session(self, inner: Any) -> AnySession: - """Make this layer's session around ``inner``, on the side of a ``ChunkPlayer`` the layer sits on.""" + def make_session(self, inner: Session) -> Session: + """Make this layer's session around ``inner``.""" raise NotImplementedError('Override make_session or wrap') # The name this layer travels under, set by every deliverable subclass. ``WIRE_LAYERS`` is keyed by # 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. @@ -300,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.""" @@ -325,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 51008dfc4..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, Answer, ChunkSession, DelegatingChunkSession, Layer, _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: ChunkSession) -> ChunkSession: - 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(DelegatingChunkSession): - """Session wrapped with a codec: encodes observations, decodes actions.""" - - def __init__(self, inner: ChunkSession, codec: 'Codec'): - super().__init__(inner) - self._codec = codec - - def __call__(self, obs, time_ns): - answer = self._inner(self._codec.encode(obs), time_ns) - # A codec decodes a chunk, so it belongs under the player. Above one it would be handed commands - # paired with a resume time, and each codec would drop or corrupt them in its own way. - assert isinstance(answer, Answer), 'compose a Codec under the ChunkPlayer, not above it' - return answer.map(self._codec.decode) - - @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.""" diff --git a/positronic/policy/executor.py b/positronic/policy/executor.py index fe3a1c0b1..daab032f1 100644 --- a/positronic/policy/executor.py +++ b/positronic/policy/executor.py @@ -9,17 +9,7 @@ from functools import partial from typing import Any -from positronic.policy.base import ( - Answer, - ChunkSession, - DelegatingChunkSession, - DelegatingPolicy, - Fn, - NotAnswered, - Policy, - Runtime, - Session, -) +from positronic.policy.base import Answer, Fn, NotAnswered, Runtime class Executor(Runtime): @@ -111,50 +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(DelegatingChunkSession): - def __init__(self, inner: ChunkSession, rt: Executor): - super().__init__(inner) - self._rt = rt - - def __call__(self, obs: Mapping[str, Any], time_ns: int) -> Answer: - answer = self._inner(obs, time_ns) - self._rt.wait() - return answer - - 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) -> ChunkSession: - assert rt is None, 'a blocking policy serves its own functions; nothing above it runs them' - own = Executor(self._inner.functions) - inner = None - try: - inner = self._inner.new_session(context, own) - # Waiting out the work is a question about a chunk, so nothing above a ``ChunkPlayer`` blocks. - assert not isinstance(inner, Session), f'{type(inner).__name__} answers commands, not a chunk' - except BaseException: - own.close() - if inner is not None: - inner.close() - raise - return _BlockingPolicy._Session(inner, own) - - @property - def functions(self) -> Mapping[str, Callable[..., Any]]: - return {} - - -def blocking(policy: Policy) -> Policy: - """``policy`` with its heavy work waited out: the handle a call answers is already done. - - 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 6ac1cb2e9..7a35cf20b 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,3 +1,4 @@ +import contextlib import time from collections.abc import Generator, Iterator, Mapping from typing import Any @@ -11,7 +12,7 @@ from positronic.dataset.serializers import expand_suffixed from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Task -from positronic.policy.base import Policy, Session +from positronic.policy.base import Policy from positronic.policy.executor import Executor from positronic.utils import flatten_dict, frozen_view @@ -24,31 +25,22 @@ 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) - session = None + self._closing = contextlib.ExitStack() try: - session = policy.new_session(context, self._runtime) - # A stack that answers chunks has no player in it, so nothing turns them into commands. - assert isinstance(session, Session), ( - f'{type(session).__name__} answers a chunk; a policy played here needs a ChunkPlayer in its stack' - ) + 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() - if session is not None: - session.close() + self._closing.close() raise - self._session = session - - @property - def meta(self) -> dict[str, Any]: - return self._session.meta @staticmethod def _owned(obs: dict[str, Any]) -> dict[str, Any]: @@ -81,12 +73,11 @@ def wait(self, should_stop: pimm.SignalReceiver[bool]) -> None: 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: @@ -219,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 @@ -327,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 diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 50127557a..be0c253a3 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -1,8 +1,8 @@ """Composable policy layers — chunk playback, fault handling and temporal frame stacking. Layers are serving-time concerns wrapped around a policy with ``|`` (left is outermost). ``ChunkPlayer`` -converts a chunk into the commands a caller runs; the layers around it read time from the observation -(``obs_time_ns``) or from the call's ``time_ns``, both in nanoseconds. +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 @@ -12,7 +12,7 @@ from positronic import keys from positronic.drivers.roboarm import RobotStatus -from positronic.policy.base import Answer, ChunkSession, 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. @@ -66,21 +66,32 @@ def to_spec(self): class ChunkPlayer(Layer): - """Hold the chunk the sessions below answer, anchor it to the call that received it, and emit each - waypoint at its own time.""" + """Hold the chunk ``INFER`` answers, anchor it to the call that received it, and emit each waypoint at + its own time. + + 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 = 'chunk_player' + PLAYS_CHUNKS = True POLL_SEC = 0.01 + class _Policy(DelegatingPolicy): + """The player is the session of a chunk policy, so nothing below it opens one.""" + + 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, inner: ChunkSession): - self._inner = inner + def __init__(self, infer: Fn): + self._infer = infer self._waypoints: deque[ChunkPlayer._Session._Waypoint] = deque() - # The one call this layer keeps in flight, and whether a ``cancel`` has orphaned the chunk it + # 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 @@ -89,7 +100,7 @@ def __call__(self, obs, time_ns): """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._inner(obs, time_ns) + self._answer = self._infer(obs) if self._answer.done(): chunk = self._take() if chunk is not None: @@ -107,7 +118,7 @@ def _take(self) -> list[dict[str, Any]] | dict[str, Any] | None: 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 layer made brings a chunk back' + 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() @@ -116,7 +127,7 @@ def _take(self) -> list[dict[str, Any]] | dict[str, Any] | None: 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 session may answer a bare dict, and a no-codec path may omit ``timestamp`` + 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. """ @@ -126,9 +137,9 @@ def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> N 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 sessions below are timing actions against a clock of their own' + f'the work below is timing actions against a clock of its own' ) - # The single explicit seconds->ns seam: the sessions below time actions in float seconds, and + # 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. @@ -140,21 +151,13 @@ def _load(self, chunk: list[dict[str, Any]] | dict[str, Any], time_ns: int) -> N for action in chunk ) - @property - def meta(self): - return self._inner.meta - def cancel(self): 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 - self._inner.cancel() - - def close(self): - self._inner.close() - def make_session(self, inner: ChunkSession) -> Session: - return ChunkPlayer._Session(inner) + def wrap(self, policy: Policy) -> Policy: + return ChunkPlayer._Policy(policy) def to_spec(self): return {'name': self.WIRE_NAME} diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index 67c166a4f..990485970 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -46,8 +46,10 @@ 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 @@ -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 Answer, AnySession, ChunkSession, 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,49 +367,33 @@ def _release_stream(self) -> None: self._live -= 1 -class _RecordingTap(Layer): - """A named tap. Wraps a single session to log its observations and actions.""" - - def __init__(self, rec: Recorder, name: str): - self._rec = rec - self._name = name - - def make_session(self, inner: AnySession) -> AnySession: - stream = self._rec._open_stream() - # The side a tap sits on is the side its inner is on, so a chain of taps carries it outwards. - tap = _CommandTapSession if isinstance(inner, Session) else _ChunkTapSession - return tap(inner, self._rec, self._name, stream) - - class _Frame(NamedTuple): - """What one call is logged against: its observation, the rerun timelines it set, and its step. - - A chunk is logged when the caller reads it, so its frame outlives the call that started the work. - """ + """What one call is logged against: its observation, the rerun timelines it set, and its step.""" obs: Mapping[str, Any] timelines: dict[str, Any] step: int -class _RecordingTapSession: +class _TapLog: """Logs the observation flowing down and the actions flowing up at one point. - ``_CommandTapSession`` and ``_ChunkTapSession`` add the call, one per side of a ``ChunkPlayer``. + 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: AnySession, rec: Recorder, name: str, stream: rr.RecordingStream): - self._inner = 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, frame: '_Frame') -> None: + def _set_timelines(self, frame: _Frame) -> None: for timeline, value in frame.timelines.items(): set_timeline_time(timeline, value) set_timeline_sequence('step', frame.step) @@ -416,7 +413,7 @@ def _log(self, prefix: str, data: Mapping[str, Any]) -> None: log_numeric_series(path, num) self._rec._numeric_paths.append(path) - def _log_action_chunk(self, prefix: str, actions: list[dict], frame: '_Frame') -> 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)] @@ -467,10 +464,6 @@ def _log_action_chunk(self, prefix: str, actions: list[dict], frame: '_Frame') - _log_action_series(f'{prefix}/series/{suffix}', field_arr, group_h, base_ns, names=None) self._rec._series_paths.append(f'{prefix}/series/{suffix}') - def _actions(self, answer) -> list[dict]: - """The actions the chunk logger plots out of what the inner answered.""" - raise NotImplementedError - def _send_blueprint(self) -> None: rec = self._rec bp = rec._blueprint or _build_blueprint( @@ -479,65 +472,92 @@ def _send_blueprint(self) -> None: if bp is not None: rr.send_blueprint(bp) - def _tapped(self, obs, time_ns) -> tuple[Any, _Frame]: - """What the inner answered, and the frame this call is logged against.""" - 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(rec._timeline_values), self._step) + frame = _Frame(obs, dict(values), self._step) with self._stream: self._set_timelines(frame) self._log(self._name, obs) - return self._inner(obs, time_ns), frame + yield frame finally: - rec._depth -= 1 - if rec._depth == 0: - rec._timeline_values = {} + if token is not None: + _TIMELINE_VALUES.reset(token) - def _log_actions(self, actions: list[dict], frame: _Frame) -> None: + 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 outermost tap.""" - if self._rec._depth == 0 and self._step == 0: + 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 cancel(self): - self._inner.cancel() + +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): self._inner.close() - self._rec._release_stream() + 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 _CommandTapSession(_RecordingTapSession, Session): +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 __call__(self, obs, time_ns) -> tuple[Mapping[str, Any], int]: - answer, frame = self._tapped(obs, time_ns) - commands, _ = answer - self._log_actions([dict(commands)] if commands else [], frame) - self._end_step() - return answer - + def make_session(self, inner: Session) -> Session: + return _CommandTapSession(inner, _TapLog(self._rec, self._name)) -class _ChunkTapSession(_RecordingTapSession, ChunkSession): - """A tap under a ``ChunkPlayer``: it logs the chunk a call brings back, when the caller reads it.""" - def __call__(self, obs, time_ns) -> Answer: - answer, frame = self._tapped(obs, time_ns) - self._end_step() - return answer.map(partial(self._log_chunk, frame)) +class _ChunkTap(_Tap, ChunkLayer): + """A tap under a ``ChunkPlayer``: it logs the observation the model sees and the chunk it answers.""" - def _log_chunk(self, frame: _Frame, chunk): - self._log_actions([dict(chunk)] if isinstance(chunk, Mapping) else (chunk or []), frame) - return chunk + @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 bb714e4ee..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, AnySession, ChunkSession, Layer, Policy, Runtime +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,33 +52,6 @@ def round_trip( telemetry.record_span(telemetry_keys.SPAN_POLICY_INFER, infer_start_ns, time.time_ns()) -class RemoteSession(ChunkSession): - """Per-episode session that forwards observations to a remote inference server. - - ``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 - - def __call__(self, obs: cabc.Mapping[str, Any], time_ns: int) -> Answer: - """One round trip to the server, and the caller's handle on the trajectory it brings back. - - A server answer of one action becomes a 1-element list, which is the form a chunk takes. - """ - answer = self._rt.fns[INFER](self._session, obs, self._compress_images) - return answer.map(lambda result: [result] if isinstance(result, dict) else result) - - @property - def meta(self) -> dict[str, Any]: - return flatten_dict({keys.TYPE: 'remote', keys.SERVER: self._session.metadata}) - - def close(self): - self._session.close() - - class _Endpoint(Policy): """The wire connection to one inference server: sessions forward observations under the border's settings. @@ -100,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]: @@ -170,12 +141,11 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, rt=None) -> AnySession: - 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/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 3550b7f79..1530eb555 100644 --- a/positronic/policy/tests/test_executor.py +++ b/positronic/policy/tests/test_executor.py @@ -10,17 +10,8 @@ import pytest -from positronic.policy.base import ( - Answer, - ChunkSession, - DelegatingChunkSession, - Done, - 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 @@ -243,153 +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(ChunkSession): - def __init__(self): - self.calls = 0 - - def __call__(self, obs, time_ns): - self.calls += 1 - return Done([{'action': obs}]) - - def __init__(self): - self.session = _PlainPolicy._Session() - - def new_session(self, context=None, rt=None) -> ChunkSession: - return self.session - - -_ECHO = 'echo' - - -class _EchoPolicy(Policy): - """Serves ``echo``: its session hands the caller the handle on one call of it.""" - - class _Session(ChunkSession): - def __init__(self, rt): - self._rt = rt - self.calls = 0 - - def __call__(self, obs, time_ns): - self.calls += 1 - return self._rt.fns[_ECHO](obs) - - def __init__(self): - self.session: _EchoPolicy._Session - - def new_session(self, context=None, rt=None) -> ChunkSession: - assert rt is not None - self.session = _EchoPolicy._Session(rt) - 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(DelegatingChunkSession): - def __init__(self, inner: ChunkSession, 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) -> ChunkSession: - 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_commands_is_refused_and_closed(self): - """Nothing above a ``ChunkPlayer`` has work to wait out, and the session opened to find that out - holds a connection of its own.""" - - class _Commanding(Session): - def __init__(self): - self.closed = False - - def __call__(self, obs, time_ns): - return {}, time_ns - - def close(self): - self.closed = True - - session = _Commanding() - - class _Commander(Policy): - def new_session(self, context=None, rt=None): - return session - - with pytest.raises(AssertionError, match='answers commands'): - blocking(_Commander()).new_session() - assert session.closed, 'the session it opened to check was dropped without closing' - - 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).result() == [{'action': {'x': 1}}] - assert policy.session.calls == 1 - - def test_the_handle_a_call_answers_is_already_done(self, opened): - """The wait is what ``blocking`` adds: the caller reads the chunk without asking a second time.""" - policy = _EchoPolicy() - - answer = opened(blocking(policy))({'x': 1}, 0.0) - - assert answer.done() - assert answer.result() == {'x': 1} - 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() - - assert opened(layer.wrap(blocking(policy)))({'x': 1}, 0.0).result() == {'x': 1} - assert (layer.calls, policy.session.calls) == (1, 1) - - 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()).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() - 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 7f5c16fd6..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,7 +42,7 @@ 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 Answer, ChunkSession, DelegatingChunkSession, DelegatingPolicy, Done, Policy +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 ChunkPlayer, StopOnFault @@ -64,18 +65,6 @@ CAPTURED_SIGNALS = (keys.EE_POSE, keys.JOINTS, keys.GRIP) -class _ScriptedSession(ChunkSession): - def __call__(self, obs, time_ns): - current = np.asarray(obs[keys.EE_POSE][:3], dtype=np.float32) - delta = TARGET_POS - current - chunk = [] - for i in range(10): - step = current + delta * 0.5 * ((i + 1) / 10.0) - pose = Transform3D(translation=step.astype(np.float32), rotation=Rotation.identity) - chunk.append({keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': round(0.50 + 0.01 * i, 4)}) - return Done(chunk) - - class ScriptedProportionalPolicy(Policy): """Pure proportional controller toward ``TARGET_POS``. @@ -83,14 +72,25 @@ class ScriptedProportionalPolicy(Policy): clock, no images. The codec stamps and truncates; ``ChunkPlayer`` anchors the chunk and plays it. """ - def new_session(self, context=None, rt=None): - return _ScriptedSession() + @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 = [] + for i in range(10): + step = current + delta * 0.5 * ((i + 1) / 10.0) + pose = Transform3D(translation=step.astype(np.float32), rotation=Rotation.identity) + chunk.append({keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': round(0.50 + 0.01 * i, 4)}) + return chunk class _SimulatedLatency(DelegatingPolicy): - """A fixed inference latency in world time: every chunk is stamped for, and answers at, ``latency_sec`` - after the call that produced it, whatever the machine took. Sits under the player, so nothing below it - sees an observation while a chunk is in flight.""" + """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.""" def __init__(self, inner: Policy, latency_sec: float, clock: pimm.Clock): super().__init__(inner) @@ -98,7 +98,7 @@ def __init__(self, inner: Policy, latency_sec: float, clock: pimm.Clock): self._clock = clock class _Held(Answer): - """The chunk under it, held back until the world reaches ``release_at_ns``.""" + """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 @@ -111,21 +111,22 @@ def done(self) -> bool: def result(self): return self._inner.result() - class _Session(DelegatingChunkSession): - def __init__(self, inner: ChunkSession, latency_ns: int, clock: pimm.Clock): - 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._clock = clock - def __call__(self, obs, time_ns): - # The chunk answers at ``latency_ns``, so the sessions below time it from that instant. - release_at_ns = time_ns + self._latency_ns - return _SimulatedLatency._Held(self._inner(obs, release_at_ns), release_at_ns, self._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, context=None, rt=None) -> ChunkSession: - inner = self._inner.new_session(context, rt) - assert isinstance(inner, ChunkSession) - return _SimulatedLatency._Session(inner, self._latency_ns, self._clock) + def new_session(self, rt): + return self._inner.new_session(_SimulatedLatency._Runtime(rt, self._latency_ns, self._clock)) class FakeRobot(pimm.ControlSystem): @@ -212,7 +213,9 @@ def _run_pipeline(tmp_path: Path) -> dict: simulated=True, ) harness = Harness( - (StopOnFault() | ChunkPlayer()).wrap(_SimulatedLatency(policy, INFERENCE_LATENCY_S, world.clock)), + # 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) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 1024f6b53..679fa9340 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -22,11 +22,11 @@ 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 Answer, ChunkSession, DelegatingSession, Done, 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 MAX_ROUND_SEC, MIN_ROUND_SEC, WAIT_PERIOD_SEC, Harness, _EpisodeInference from positronic.policy.layers import ChunkPlayer, StopOnFault -from positronic.policy.remote import INFER, RemoteSession, round_trip +from positronic.policy.remote import round_trip from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler WAIT_PERIOD_NS = round(WAIT_PERIOD_SEC * 1e9) @@ -80,17 +80,6 @@ def make_embodiment( ) -class _SpySession(ChunkSession): - def __init__(self, policy): - self._policy = policy - - def __call__(self, obs, time_ns): - self._policy.last_obs = obs - return Done([ - {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: @@ -102,27 +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(ChunkSession): - def __init__(self, policy): - self._policy = policy - self._meta = dict(policy._meta) + yield {INFER: self._infer} - def __call__(self, obs, time_ns): - self._policy.last_obs = obs - self._policy.observations.append(obs) - return Done([ - {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.target_grip, 'timestamp': 0.0} - ]) - - @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): @@ -149,27 +126,16 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, rt=None) -> ChunkSession: + @contextmanager + def episode(self, context=None): self.reset_calls += 1 self.last_reset_context = context - return _StubSession(self) - - -class _ChunkSession(ChunkSession): - def __init__(self, policy): - self._policy = policy + yield {INFER: self._infer} - def __call__(self, obs, time_ns): - self._policy.counter += 1 - dt = 0.005 - return Done([ - { - 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): @@ -179,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 @@ -207,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``. """ @@ -217,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]]: @@ -600,20 +570,14 @@ 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(ChunkSession): - def __init__(self, command): - self._command = command - - def __call__(self, obs, time_ns): - return Done([{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): @@ -697,20 +661,9 @@ def test_an_uncharged_wait_ends_when_the_world_comes_down(world): never_answers = threading.Event() class _HangingPolicy(Policy): - class _Session(ChunkSession): - def __init__(self, rt): - self._rt = rt - - def __call__(self, obs, time_ns): - return self._rt.fns[INFER]() - - def new_session(self, context=None, rt=None): - assert rt is not None - return _HangingPolicy._Session(rt) - - @property - def functions(self): - return {INFER: never_answers.wait} + @contextmanager + def episode(self, context=None): + yield {INFER: lambda obs: never_answers.wait()} inference = _EpisodeInference(ChunkPlayer().wrap(_HangingPolicy()), {}, charges_wall_time=False, clock=world.clock) try: @@ -726,13 +679,10 @@ 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(ChunkSession): - 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(ChunkPlayer().wrap(_RaisingPolicy()), make_embodiment()) p = _pair_all(world, harness) @@ -1169,13 +1119,10 @@ 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(ChunkSession): - def __call__(self, obs, time_ns): - raise RuntimeError('inference boom') - class _BoomPolicy(Policy): - def new_session(self, context=None, rt=None): - return _BoomSession() + @contextmanager + def episode(self, context=None): + yield {INFER: _boom} harness = Harness(ChunkPlayer().wrap(_BoomPolicy()), make_embodiment()) p = _pair_all(world, harness) @@ -1396,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) @@ -1519,13 +1468,10 @@ 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(ChunkSession): - def __call__(self, obs, time_ns): - return Done([]) - 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(ChunkPlayer().wrap(EmptyChunkPolicy()), make_embodiment()) cmd_recorder = RecordingEmitter() @@ -1600,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' @@ -1782,7 +1728,7 @@ 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).""" + (a real wire round trip over a fake inference session).""" policy = ChunkPlayer().wrap(RemoteStubPolicy()) harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) @@ -1996,7 +1942,7 @@ def test_a_policy_that_answers_chunks_refuses_to_open(world): 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)) - with pytest.raises(AssertionError, match='needs a ChunkPlayer'): + with pytest.raises(NotImplementedError, match='put a ChunkPlayer above it'): drive_scheduler(scheduler, steps=20) assert call.done() @@ -2011,8 +1957,8 @@ class _ReplanEarly(Layer): class _Session(Session): POLL_SEC = 0.001 - def __init__(self, inner: ChunkSession): - self._inner = inner + 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 @@ -2029,7 +1975,7 @@ def __call__(self, obs, time_ns): def _replan(self, obs, time_ns: int) -> None: if self._answer is None: - self._answer = self._inner(obs, time_ns) + 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 @@ -2042,8 +1988,14 @@ def _replan(self, obs, time_ns: int) -> None: 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: ChunkSession): - 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): @@ -2366,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): @@ -2414,26 +2366,26 @@ 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(ChunkSession): + 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) if self._calls == 1: - return Done([ + return [ {keys.ROBOT_COMMAND: command, keys.TARGET_GRIP: 0.5, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10) - ]) - return Done([{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() + ] + return [{keys.ROBOT_COMMAND: command, keys.ACTION_TIMESTAMP: i * 0.01} for i in range(10)] harness = Harness(ChunkPlayer().wrap(_GripThenArmPolicy()), make_embodiment()) grip_recorder = RecordingEmitter() diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 0117bbf97..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, 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 Answer, ChunkSession, Done, Layer, Policy, Session +from positronic.policy.base import INFER, Answer, Fn, Layer, Policy, Runtime, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -25,46 +27,68 @@ ) from positronic.policy.layers import ChunkPlayer, StopOnFault, TemporalStack from positronic.policy.observation import ObservationCodec +from positronic.policy.tests.conftest import Done -class _Pending(Answer): +class _Const: + """Answers the same chunk to every call.""" + + def __init__(self, actions): + self._actions = actions + self.call_count = 0 + + 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.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=None, failure: BaseException | None = None): + def __init__(self, value: Any = None, failure: BaseException | None = None): self._value = value self._failure = failure self._done = False - def answer(self): + def answer(self) -> None: self._done = True - def done(self): + def done(self) -> bool: return self._done - def result(self): + def result(self) -> Any: if self._failure is not None: raise self._failure return self._value -class _ConstSession(ChunkSession): - def __init__(self, actions): - self._actions = actions - self.call_count = 0 +class _FnRuntime(Runtime): + """Serves one ``INFER``, exactly as the test wrote it.""" - def __call__(self, obs, time_ns): - self.call_count += 1 - return Done(self._actions) + def __init__(self, infer: Fn): + self._fns: Mapping[str, Fn] = {INFER: infer} + @property + def fns(self) -> Mapping[str, Fn]: + return self._fns -class _ConstPolicy(Policy): - def __init__(self, actions): - self._actions = actions - self._session = _ConstSession(actions) - def new_session(self, context=None, rt=None) -> _ConstSession: - self._session = _ConstSession(self._actions) - return self._session +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): @@ -151,26 +175,26 @@ def test_a_status_no_arm_answers_to_raises(self): with pytest.raises(ValueError): session({keys.OBS_TIME_NS: 0, keys.ROBOT_STATUS: 99}, 0) - def test_recovery_plans_afresh_instead_of_resuming(self): + 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() | ChunkPlayer()).wrap(inner).new_session() + session = open_session((StopOnFault() | ChunkPlayer()).wrap(inner)) 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._session.call_count == 2 + assert inner.infer.call_count == 2 -class _ScriptedSession(ChunkSession): +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) @@ -178,9 +202,9 @@ def __call__(self, obs, time_ns): 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 = _ScriptedSession([pending]) - session = ChunkPlayer().make_session(inner) + pending = Pending([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) + inner = _Scripted([pending]) + session = player(inner) 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)) @@ -191,9 +215,9 @@ def test_a_call_that_has_not_answered_leaves_the_player_holding(self): 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 = _ScriptedSession([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) - session = ChunkPlayer().make_session(inner) + 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) @@ -207,8 +231,8 @@ def test_a_cancel_drops_the_chunk_of_the_call_in_flight(self): 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 = ChunkPlayer().make_session(_ScriptedSession([pending])) + pending = Pending(failure=TimeoutError('server stalled')) + session = player(_Scripted([pending])) session(_obs(), int(1e9)) pending.answer() @@ -220,9 +244,9 @@ def test_a_cancelled_call_still_raises_what_it_failed_with(self): 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 = _ScriptedSession([pending, Done([{'v': 2, keys.ACTION_TIMESTAMP: 0.0}])]) - session = ChunkPlayer().make_session(inner) + 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() @@ -232,16 +256,16 @@ def test_a_cancel_ends_with_the_call_it_was_made_against(self): 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): + 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 = ChunkPlayer().wrap(inner).new_session() + 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._session.call_count == 1 + assert inner.infer.call_count == 1 - def test_waypoints_due_together_keep_every_channel_they_name(self): + 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}, @@ -249,7 +273,7 @@ def test_waypoints_due_together_keep_every_channel_they_name(self): {'arm': 2, keys.ACTION_TIMESTAMP: 0.02}, {keys.ACTION_TIMESTAMP: 0.03}, ]) - session = ChunkPlayer().wrap(inner).new_session() + 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)) @@ -257,84 +281,81 @@ def test_waypoints_due_together_keep_every_channel_they_name(self): 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 = _ScriptedSession([ + 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 = ChunkPlayer().make_session(inner) + 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): + 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 = ChunkPlayer().wrap(inner).new_session() + 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): + 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 = ChunkPlayer().wrap(inner).new_session() + 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._session.call_count == 2 + assert inner.infer.call_count == 2 - def test_a_single_action_is_drained_by_the_call_that_loads_it(self): + 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 = ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])).new_session() + 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): + 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 = ChunkPlayer().make_session(inner.new_session()) + 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): + 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 = ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 1.77e18 / 1e9}])).new_session() + session = open_session(ChunkPlayer().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 1.77e18 / 1e9}]))) - with pytest.raises(ValueError, match='clock of their own'): + 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): + 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() + 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): - """A codec composes above a player, and refuses the commands it would silently corrupt.""" + """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) - session = pipeline.wrap(_ConstPolicy([{'action': 'test'}])).new_session() - with pytest.raises(AssertionError, match='under the ChunkPlayer'): - session(_obs(), int(1e9)) + 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 = ChunkPlayer() | codec assert isinstance(pipeline, Layer) # 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. - policy = pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)])) - session = policy.new_session() + 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)) @@ -374,21 +395,19 @@ def test_parallel_frame_codecs_keep_the_frame_they_share(self): ) -class _CaptureSession(ChunkSession): +class _CapturePolicy(Policy): + """Records every observation its inference is given.""" + def __init__(self): self.seen = [] - def __call__(self, obs, time_ns): - self.seen.append(obs) - return Done([]) - + @contextmanager + def episode(self, context=None): + yield {INFER: self._infer} -class _CapturePolicy(Policy): - def __init__(self): - self.session = _CaptureSession() - - def new_session(self, context=None, rt=None): - return self.session + def _infer(self, obs): + self.seen.append(obs) + return [] def _stack_obs(now_sec, value): @@ -435,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() @@ -652,18 +671,17 @@ def test_pipe_composes_no_further(self): 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(ChunkPlayer() | spec.remote | ActionTimestamp(fps=10.0) | spec.PolicySource(inner)) assert isinstance(policy, Policy) - session = policy.new_session() + 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(ChunkPlayer() | spec.PolicySource(inner)) - session = policy.new_session() + 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): diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index d7191fa0a..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 Answer, ChunkSession, Done, Policy +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(ChunkSession): - def __init__(self, result): - self._result = result - - def __call__(self, obs, time_ns): - return Done(self._result) - - -class _ChunkPolicy(Policy): - def __init__(self, actions: list[dict]): - self._actions = actions - - def new_session(self, context=None, rt=None) -> ChunkSession: - return _FixedSession(list(self._actions)) +class _FixedPolicy(Policy): + """Answers the same chunk to every call.""" + def __init__(self, chunk): + self._chunk = chunk -class _SinglePolicy(Policy): - def new_session(self, context=None, rt=None) -> ChunkSession: - 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) -> ChunkSession: - return _FixedSession({}) +class _MetaPolicy(_FixedPolicy): + def __init__(self): + super().__init__({}) @property def meta(self): @@ -157,19 +149,18 @@ def meta(self): _T0_OBS = {obs_keys.OBS_TIME_NS: 0} -def _chunk(session, obs=_T0_OBS, time_ns: int = 0): - """The chunk a session under a player answers for ``obs``.""" - answer = session(obs, time_ns) - assert isinstance(answer, Answer), 'the session answers commands, not a chunk' - return answer.result() +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 = _chunk(policy.new_session()) + 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 @@ -178,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 = _chunk(policy.new_session()) + policy = codec.wrap(_FixedPolicy(list(actions))) + result = _chunk(policy) assert len(result) == 6 # 5 actions + timestamp sentinel @@ -187,16 +178,16 @@ 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 = _chunk(policy.new_session()) + 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 = _chunk(policy.new_session()) + 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) @@ -206,8 +197,8 @@ 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 = _chunk(policy.new_session()) + 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): @@ -257,8 +248,8 @@ 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 = _chunk(policy.new_session()) + 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 @@ -266,8 +257,8 @@ def test_action_timestamp_and_horizon_compose(): def test_single_action_has_zero_timestamp(): codec = ActionTiming(fps=15.0) - policy = codec.wrap(_SinglePolicy()) - result = _chunk(policy.new_session()) + 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 a6c8e05ad..7ebaecbe3 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -1,84 +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 Answer, ChunkSession, Done, Policy +from positronic.policy.base import INFER, Policy from positronic.policy.layers import ChunkPlayer from positronic.policy.recording import ( + _TIMELINE_VALUES, Recorder, _build_blueprint, - _ChunkTapSession, _command_field_arrays, _CommandTapSession, _flat_wire, - _RecordingTapSession, _squeeze_batch, _stack_numeric, ) -def _chunk(session, obs, time_ns: int = 0): - """The chunk a session under a player answers for ``obs``.""" - answer = session(obs, time_ns) - assert isinstance(answer, Answer), 'the session answers commands, not a chunk' - return answer.result() - - -class _TrackingSession(ChunkSession): - def __init__(self, actions, meta): - self._actions = actions - self._meta = meta - - def __call__(self, obs, time_ns): - return Done(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(ChunkSession): - """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 Done(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(): @@ -127,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) - _chunk(rec.tap('t').wrap(_TrackingPolicy(actions)).new_session(), {'x': 1.0, keys.WALL_TIME_NS: 1}) + _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' @@ -162,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 = _chunk(session, {'x': 1.0, keys.WALL_TIME_NS: 1_000_000}) - 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', @@ -195,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 @@ -216,42 +196,34 @@ 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 = _chunk(session, {'x': 1.0, keys.WALL_TIME_NS: 1}) + 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_an_empty_chunk(tmp_path): - class _EmptySession(ChunkSession): - def __call__(self, obs, time_ns): - return Done([]) - - class _EmptyPolicy(Policy): - def new_session(self, context=None, rt=None): - return _EmptySession() + rec = Recorder(tmp_path) + policy = rec.chunk_tap('t').wrap(_TrackingPolicy([])) - session = Recorder(tmp_path).tap('t').wrap(_EmptyPolicy()).new_session() - assert isinstance(session, _RecordingTapSession) - assert _chunk(session, {'x': 1.0}) == [] - assert session._step == 1 + assert _chunk(policy, {'x': 1.0}, times=2) == [] + assert rec._series_paths == [] 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}] - session = Recorder(tmp_path).tap('t').wrap(_TrackingPolicy(chunk)).new_session() - assert isinstance(session, _ChunkTapSession) + policy = Recorder(tmp_path).chunk_tap('t').wrap(_TrackingPolicy(chunk)) - assert session({'x': 1.0}, 0).result() == chunk + assert _chunk(policy, {'x': 1.0}) == chunk -def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path): +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) - session = (rec.tap('raw') | ChunkPlayer()).wrap(_TrackingPolicy(actions)).new_session() + session = open_session((rec.tap('raw') | ChunkPlayer()).wrap(_TrackingPolicy(actions))) assert isinstance(session, _CommandTapSession) assert session({'x': 1.0, keys.WALL_TIME_NS: 1}, 0) == ({'v': 1.0}, int(0.5e9)) @@ -261,25 +233,22 @@ def test_a_tap_above_the_player_logs_the_command_of_each_round(tmp_path): def test_two_taps_share_one_file_per_episode(tmp_path): rec = Recorder(tmp_path) - policy = (rec.tap('raw') | rec.tap('server')).wrap(_TrackingPolicy()) + policy = (rec.chunk_tap('raw') | rec.chunk_tap('server')).wrap(_TrackingPolicy()) - session = policy.new_session() - assert len(list(tmp_path.glob('*.rrd'))) == 1 - assert rec._live == 2 - - 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 @@ -289,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): @@ -317,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 8ca0b6e42..284e4240b 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -164,17 +164,18 @@ def main( rec = Recorder(pos3.sync(output_dir)) tapped = rec.tap(_TAP).wrap(policy) - runtime = Executor(policy.functions) - session = tapped.new_session({keys.TASK: task} if task else None, runtime) - meta = dict(session.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 session, whose socket the call still holds. - runtime.close() - session.close() + 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 diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 12418a62e..03bef768d 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -20,7 +20,7 @@ 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 ChunkSession, Done, Policy +from positronic.policy import INFER, Policy from positronic.policy.codec import ActionTimestamp from positronic.policy.layers import ChunkPlayer from positronic.policy.tests.test_harness import StubPolicy @@ -369,28 +369,23 @@ 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(ChunkSession): - def __init__(self, policy: _JointposChunks): - self._policy = policy - - def __call__(self, obs, time_ns): - self._policy.chunks += 1 - return Done([ - {keys.ROBOT_COMMAND: self._policy.command, 'target_grip': self._policy.chunks * 100.0 + i} - for i in range(self._policy.chunk_len) - ]) + def _infer(self, obs): + self.chunks += 1 + return [ + {keys.ROBOT_COMMAND: self.command, 'target_grip': self.chunks * 100.0 + i} for i in range(self.chunk_len) + ] @pytest.mark.timeout(60.0) 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 ``ChunkPlayer`` 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'] @@ -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/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index a7e922384..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 ChunkSession, Codec, Done, Layer, Policy +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(ChunkSession): - 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 Done([{'action': action_array}]) - return Done([{'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 467a03dee..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,7 +19,7 @@ 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 ChunkSession, Done, Policy +from positronic.policy import INFER, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -200,17 +202,14 @@ def stop(self): ########################################################################################### -class _Gr00tSession(ChunkSession): - 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 Done([{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): diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index 730eefa4e..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 ChunkSession, Done, Policy +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(ChunkSession): - 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) -> Done: - 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 Done([{'action': action}]) - return Done([{'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_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index bea4927f9..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,8 +14,7 @@ from positronic import keys from positronic.cfg import codecs -from positronic.policy import ChunkSession, Codec, Policy -from positronic.policy.base import Answer, Runtime +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 @@ -81,36 +80,15 @@ def _infer(policy: PreTrainedPolicy, device: str, obs: dict[str, Any]) -> list[d class LerobotPolicy(Policy): - _INFER = 'infer' - - class _Session(ChunkSession): - """Per-episode session that gives the model call to the runtime.""" - - def __init__(self, rt: Runtime, meta: dict[str, Any]): - self._rt = rt - self._meta = meta - - def __call__(self, obs: dict[str, Any], time_ns: int) -> Answer: - return self._rt.fns[LerobotPolicy._INFER](obs) - - @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]: diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index c2e4bf5d8..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 ChunkSession, Done, Policy +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(ChunkSession): - 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) -> Done: - # 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 Done([{'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) -> ChunkSession: - 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/openpi/server.py b/positronic/vendors/openpi/server.py index 74c5ddb9d..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,7 +16,7 @@ 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 ChunkSession, Codec, Done, Policy +from positronic.policy import INFER, Codec, Policy from positronic.policy.codec import ChangeEEFrame, RestrictImageSize from positronic.policy.layers import ChunkPlayer, StopOnFault from positronic.policy.spec import ModelSource, remote @@ -124,14 +126,9 @@ def stop(self): ########################################################################################### -class _OpenpiSession(ChunkSession): - def __init__(self, client: WebsocketClientPolicy): - self._client = client - - def __call__(self, obs, time_ns): - response = self._client.infer(obs) - actions = response['actions'] - return Done([{'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()