From eb124a7101e73b14a2d3907b00eabb105e74592e Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Mon, 10 Aug 2026 11:29:46 +0300 Subject: [PATCH 01/64] Play the policy's trajectory in the harness, one command per channel per round MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The wire between harness and driver carries the single command due now, not a future trajectory. `TrajectoryPlayer` moves above the wire into the harness, every driver executes the latest updated message and holds otherwise, and emission time is execution time — so the recorder taps command channels as plain per-sample signals, and `TrajectoryOverrideSerializer` and the self-timestamped-stream mechanism it needed are deleted. The session call runs on a one-worker executor, so the harness keeps playing while a model runs. `InferenceGate`, installed by the composition machinery below any `SchedulingWrapper`, is what charges that call the trial's inference latency: the wrapper resumes at the release instant and anchors there, so `_bump_schedule_end`, `_inference_delay` and the post-inference shift are deleted rather than ported. --- .basedpyright/baseline.json | 442 --------------- ARCHITECTURE.md | 29 +- pimm/core.py | 4 +- positronic/cli/eval/run.py | 4 +- positronic/data_collection.py | 6 +- positronic/dataset/ds_writer_agent.py | 90 +-- positronic/dataset/serializers.py | 24 +- .../dataset/tests/test_ds_writer_agent.py | 79 +-- positronic/drivers/gripper/dh.py | 13 +- positronic/drivers/gripper/robotiq.py | 14 +- positronic/drivers/roboarm/command.py | 68 +-- positronic/drivers/roboarm/franka.py | 23 +- positronic/drivers/roboarm/kinova/driver.py | 17 +- positronic/drivers/roboarm/so101/driver.py | 29 +- positronic/drivers/roboarm/yam.py | 40 +- positronic/keys.py | 5 + .../offboard/tests/test_remote_policy.py | 4 +- positronic/offboard/tests/test_server.py | 2 +- positronic/policy/base.py | 118 +++- positronic/policy/harness.py | 265 +++++---- positronic/policy/recording.py | 12 +- positronic/policy/remote.py | 12 +- .../policy/tests/golden_pipeline.json.gz | Bin 7724 -> 7893 bytes .../policy/tests/test_golden_pipeline.py | 42 +- positronic/policy/tests/test_harness.py | 513 +++++++++++------- positronic/policy/tests/test_policy_io.py | 6 +- positronic/policy/tests/test_recording.py | 6 +- positronic/policy/tests/test_wrappers.py | 125 ++++- positronic/policy/wrappers.py | 4 +- positronic/replay_record.py | 11 +- positronic/robot_controller.py | 7 +- positronic/simulator/env_server/adapter.py | 52 +- positronic/simulator/env_server/proxy.py | 10 +- .../simulator/env_server/tests/mujoco_env.py | 5 +- .../env_server/tests/test_remote_env.py | 4 +- positronic/simulator/mujoco/sim.py | 43 +- positronic/tests/test_data_collection.py | 15 +- positronic/utils/registration.py | 2 +- positronic/vendors/dreamzero/server.py | 2 +- positronic/vendors/gr00t/server.py | 2 +- positronic/vendors/lerobot/policy.py | 2 +- positronic/vendors/lerobot_0_3_3/policy.py | 2 +- positronic/vendors/molmoact2/policy.py | 2 +- positronic/vendors/openpi/server.py | 2 +- positronic/wire.py | 12 +- utilities/fake_dataset_generator.py | 2 +- 46 files changed, 909 insertions(+), 1262 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 33274d34d..efb280b64 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -1,23 +1,5 @@ { "files": { - "./pimm/core.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 8, - "endColumn": 17, - "lineCount": 1 - } - }, - { - "code": "reportReturnType", - "range": { - "startColumn": 15, - "endColumn": 23, - "lineCount": 1 - } - } - ], "./pimm/shared_memory.py": [ { "code": "reportIndexIssue", @@ -995,22 +977,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 62, - "endColumn": 77, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 96, - "endColumn": 112, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -1285,22 +1251,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 31, - "endColumn": 58, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 63, - "endColumn": 70, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1317,14 +1267,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 47, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1587,14 +1529,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 110, - "endColumn": 112, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1635,14 +1569,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 110, - "endColumn": 112, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1658,14 +1584,6 @@ "endColumn": 73, "lineCount": 1 } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 106, - "endColumn": 108, - "lineCount": 1 - } } ], "./positronic/dataset/edits.py": [ @@ -2357,22 +2275,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 24, - "endColumn": 28, - "lineCount": 1 - } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 24, - "endColumn": 28, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4125,22 +4027,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 28, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 40, - "endColumn": 44, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4157,14 +4043,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 33, - "endColumn": 52, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4215,22 +4093,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 27, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 39, - "endColumn": 43, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4339,22 +4201,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 38, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 43, - "endColumn": 47, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4371,14 +4217,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 30, - "endColumn": 44, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -4403,14 +4241,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 77, - "endColumn": 83, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4419,14 +4249,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 29, - "endColumn": 36, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -4847,22 +4669,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 27, - "endColumn": 34, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 39, - "endColumn": 43, - "lineCount": 1 - } - }, { "code": "reportReturnType", "range": { @@ -4891,38 +4697,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 35, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 24, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 41, - "endColumn": 45, - "lineCount": 1 - } - }, { "code": "reportGeneralTypeIssues", "range": { @@ -5875,22 +5649,6 @@ } ], "./positronic/policy/harness.py": [ - { - "code": "reportOptionalCall", - "range": { - "startColumn": 18, - "endColumn": 56, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 39, - "endColumn": 55, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -6003,38 +5761,6 @@ } ], "./positronic/policy/tests/test_golden_pipeline.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 35, - "endColumn": 39, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 19, - "endColumn": 26, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 35, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -6643,14 +6369,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 97, - "endColumn": 101, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6659,14 +6377,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 52, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6675,14 +6385,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 53, - "endColumn": 55, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6691,14 +6393,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 53, - "endColumn": 60, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6707,14 +6401,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 95, - "endColumn": 102, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -6722,14 +6408,6 @@ "endColumn": 56, "lineCount": 1 } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 57, - "endColumn": 63, - "lineCount": 1 - } } ], "./positronic/server/dataset_utils.py": [ @@ -7287,22 +6965,6 @@ } ], "./positronic/simulator/env_server/tests/mujoco_env.py": [ - { - "code": "reportArgumentType", - "range": { - "startColumn": 42, - "endColumn": 54, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 43, - "endColumn": 58, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -7441,14 +7103,6 @@ "lineCount": 1 } }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 15, - "endColumn": 33, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -7763,38 +7417,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 23, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 45, - "endColumn": 49, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 24, - "endColumn": 31, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 47, - "endColumn": 51, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -8271,22 +7893,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 41, - "endColumn": 53, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 38, - "endColumn": 53, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -8327,14 +7933,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 33, - "endColumn": 48, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -8651,38 +8249,6 @@ "lineCount": 1 } }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 19, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 19, - "endColumn": 33, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 32, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportIndexIssue", - "range": { - "startColumn": 33, - "endColumn": 47, - "lineCount": 1 - } - }, { "code": "reportAttributeAccessIssue", "range": { @@ -8690,14 +8256,6 @@ "endColumn": 45, "lineCount": 1 } - }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 46, - "endColumn": 53, - "lineCount": 1 - } } ], "./positronic/wire.py": [ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 168521c79..b9b764d36 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -36,7 +36,7 @@ library must supply the tools that make late binding possible — codecs, per-tr projections over raw recordings. **Every decision lives with the party that has the information.** Only a driver knows its motion -capabilities, so drivers plan through waypoints. Only a sensor knows its own cadence, so sensors +capabilities, so drivers own how they reach a commanded setpoint. Only a sensor knows its own cadence, so sensors run at their own rate instead of a rate the loop imposes. Only a policy knows what its model was trained on, so translation to model I/O ships with the policy. @@ -116,21 +116,28 @@ 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. -**Trajectory is the command.** Ownership puts execution with the driver: the policy emits a -trajectory of waypoints with absolute timestamps, and the driver plays it at its own control rate, -planning through the waypoints as well as it knows how. Signals are last-value-wins, so a new -trajectory overwrites the current one — the previous command is merely context for the next. -Continuous-update schemes (RTC, temporal ensembling) therefore need no special mechanism: they are -wrappers that rewrite the command more often. An empty trajectory cancels the channel and the -device holds. +**The wrapper 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 +holds otherwise, and emission time *is* execution time. Continuous-update schemes (RTC, temporal +ensembling) therefore need no special mechanism: they are wrappers that hand back a new trajectory +more often, and the harness keeps playing the old one until they do. **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, demuxes the returned waypoints per command channel, and runs episode -lifecycle — nothing else. Scheduling, blending, history stacking and error recovery live in the -wrapper stack around the policy; a session returning `None` means "keep executing the current +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 wrapper stack around the policy; a session returning `None` means "keep executing the current trajectory". +**The platform charges inference, not the wrapper.** A scheduling wrapper may be someone else's +submission, so what a model call costs cannot be left to the wrapper's own bookkeeping. The harness +runs the call on a worker thread and installs an `InferenceGate` around the wrapper's path to the +model: the gate holds the answer until the cost is paid, so the wrapper resumes — and anchors its +chunk — at the release instant, and cannot see a result early however it is written. + **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every model-facing view — action space, control frame, vendor format — is a codec's projection. diff --git a/pimm/core.py b/pimm/core.py index 34c3df7ea..5a3be162b 100644 --- a/pimm/core.py +++ b/pimm/core.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass -from typing import Generic, TypeVar, cast, final +from typing import Any, Generic, TypeVar, cast, final T = TypeVar('T') U = TypeVar('U') @@ -11,7 +11,7 @@ class NoValueException(Exception): pass -NODEFAULT = object() +NODEFAULT: Any = object() @dataclass diff --git a/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index fb81e2c56..435314e92 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -12,7 +12,7 @@ import pimm import positronic.cfg.policy as policy_cfg -from positronic import telemetry, telemetry_keys, utils, wire +from positronic import keys, telemetry, telemetry_keys, utils, wire from positronic.cfg.eval import placeholder from positronic.cli.eval.submit import submit from positronic.dataset.ds_writer_agent import TimeMode @@ -223,7 +223,7 @@ def run( raise SystemExit(f'--eval={eval!r} is a name, not a config: pass --policy-image to run it on the platform') # The eval config owns the trial sweep (seed, task range); ``inference_latency`` is the CLI's per-run knob # (sim inference-cost simulation). Overlay it onto every trial context, then self-drive the eval. - eval = replace(eval, trials=[{**trial, 'inference_latency': inference_latency} for trial in eval.trials]) + eval = replace(eval, trials=[{**trial, keys.INFERENCE_LATENCY: inference_latency} for trial in eval.trials]) main(policy=policy, evals=[eval], output_dir=output_dir, timing=timing) return None diff --git a/positronic/data_collection.py b/positronic/data_collection.py index fb18915d9..c4cef98cf 100644 --- a/positronic/data_collection.py +++ b/positronic/data_collection.py @@ -160,9 +160,9 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self.sound.emit(abort_wav_path) tracker.turn_off() recording = False - self.robot_commands.emit([(clock.now_ns(), roboarm.command.Reset())]) + self.robot_commands.emit(roboarm.command.Reset()) - self.target_grip.emit([(clock.now_ns(), button_handler.get_value('right_trigger'))]) + self.target_grip.emit(button_handler.get_value('right_trigger')) cp_msg = self.controller_positions.read() if cp_msg.updated: target_robot_pos = tracker.update(cp_msg.data['right']) @@ -175,7 +175,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self.sound.emit(error_wav_path) if not in_error and cp_msg.updated: cmd = roboarm.command.CartesianPosition(target_robot_pos) - self.robot_commands.emit([(clock.now_ns(), cmd)]) + self.robot_commands.emit(cmd) yield pimm.Sleep(0.001) diff --git a/positronic/dataset/ds_writer_agent.py b/positronic/dataset/ds_writer_agent.py index f6ec10428..54b204594 100644 --- a/positronic/dataset/ds_writer_agent.py +++ b/positronic/dataset/ds_writer_agent.py @@ -10,7 +10,7 @@ from .dataset import DatasetWriter from .episode import EpisodeWriter -from .serializers import Serializer, StatefulSerializer, Timestamped, _PureSerializer, expand_suffixed +from .serializers import Serializer, StatefulSerializer, _PureSerializer, expand_suffixed logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) @@ -59,70 +59,7 @@ def ABORT(): return DsWriterCommand(DsWriterCommandType.ABORT_EPISODE) -class TrajectoryOverrideSerializer(StatefulSerializer): - """Flatten policy trajectories into a single monotonic per-point stream. - - A policy emits whole trajectories ``[(abs_ts_ns, value), ...]``. A newer - trajectory overrides the overlapping tail of the previous one - (last-writer-wins on the timeline): given ``1:[1..10]`` then ``2:[5..15]`` - the recorded stream is ``1@1..4`` then ``2@5..15``. A point is committed - only once a newer trajectory starting after it proves it final; the - remainder is drained by :meth:`flush` at episode end. - - HACK: lossy. Drops the notion of a *predicted* trajectory and cannot - represent overlapping schedulers (RTC/temporal ensembling) that replan into - the already-committed past — such points are dropped to keep timestamps - strictly increasing. Faithful full-command recording needs an - object-valued Signal (``Kind.OBJECT``); tracked in TODO(positronic#NNN). - """ - - def __init__(self, inner: Serializer | None): - self._inner = inner - self._buffer: list[tuple[int, Any]] = [] # latest trajectory, (abs_ts_ns, value), ts-sorted - self._last_ts: int | None = None - - def reset(self) -> None: - self._buffer = [] - self._last_ts = None - - def _encode(self, value: Any) -> Any: - return self._inner(value) if self._inner is not None else value - - def _committable(self, points: list[tuple[int, Any]]) -> list[Timestamped]: - # Guard only bites in the overlap-degrade case (RTC replanning into the - # past); under ChunkedSchedule the prefix is always already ahead. - if self._last_ts is not None: - points = [(ts, v) for ts, v in points if ts > self._last_ts] - if points: - self._last_ts = points[-1][0] - return [Timestamped(ts, self._encode(v)) for ts, v in points] - - def __call__(self, message: list[tuple[int, Any]]) -> list[Timestamped]: - if not message: - # Empty trajectory is the cancel signal (the Harness emits it at episode end): - # drop the buffered tail so flush() does not commit canceled waypoints. - self._buffer = [] - return [] - - start = message[0][0] - # Buffer is ts-sorted: everything before the new trajectory's start is - # final; the rest is overridden and dropped by the reassignment below. - cut = next((i for i, (ts, _) in enumerate(self._buffer) if ts >= start), len(self._buffer)) - committed = self._committable(self._buffer[:cut]) - self._buffer = list(message) - return committed - - def flush(self, now_ns: int | None = None) -> list[Timestamped]: - # At episode end, commit only points already due (ts <= now_ns); the - # remaining future-scheduled tail never executed, so drop it. ``now_ns`` - # is None only for callers wanting the legacy "commit everything". - points = self._buffer if now_ns is None else [(ts, v) for ts, v in self._buffer if ts <= now_ns] - out = self._committable(points) - self._buffer = [] - return out - - -def _append(ep_writer: EpisodeWriter, name: str, value: Any, ts_ns: int, extra_ts: dict[str, int] | None = None): +def _append(ep_writer: EpisodeWriter, name: str, value: Any, ts_ns: int, extra_ts: dict[str, int]): for full_name, v in expand_suffixed(name, value): if v is None: continue @@ -212,7 +149,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 opened = False if start: was_open = ep_writer is not None - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) opened = ep_writer is not None and not was_open if ep_writer is not None: @@ -240,23 +177,16 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 value = msg.data if serializer is not None: value = serializer(value) - # Gate on `Timestamped` so plain list-valued samples - # (e.g. list-state vectors) still go through `_append`. - # Empty list matches too — used as the cancel signal. - if isinstance(value, list) and (not value or isinstance(value[0], Timestamped)): - for sample in value: - _append(ep_writer, name, sample.value, sample.ts, None) - else: - _append(ep_writer, name, value, primary_ts, extra_ts) + _append(ep_writer, name, value, primary_ts, extra_ts) if closing: - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) yield pace() finally: cmd_msg = self.command.read() if cmd_msg.updated: - ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter, cmd_msg.ts) + ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) if ep_writer is not None: try: @@ -265,9 +195,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 ep_writer.__exit__(None, None, None) logger.info(f'DsWriterAgent: [ABORT] Episode {ep_counter}') - def _handle_command( # noqa: C901 - self, cmd: DsWriterCommand, ep_writer: EpisodeWriter | None, ep_counter: int, now_ns: int | None = None - ): + def _handle_command(self, cmd: DsWriterCommand, ep_writer: EpisodeWriter | None, ep_counter: int): match cmd.type: case DsWriterCommandType.START_EPISODE: if ep_writer is None: @@ -282,10 +210,6 @@ def _handle_command( # noqa: C901 logger.warning('Episode already started, ignoring start command') case DsWriterCommandType.STOP_EPISODE: if ep_writer is not None: - with self._telemetry_span(): - for name, ser in self._serializers.items(): - for sample in ser.flush(now_ns): - _append(ep_writer, name, sample.value, sample.ts, None) for k, v in cmd.static_data.items(): ep_writer.set_static(k, v) ep_writer.__exit__(None, None, None) diff --git a/positronic/dataset/serializers.py b/positronic/dataset/serializers.py index 90d92fafd..d04578446 100644 --- a/positronic/dataset/serializers.py +++ b/positronic/dataset/serializers.py @@ -1,5 +1,4 @@ from collections.abc import Callable, Iterator -from dataclasses import dataclass from typing import Any import numpy as np @@ -26,20 +25,9 @@ # - use "" (empty string) to keep the base name as-is # - any dict entry with value None is skipped # * None -> the sample is dropped -# * a list[Timestamped] -> a self-timestamped stream (recording only): each item is -# recorded at its own ``ts_ns``. An empty list defers; a StatefulSerializer may emit -# the remainder later from ``flush()``. The per-item ``value`` follows the rules above. Serializer = Callable[[Any], Any | dict[str, Any]] -@dataclass -class Timestamped: - """A sample paired with its own absolute timestamp (ns).""" - - ts: int - value: Any - - class StatefulSerializer: """Base for serializers registered with ``DsWriterAgent``. @@ -51,19 +39,9 @@ class StatefulSerializer: def reset(self) -> None: pass - def __call__(self, value: Any) -> Any | dict[str, Any] | list['Timestamped']: + def __call__(self, value: Any) -> Any | dict[str, Any]: raise NotImplementedError - def flush(self, now_ns: int | None = None) -> list['Timestamped']: - """Drain any buffered samples at episode end (mirror of ``reset``). - - Called once on ``STOP_EPISODE`` before the episode is finalized. ``now_ns`` - is the episode-end time; serializers that buffer future-scheduled samples - use it to drop the un-executed tail. The default keeps stateless - serializers a no-op. - """ - return [] - class _PureSerializer(StatefulSerializer): """Wraps a plain callable so every serializer has a uniform interface.""" diff --git a/positronic/dataset/tests/test_ds_writer_agent.py b/positronic/dataset/tests/test_ds_writer_agent.py index 95a97ba5c..ecaa952fd 100644 --- a/positronic/dataset/tests/test_ds_writer_agent.py +++ b/positronic/dataset/tests/test_ds_writer_agent.py @@ -9,13 +9,7 @@ import pimm from positronic import geom, keys, telemetry, telemetry_keys from positronic.dataset import DatasetWriter, EpisodeWriter -from positronic.dataset.ds_writer_agent import ( - DsWriterAgent, - DsWriterCommand, - DsWriterCommandType, - TimeMode, - TrajectoryOverrideSerializer, -) +from positronic.dataset.ds_writer_agent import DsWriterAgent, DsWriterCommand, DsWriterCommandType, TimeMode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers from positronic.drivers.roboarm import RobotStatus, State @@ -482,7 +476,7 @@ def test_pickles_with_every_constructor_argument_filled(): virtual_time=True, telemetry_span=partial(telemetry.span, telemetry_keys.SPAN_RECORD_IO), ) - agent.add_signal('robot_command', TrajectoryOverrideSerializer(Serializers.robot_command)) + agent.add_signal('robot_command', Serializers.robot_command) agent.add_signal('robot_state', Serializers.robot_state) loaded = pickle.loads(pickle.dumps(agent)) @@ -492,26 +486,8 @@ def test_pickles_with_every_constructor_argument_filled(): pass -def test_trajectory_override_serializer(): - s = TrajectoryOverrideSerializer(None) - s.reset() - - # First trajectory: nothing is final yet (could be overridden). - assert s([(1, 'a'), (2, 'b'), (3, 'c')]) == [] - - # Next trajectory starts at ts=2 -> only ts<2 ('a') is final; 'b','c' overridden. - out = s([(2, 'B'), (3, 'C'), (4, 'D')]) - assert [(t.ts, t.value) for t in out] == [(1, 'a')] - - # Episode end drains the still-live buffer. - assert [(t.ts, t.value) for t in s.flush()] == [(2, 'B'), (3, 'C'), (4, 'D')] - - def test_serializer_plain_list_value(world): - """A serializer returning a plain list (non-`Timestamped`) is appended as one sample. - - The trajectory-stream dispatch must not hijack legitimate list-valued samples. - """ + """A serializer returning a plain list is appended as one sample, values and all.""" ds = FakeDatasetWriter() def to_list(_): @@ -527,52 +503,3 @@ def to_list(_): w = ds.created[-1] assert [(s, v) for (s, v, _, _) in w.appends] == [('v', [1, 2, 3])] - - -def test_trajectory_override_serializer_empty_cancels_buffer(): - """Empty trajectory is the Harness STOP cancel signal: drop the buffered tail.""" - s = TrajectoryOverrideSerializer(None) - s.reset() - - # Buffer a trajectory (nothing committed yet). - assert s([(1, 'a'), (2, 'b'), (3, 'c')]) == [] - # Empty trajectory = cancel: nothing committed AND buffer cleared. - assert s([]) == [] - # Subsequent flush must not emit the canceled waypoints. - assert s.flush() == [] - - -def test_trajectory_override_serializer_flush_cutoff(): - """flush(now_ns) commits only points already due; the future tail is dropped.""" - s = TrajectoryOverrideSerializer(None) - s.reset() - - # Buffer a chunk scheduled at ts 1..4 (nothing committed yet). - assert s([(1, 'a'), (2, 'b'), (3, 'c'), (4, 'd')]) == [] - - # Episode ends at ts=2: only the due points (ts <= 2) are committed; 'c','d' dropped. - assert [(t.ts, t.value) for t in s.flush(now_ns=2)] == [(1, 'a'), (2, 'b')] - - # No cutoff keeps the legacy "commit everything" behavior. - s.reset() - assert s([(1, 'a'), (2, 'b')]) == [] - assert [(t.ts, t.value) for t in s.flush()] == [(1, 'a'), (2, 'b')] - - -def test_stop_commits_due_drops_future_trajectory(world): - """A mid-trajectory STOP commits already-due samples and drops the un-executed tail.""" - ds = FakeDatasetWriter() - agent, cmd_em, emitters = build_agent_with_pipes({'traj': TrajectoryOverrideSerializer(None)}, ds, world) - - future = 10**18 # far beyond the test clock, so it stays an un-executed tail - script = [ - (partial(cmd_em.emit, DsWriterCommand(DsWriterCommandType.START_EPISODE)), 0.001), - (partial(emitters['traj'].emit, [(0, 'due'), (future, 'tail')]), 0.001), - (partial(cmd_em.emit, DsWriterCommand(DsWriterCommandType.STOP_EPISODE)), 0.001), - ] - run_scripted_agent(agent, script, world=world) - - w = ds.created[-1] - # 'due' (ts <= stop time) is committed; the future 'tail' is dropped. - assert [(s, v) for (s, v, _, _) in w.appends] == [('traj', 'due')] - assert w.exited is True diff --git a/positronic/drivers/gripper/dh.py b/positronic/drivers/gripper/dh.py index df62fa463..6a27f5864 100644 --- a/positronic/drivers/gripper/dh.py +++ b/positronic/drivers/gripper/dh.py @@ -3,7 +3,6 @@ import pimm from positronic.drivers import vendor_import -from positronic.drivers.roboarm.command import Trajectory, TrajectoryPlayer with vendor_import('pymodbus', 'Gripper support'): import pymodbus.client as ModbusClient @@ -13,7 +12,7 @@ class DHGripper(pimm.ControlSystem): def __init__(self, port: str): self.port = port self.grip: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.SignalReceiver[Trajectory[float]] = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.force: pimm.SignalReceiver = pimm.ControlSystemReceiver(self, default=100) self.speed: pimm.SignalReceiver = pimm.ControlSystemReceiver(self, default=100) @@ -33,18 +32,14 @@ def _state_r(): while _state_g() != 1 and _state_r() != 1: yield pimm.Sleep(0.1) - player = TrajectoryPlayer() last_grip = 0.0 # TODO: We must translate these to physical units (N and m/s) while not should_stop.value: try: grip_msg = self.target_grip.read() - if grip_msg.updated: - player.set(grip_msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - last_grip = grip + if grip_msg is not None and grip_msg.updated: + last_grip = grip_msg.data width = round((1 - max(0, min(last_grip, 1))) * 1000) client.write_register(0x103, c_uint16(width).value, slave=1) client.write_register(0x101, c_uint16(self.force.value).value, slave=1) @@ -78,7 +73,7 @@ def _state_r(): force.emit(100) for width in np.sin(np.linspace(0, 10 * np.pi, 60)) + 1: - target_grip.emit([(world.clock.now_ns(), width)]) + target_grip.emit(width) time.sleep(0.5) try: print(f'Real grip position: {grip.value}') diff --git a/positronic/drivers/gripper/robotiq.py b/positronic/drivers/gripper/robotiq.py index f7bdfb59b..893f3def5 100644 --- a/positronic/drivers/gripper/robotiq.py +++ b/positronic/drivers/gripper/robotiq.py @@ -4,7 +4,6 @@ import pimm from positronic.drivers import vendor_import -from positronic.drivers.roboarm.command import Trajectory, TrajectoryPlayer with vendor_import('pymodbus', 'Gripper support'): import pymodbus.client as ModbusClient @@ -22,7 +21,7 @@ class Robotiq2F(pimm.ControlSystem): def __init__(self, port: str): self._port = port self.grip = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.ControlSystemReceiver[Trajectory[float]] = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.force = pimm.ControlSystemReceiver(self, default=255) # device scale 0..255 self.speed = pimm.ControlSystemReceiver(self, default=255) # device scale 0..255 @@ -37,15 +36,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p client.write_registers(_REG_CMD, [0x0000, 0x0000, 0x0000], device_id=_SLAVE) client.write_registers(_REG_CMD, [0x0100, 0x0000, 0x0000], device_id=_SLAVE) - player = TrajectoryPlayer() - while not should_stop.value: pos_msg = self.target_grip.read() - if pos_msg.updated: - player.set(pos_msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - pos = int(max(0, min(1, grip)) * 255) + if pos_msg is not None and pos_msg.updated: + pos = int(max(0, min(1, pos_msg.data)) * 255) spd = int(max(0, min(255, self.speed.value))) frc = int(max(0, min(255, self.force.value))) @@ -81,7 +75,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p while True: if time.time() - start > i * 1.0: - tgt.emit([(world.clock.now_ns(), waypoints[i])]) + tgt.emit(waypoints[i]) i += 1 if i >= len(waypoints): break diff --git a/positronic/drivers/roboarm/command.py b/positronic/drivers/roboarm/command.py index d231efce9..db136733e 100644 --- a/positronic/drivers/roboarm/command.py +++ b/positronic/drivers/roboarm/command.py @@ -82,44 +82,11 @@ def apply(self, current: geom.Transform3D) -> geom.Transform3D: _T = TypeVar('_T') -# The wire shape of every command channel: waypoints stamped with absolute clock ns. A single immediate -# command is the one-waypoint trajectory ``[(clock.now_ns(), value)]``; ``[]`` cancels whatever is in flight. +# A schedule the harness plays: waypoints stamped with absolute clock ns, ascending. Command channels +# themselves carry one value -- the command due at the moment it is emitted. Trajectory: TypeAlias = list[tuple[int, _T]] -def _combine(acc: CommandType, cmd: CommandType) -> CommandType: - match (acc, cmd): - case (CartesianDelta(a, frame_a), CartesianDelta(b, frame_b)): - if not np.allclose(frame_a.as_matrix, frame_b.as_matrix): - raise ValueError('Cannot accumulate cartesian deltas expressed in different frames') - return CartesianDelta(_compose_delta(a, b), frame_a) - case (JointDelta(a), JointDelta(b)): - return JointDelta(a + b) - case (CartesianDelta() | JointDelta(), _) | (_, CartesianDelta() | JointDelta()): - raise ValueError(f'Cannot reduce {type(acc).__name__} then {type(cmd).__name__} in one tick') - case _: - return cmd - - -def reduce(due: Trajectory[CommandType]) -> CommandType: - """Collapse the commands due in one control tick into the single command to execute. - - Folds the batch in timestamp order. A run of same-space deltas accumulates (their motion is summed, so a - missed tick is caught up rather than dropped); a run of absolute commands keeps the last. Mixing an absolute - with a delta, or two delta spaces, has no faithful single-command form -- a delta binds to the pose measured - when it is consumed, which an absolute target or a foreign space cannot supply -- and raises. - """ - result = due[0][1] - for _, cmd in due[1:]: - result = _combine(result, cmd) - return result - - -def _reduce_last(due: Trajectory[Any]) -> Any: - """The trailing value wins -- the right collapse for absolute setpoints and gripper targets.""" - return due[-1][1] - - def to_wire(command: CommandType) -> dict[str, Any]: match command: case Reset(): @@ -139,32 +106,31 @@ def to_wire(command: CommandType) -> dict[str, Any]: class TrajectoryPlayer: - """Plays back a timestamped trajectory at the driver's control rate. - - Call ``set()`` when a new trajectory arrives, then ``advance(now)`` each tick to get the single command to - apply: every waypoint whose timestamp has been reached is collapsed by ``reduce`` into one value (the last - one by default; the arm channels pass ``command.reduce`` to accumulate due deltas instead of dropping them). - """ + """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value + to emit.""" - def __init__(self, reduce=_reduce_last): + def __init__(self): self._trajectory: Trajectory[Any] = [] self._index: int = 0 - self._reduce = reduce def set(self, trajectory: Trajectory[Any]): self._trajectory = trajectory self._index = 0 + def next_due(self) -> int | None: + """Timestamp of the earliest waypoint not yet played, or ``None`` once the schedule is exhausted.""" + return self._trajectory[self._index][0] if self._index < len(self._trajectory) else None + def advance(self, current_time: int): - """Collapse every waypoint whose timestamp <= current_time into the single value to apply, or None.""" - due = [] - while self._index < len(self._trajectory): - ts, value = self._trajectory[self._index] - if ts > current_time: - break + """The single value due at ``current_time``, or ``None`` when no waypoint has come due since the last + call. Several waypoints due at once collapse to the last: an absolute setpoint supersedes the ones it + overtook, and only a late round makes it happen. + """ + value = None + while self._index < len(self._trajectory) and self._trajectory[self._index][0] <= current_time: + value = self._trajectory[self._index][1] self._index += 1 - due.append((ts, value)) - return self._reduce(due) if due else None + return value def from_wire(wire: dict[str, Any]) -> CommandType: diff --git a/positronic/drivers/roboarm/franka.py b/positronic/drivers/roboarm/franka.py index 4614cbbc6..af4cd1a42 100644 --- a/positronic/drivers/roboarm/franka.py +++ b/positronic/drivers/roboarm/franka.py @@ -141,9 +141,7 @@ def __init__( self._home_joints_variation = ( home_joints_variation if home_joints_variation is not None else [0.03, 0.05, 0.08, 0.08, 0.10, 0.10, 0.10] ) - self.commands: pimm.SignalReceiver[command.Trajectory[command.CommandType]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) self.state: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) self._load = load @@ -332,7 +330,6 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p yield from self._reset(robot, robot_state, rate_limiter, should_stop) in_error = False - player = command.TrajectoryPlayer(reduce=command.reduce) while not should_stop.value: st = robot.state() @@ -344,23 +341,17 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p logging.warning(f'Robot error: {st.error_message}') cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) if in_error: # The driver always clears a recoverable error itself; making it optional (hold in # ERROR for out-of-band recovery instead) is a config knob to add when an embodiment - # needs it. + # needs it. A command that arrived with the error is skipped, not queued. robot.recover_from_errors() - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale - # waypoint once the error clears. - player.set([]) yield rate_limiter.wait() continue - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): yield from self._reset(robot, robot_state, rate_limiter, should_stop) case command.CartesianPosition(pose): @@ -376,8 +367,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p robot.set_target_joints(positions) case command.JointDelta(velocities=joint_delta): robot.set_target_joints(st.q + joint_delta) - case _: - raise NotImplementedError(f'Unsupported command {cmd}') + case other: + raise NotImplementedError(f'Unsupported command {other}') yield rate_limiter.wait() @@ -419,7 +410,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p if time.monotonic() > start + duration: print(f'Moving to {pos + origin.translation}') target = command.CartesianPosition(geom.Transform3D(pos + origin.translation, origin.rotation)) - commands.emit([(world.clock.now_ns(), target)]) + commands.emit(target) i += 1 else: time.sleep(0.01) diff --git a/positronic/drivers/roboarm/kinova/driver.py b/positronic/drivers/roboarm/kinova/driver.py index baf573687..3f3a11040 100644 --- a/positronic/drivers/roboarm/kinova/driver.py +++ b/positronic/drivers/roboarm/kinova/driver.py @@ -68,9 +68,7 @@ def __init__(self, ip: str, relative_dynamics_factor=0.2, home_joints: list[floa self.relative_dynamics_factor = relative_dynamics_factor self.solver = KinematicsSolver() self.home_joints = home_joints if home_joints is not None else [0.0, -0, 0.5, -1.5, 0.0, -0.5, 1.57079633] - self.commands: pimm.SignalReceiver[command.Trajectory[command.CommandType]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) self.state: pimm.SignalEmitter[KinovaState] = pimm.ControlSystemEmitter(self) def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: @@ -89,15 +87,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p joint_controller.compute_torque(q, dq, tau) current_command = np.zeros(api.actuator_count, dtype=np.float32) - player = command.TrajectoryPlayer(reduce=command.reduce) - while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): joint_controller.set_target_qpos(self.home_joints) case command.CartesianPosition(pose): @@ -110,8 +103,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p case command.JointPosition(positions): qpos = np.array(positions, dtype=np.float32) joint_controller.set_target_qpos(qpos) - case _: - print(f'Unsuported command: {cmd}') + case other: + print(f'Unsuported command: {other}') torque_command = joint_controller.compute_torque(q, dq, tau) np.divide(torque_command, torque_constant, out=current_command) diff --git a/positronic/drivers/roboarm/so101/driver.py b/positronic/drivers/roboarm/so101/driver.py index f9578630f..35b6fdd5f 100644 --- a/positronic/drivers/roboarm/so101/driver.py +++ b/positronic/drivers/roboarm/so101/driver.py @@ -59,12 +59,8 @@ def __init__(self, motor_bus: MotorBus, home_joints: list[float] | None = None): self.kinematic = Kinematics(_SO101_URDF_PATH, _SO101_EE_JOINT) self.joint_limits = self.kinematic.joint_limits self.home_joints = home_joints if home_joints is not None else [0.0, 0.0, 0.0, 0.0, 0.0] - self.commands: pimm.SignalReceiver[roboarm_command.Trajectory[roboarm_command.CommandType]] = ( - pimm.ControlSystemReceiver(self, default=[]) - ) - self.target_grip: pimm.SignalReceiver[roboarm_command.Trajectory[float]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.commands = pimm.ControlSystemReceiver[roboarm_command.CommandType](self) + self.target_grip = pimm.ControlSystemReceiver[float](self) self._last_grip: float = 0.0 self.grip: pimm.SignalEmitter[float] = pimm.ControlSystemEmitter(self) @@ -88,22 +84,13 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p rate_limit = pimm.RateLimiter(hz=1000, clock=clock) state = SO101State() - player = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - grip_player = roboarm_command.TrajectoryPlayer() - while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) grip_msg = self.target_grip.read() - if grip_msg.updated: - grip_player.set(grip_msg.data) - grip = grip_player.advance(clock.now_ns()) - if grip is not None: - self._last_grip = grip - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if grip_msg is not None and grip_msg.updated: + self._last_grip = grip_msg.data + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case roboarm_command.Reset(): raise NotImplementedError('Reset not implemented') case roboarm_command.CartesianPosition(pose): @@ -120,8 +107,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p q_norm = self.rad_to_norm(qpos) q_with_gripper = np.concatenate([q_norm, [self._last_grip]]) self.motor_bus.set_target_position(q_with_gripper) - case _: - raise ValueError(f'Unknown command: {cmd}') + case other: + raise ValueError(f'Unknown command: {other}') q = self.motor_bus.position dq = self.motor_bus.velocity[:-1] diff --git a/positronic/drivers/roboarm/yam.py b/positronic/drivers/roboarm/yam.py index 829514fc9..3db22f89b 100644 --- a/positronic/drivers/roboarm/yam.py +++ b/positronic/drivers/roboarm/yam.py @@ -184,8 +184,8 @@ def __init__( self._sim = sim self._connect = connect - self.commands = pimm.ControlSystemReceiver[command.Trajectory[command.CommandType]](self, default=[]) - self.target_grip = pimm.ControlSystemReceiver[command.Trajectory[float]](self, default=[]) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self, default=None) + self.target_grip = pimm.ControlSystemReceiver[float](self, default=None) self.state = pimm.ControlSystemEmitter[YamState](self) self.grip = pimm.ControlSystemEmitter[float](self) self.robot_meta = pimm.ControlSystemEmitter[dict[str, Any]](self) @@ -199,33 +199,21 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p robot_state = YamState() limiter = pimm.RateLimiter(clock, hz=100) - player = command.TrajectoryPlayer(reduce=command.reduce) - grip_player = command.TrajectoryPlayer() q_target = self._reset(arm, kin, robot_state) grip_target = 0.0 while not should_stop.value: - if (cmd_msg := self.commands.read()) is not None and cmd_msg.updated: - player.set(cmd_msg.data) + cmd_msg = self.commands.read() if (grip_msg := self.target_grip.read()) is not None and grip_msg.updated: - grip_player.set(grip_msg.data) - - grip = grip_player.advance(clock.now_ns()) - if grip is not None: - grip_target = float(grip) + grip_target = float(grip_msg.data) obs = arm.get_observations() q = obs['joint_pos'] - cmd = player.advance(clock.now_ns()) - if cmd is not None: - match cmd: + if cmd_msg is not None and cmd_msg.updated: + match cmd_msg.data: case command.Reset(): - # Drop the in-flight trajectories so the homed arm holds position rather than - # resuming stale waypoints on the next tick. - player.set([]) - grip_player.set([]) q_target = self._reset(arm, kin, robot_state) grip_target = 0.0 obs = arm.get_observations() @@ -239,8 +227,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p case command.CartesianDelta() as delta_cmd: target = delta_cmd.apply(self._base_pose * kin.fk(q)) q_target = self._ik_or_hold(kin, target, q, q_target) - case _: - raise NotImplementedError(f'Unsupported command {cmd}') + case other: + raise NotImplementedError(f'Unsupported command {other}') arm.command_joint_pos(np.append(q_target, 1.0 - grip_target)) @@ -329,8 +317,8 @@ def close(self) -> None: with pimm.World() as world: # `World.pair` cannot express that it returns the counterpart of the port it is given, so the four # payload types are named here. - commands = cast(pimm.SignalEmitter[command.Trajectory[command.CommandType]], world.pair(robot.commands)) - target_grip = cast(pimm.SignalEmitter[command.Trajectory[float]], world.pair(robot.target_grip)) + commands = cast(pimm.SignalEmitter[command.CommandType], world.pair(robot.commands)) + target_grip = cast(pimm.SignalEmitter[float], world.pair(robot.target_grip)) state = cast(pimm.SignalReceiver[YamState], world.pair(robot.state)) grip = cast(pimm.SignalReceiver[float], world.pair(robot.grip)) @@ -356,12 +344,12 @@ def pump(seconds: float): assert home_err < 1e-4, home_err # Grip round-trip: polarity inverted on the way out (command) and on the way back (observation). - target_grip.emit([(world.clock.now_ns(), 0.8)]) + target_grip.emit(0.8) pump(0.5) assert fake.last_command is not None assert abs(fake.last_command[6] - 0.2) < 1e-6, fake.last_command # positronic 0.8 closed -> chain 0.2 assert abs(grip.value - 0.8) < 0.02, grip.value - target_grip.emit([(world.clock.now_ns(), 0.0)]) + target_grip.emit(0.0) pump(0.5) assert abs(fake.last_command[6] - 1.0) < 1e-6, fake.last_command assert abs(grip.value) < 0.02, grip.value @@ -369,7 +357,7 @@ def pump(seconds: float): # Unfold toward the workspace, then drive a Cartesian square through the driver's IK. The square # sits well inside the reach envelope, at the unfolded posture's wrist orientation. reach_q = np.array([0.0, 1.2, 1.2, 0.0, 0.6, 0.0]) - commands.emit([(world.clock.now_ns(), command.JointPosition(reach_q))]) + commands.emit(command.JointPosition(reach_q)) pump(0.5) if fake is not None: assert np.allclose(state.value.q, reach_q, atol=0.02), state.value.q @@ -383,7 +371,7 @@ def pump(seconds: float): assert solution is not None, f'IK failed for {target}' ik_err = np.linalg.norm(kin.fk(solution).translation - target.translation) assert ik_err < 5e-3, ik_err # FK↔IK consistency - commands.emit([(world.clock.now_ns(), command.CartesianPosition(target))]) + commands.emit(command.CartesianPosition(target)) pump(0.7) reached = np.linalg.norm(state.value.ee_pose.translation - target.translation) print(f'Moved to {target.translation}, error {reached * 1000:.2f} mm') diff --git a/positronic/keys.py b/positronic/keys.py index 9ee511da4..09dbba2b7 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -79,6 +79,11 @@ def is_robot_command(name: str) -> bool: COMPRESS_IMAGES = 'compress_images' POSITRONIC_VERSION = 'positronic_version' +# The sim-only knob a trial's RUN context carries: what the platform charges the scheduling wrapper's +# model call. ``False`` (or absent) holds the world still for it, a float charges that fixed delay, ``True`` +# charges the call's own wall duration. +INFERENCE_LATENCY = 'inference_latency' + POLICY_META = 'inference.policy' SERVER_META = f'{POLICY_META}.{SERVER}' diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 6a3b2c417..65af1fbb5 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -57,7 +57,9 @@ class TestPrepareObs: def test_images_pass_through_untouched_by_default(self): session = RemoteSession(_mock_ws_session()) obs = {'cam': _make_image(480, 640), 'state': np.array([1.0])} - assert session._prepare_obs(obs) is obs + prepared = session._prepare_obs(obs) + assert prepared.keys() == obs.keys() + assert all(prepared[key] is value for key, value in obs.items()) # the same arrays, not re-encoded def test_compression_reaches_nested_images(self): session = RemoteSession(_mock_ws_session(), compress_images=True) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index d5f924285..0105b1f61 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -244,7 +244,7 @@ def __call__(self, obs): class _ScriptedPolicy(Policy): """Deterministic base policy: every session returns the same untimestamped chunk.""" - def new_session(self, context=None, now=None) -> Session: + def new_session(self, context=None, now=None, gate=None) -> Session: return _ScriptedSession() diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 3bb4d66b0..4645c3bda 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,7 +1,9 @@ from __future__ import annotations +import time from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping +from enum import Enum from typing import Any Now = Callable[[], float] @@ -26,13 +28,12 @@ class Session(ABC): **Return contract**: ``list[dict] | None``. ``None`` means "no new trajectory, keep executing the current one" (used by scheduling wrappers). - 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. + A list is a new trajectory, replacing whatever is playing. Single-action + returns must be wrapped into a 1-element list by the producer. """ @abstractmethod - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: Mapping[str, Any]) -> list[dict[str, Any]] | None: """Predict actions for the given observation.""" @property @@ -72,6 +73,86 @@ def close(self): self._inner.close() +class LatencyMode(Enum): + """When the platform lets a scheduling wrapper have the model's answer.""" + + # At completion, and the world runs on meanwhile: the cost is whatever the model really took. + LIVE = 'live' + # A fixed delay after the call started, whatever the model really took. A delay of zero holds the world + # still for the whole call. + DECLARED = 'declared' + # The call's own wall duration after it started, charged on the world clock. + MEASURED = 'measured' + + +class InferenceGate: + """The platform's hold on a scheduling wrapper's path to the model. + + Installed around the wrapper's inner session, so no wrapper can reach a result before the mode's cost + has been paid. The wrapper resumes at the release instant and anchors there. + """ + + def __init__(self, now: Now, mode: LatencyMode, delay_sec: float = 0.0): + self._now = now + self._mode = mode + self._delay_sec = delay_sec + self._wall_t0 = 0.0 + self._cancelled = False + # True while a call is inside the model — a wrapper that answered on its own never sets it. + # ``t0`` is the world instant that call started, valid once ``entered``. + self.t0 = 0.0 + self.entered = False + + def wrap(self, inner: Session) -> Session: + return InferenceGate._Session(inner, self) + + def cancel(self) -> None: + """Release a parked call, whose result is on its way to a harness that no longer wants it.""" + self._cancelled = True + + def hold(self) -> float | None: + """Wall seconds the world must not advance for, or ``None`` to hold until the call completes.""" + match self._mode: + case LatencyMode.LIVE: + return 0.0 + case LatencyMode.DECLARED: + return None if self._now() >= self.t0 + self._delay_sec else 0.0 + case LatencyMode.MEASURED: + # The world may run no further ahead of the call's start than wall time has: measured + # charging only means anything with the world at or below real time during the call. + return max(0.0, (self._now() - self.t0) - (time.monotonic() - self._wall_t0)) + + def _release_at(self) -> float: + match self._mode: + case LatencyMode.LIVE: + return self.t0 + case LatencyMode.DECLARED: + return self.t0 + self._delay_sec + case LatencyMode.MEASURED: + return self.t0 + (time.monotonic() - self._wall_t0) + + class _Session(DelegatingSession): + """Charges the inner call, on whatever thread the harness dispatched it to.""" + + def __init__(self, inner: Session, gate: InferenceGate): + super().__init__(inner) + self._gate = gate + + def __call__(self, obs): + gate = self._gate + gate.t0 = gate._now() + gate._wall_t0 = time.monotonic() + gate.entered = True + result = self._inner(obs) + release = gate._release_at() + # The world clock is advanced by the harness's thread, so the park has to poll it; sleeping + # zero hands over the GIL without adding a wake-up granularity to the release instant. + while not gate._cancelled and gate._now() < release: + time.sleep(0) + gate.entered = False + return result + + class Policy(ABC): """Factory for inference sessions. @@ -81,13 +162,18 @@ class Policy(ABC): """ @abstractmethod - def new_session(self, context: dict[str, Any] | None = None, now: Now | None = None) -> Session: + def new_session( + self, context: dict[str, Any] | None = None, now: Now | None = None, gate: InferenceGate | None = None + ) -> Session: """Create a new inference session for an episode. Args: context: Episode context (task description, eval metadata, etc.). now: The runtime clock (current time in seconds), supplied by the harness and passed down to every wrapped session. ``None`` where no runtime clock exists (server-side, warmup). + gate: The platform's hold on the path to the model, supplied by the harness and installed + around the inner session of every ``SchedulingWrapper`` in the stack. ``None`` where no + runtime imposes inference cost (server-side, warmup). """ @property @@ -105,8 +191,8 @@ class DelegatingPolicy(Policy): def __init__(self, inner: Policy): self._inner = inner - def new_session(self, context=None, now=None): - return self._inner.new_session(context, now) + def new_session(self, context=None, now=None, gate=None): + return self._inner.new_session(context, now, gate) @property def meta(self): @@ -171,6 +257,15 @@ def _wrappers(self) -> tuple: return (self,) +class SchedulingWrapper(PolicyWrapper): + """A wrapper that owns the plan: it decides when to call the model and returns the trajectory the + harness plays, rather than one action for the moment. + + Being one is what earns the wrapper an ``InferenceGate`` around its inner session, so the inference + cost is imposed below it instead of trusted to it. + """ + + class _WrapperPolicy(DelegatingPolicy): """Generic policy wrapper produced by ``PolicyWrapper.wrap()``. @@ -181,8 +276,11 @@ def __init__(self, inner: Policy, wrapper: PolicyWrapper): super().__init__(inner) self._wrapper = wrapper - def new_session(self, context=None, now=None): - return self._wrapper.wrap_session(self._inner.new_session(context, now), context, now) + def new_session(self, context=None, now=None, gate=None): + inner = self._inner.new_session(context, now, gate) + if gate is not None and isinstance(self._wrapper, SchedulingWrapper): + inner = gate.wrap(inner) + return self._wrapper.wrap_session(inner, context, now) @property def meta(self): diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index cee52a980..7b3f4b1a4 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,5 +1,7 @@ +import concurrent.futures import time from collections.abc import Generator, Iterable, Iterator +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum from typing import Any @@ -10,10 +12,10 @@ from positronic import keys, telemetry, telemetry_keys from positronic.dataset.ds_writer_agent import DsWriterCommand from positronic.dataset.serializers import expand_suffixed +from positronic.drivers.roboarm.command import TrajectoryPlayer 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.wrappers import ChunkedSchedule +from positronic.policy.base import InferenceGate, LatencyMode, Policy, Session from positronic.utils import flatten_dict, frozen_view # How far from now an action may be scheduled. A chunk spans seconds, so this is loose enough that no real @@ -21,6 +23,10 @@ # chunk (decades behind) or anchored them twice (decades ahead). 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 directive is +# noticed, and with it the granularity every command timestamp is quantized to. +POLL_PERIOD_SEC = 0.01 + def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: """Reject a chunk whose timestamps are not times on the harness clock.""" @@ -137,12 +143,18 @@ def _end_span(self) -> None: class Harness(pimm.ControlSystem): - """Control system that runs the episode lifecycle and forwards trajectories to drivers. + """Control system that runs the episode lifecycle and plays the policy's trajectory to the drivers. Handles directives (RUN/FINISH/ABORT) and dataset recording. Inference intelligence — scheduling, - error recovery, blending, absolute time stamping — lives in the policy/session layer; the harness - calls the session, demuxes the action dicts into per-channel trajectories, and emits. The RUN context - is handed whole to the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). + error recovery, blending, absolute time stamping — lives in the policy/session layer: the wrapper owns + the plan, the harness plays it. Each round emits at most one command per channel, the one due now, so a + command channel carries execution rather than intent and its granularity is the round — + ``POLL_PERIOD_SEC`` on real hardware, one control period in sim. The RUN context is handed whole to the + task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). + + The session call runs on a worker thread, so playing continues while the model does; the + ``InferenceGate`` installed below the scheduling wrapper is what charges that call the trial's + inference latency. A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's @@ -176,10 +188,11 @@ def __init__( self._policy_session: Session | None = None # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. self._running = False - # Sim-only, delivered on the RUN context: ``True`` advances the sim clock by the measured wall cost - # of the inference call, a float by a fixed deterministic amount (the reproducible golden). The chunk - # is anchored before the sleep, so ``_step`` post-shifts it. - self._inference_latency: bool | float = False + # One session call at a time, on a worker so the harness keeps playing while the model runs, and the + # gate that charges it the trial's inference latency. The gate lives for one episode. + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') + self._future: Future[list[dict[str, Any]] | None] | None = None + self._gate: InferenceGate | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None # Whether this episode's first observation has landed. Until it does the deadline stands where the @@ -198,6 +211,9 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] + # TODO: a late round collapses several due waypoints to the last, which drops the motion of every + # delta but the final one. Revisit if pacing turns out not to hold one waypoint due per round. + self._players = {name: TrajectoryPlayer() for name in embodiment.commands} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) @@ -226,43 +242,36 @@ def _build_episode_meta(self, context: dict[str, Any]) -> dict[str, Any]: meta.update(context) return meta - def _emit_now(self, action: dict[str, Any], clock: pimm.Clock) -> None: - now = clock.now_ns() + def _emit(self, action: dict[str, Any]) -> None: for name, value in action.items(): - self.commands[name].emit([(now, value)]) + self.commands[name].emit(value) - def _home(self, clock: pimm.Clock) -> None: - self._emit_now(self._embodiment.home, clock) + def _home(self) -> None: + self._emit(self._embodiment.home) - def _pace(self) -> pimm.Command: + 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 the poll period to hold wall-clock rate.""" - return pimm.Yield() if self._embodiment.simulated else pimm.Sleep(0.01) - - def _bump_schedule_end(self, delta_sec: float) -> None: - """Shift the active ``ChunkedSchedule._Session`` ``_trajectory_end`` by ``delta_sec``. - - ``_step`` post-shifts a chunk's timestamps to pay for ``inference_latency``; the wrapper's - end-of-chunk gate must follow, or it re-infers before the driver has played the shifted trajectory. - - TODO: reaching the wrapper by attribute-name guessing and writing its private state breaks on any rename. + 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.""" + if self._embodiment.simulated: + return pimm.Yield() + due = [ts for player in self._players.values() if (ts := player.next_due()) is not None] + if not due: + return pimm.Sleep(POLL_PERIOD_SEC) + return pimm.Sleep(min(POLL_PERIOD_SEC, max(min(due) - clock.now_ns(), 1) / 1e9)) + + def _cancel_session(self) -> None: + """Drop everything the episode has going: the schedule being played, the call on the worker, and the + session's scheduling state so the next inference is not held back. Devices hold their last commanded + position — nothing is buffered downstream to clear. """ - s = self._policy_session - while s is not None: - if isinstance(s, ChunkedSchedule._Session) and s._trajectory_end is not None: - s._trajectory_end += delta_sec - return - s = getattr(s, '_inner', None) - - def _cancel_trajectories(self) -> None: - """Drop any in-flight chunk from drivers and from the recording's tail. - - Emits ``[]`` on every command channel, so each driver's ``TrajectoryPlayer`` clears its buffer - (devices hold position) and ``TrajectoryOverrideSerializer`` drops its uncommitted tail. Must precede - ``STOP_EPISODE``, which flushes the recording's serializers and would otherwise commit canceled - waypoints. Also cancels the session's scheduling state so the next inference is not held back. - """ - self._emit_commands([]) + for player in self._players.values(): + player.set([]) + if self._gate is not None: + self._gate.cancel() + if self._future is not None: + future, self._future = self._future, None + future.result() # nothing may close or re-enter the session while the worker is still inside it if self._policy_session is not None: self._policy_session.cancel() @@ -271,17 +280,31 @@ def _finalize_recording( ) -> Generator[pimm.Command, None, None]: """Commit the live episode: cancel the in-flight chunk, stop the recorder — stamping the episode's full static meta (plus any terminal payload) — then close its span.""" - self._cancel_trajectories() + self._cancel_session() self.ds_command.emit(DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})})) virtual_now = clock.now() # before the round below, whose sim-clock advance belongs to no rollout # Give the recorder a round to commit the STOP before the next START (they share ``ds_command``, where # last-value-wins would drop one) and before the home command, so homing stays out of the recording. - yield self._pace() + yield self._pace(clock) # After that round, so the recorder's STOP-time record.io span is still in flight and parents to the # episode. Accepted skew: a producer stepping in that shared round charges one span (≤ one control # period) to the closing episode — the cooperative scheduler cannot give the recorder a turn alone. self._telemetry.end(virtual_now) + def _new_gate(self, clock: pimm.Clock) -> InferenceGate: + """The inference cost this trial imposes on the scheduling wrapper's call. + + Hardware pays whatever the model takes. A sim trial pays what its ``inference_latency`` asks for: the + call's own wall duration (``True``), or a fixed delay — the reproducible mode, and by default zero, + which holds the world still for the whole call as sim-only harnesses do. + """ + if not self._embodiment.simulated: + return InferenceGate(clock.now, LatencyMode.LIVE) + latency = self.context.get(keys.INFERENCE_LATENCY, False) + if latency is True: # bool is an int subclass — check identity first + return InferenceGate(clock.now, LatencyMode.MEASURED) + return InferenceGate(clock.now, LatencyMode.DECLARED, float(latency)) + def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. @@ -292,7 +315,6 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: still bounded. """ self.context = context - self._inference_latency = self.context.get('inference_latency', False) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -304,7 +326,8 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._policy_session = self.policy.new_session(self.context, clock.now) + self._gate = self._new_gate(clock) + self._policy_session = self.policy.new_session(self.context, clock.now, self._gate) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -319,16 +342,17 @@ def _end_episode( """ if self._running: if abort: - self._cancel_trajectories() # abort has no finalize to do it — stop drivers before the home + self._cancel_session() # abort has no finalize to do it — stop the episode before the home self.ds_command.emit(DsWriterCommand.ABORT()) - yield self._pace() # the settling round a finalize also takes, before the home command + yield self._pace(clock) # the settling round a finalize also takes, before the home command self._telemetry.abort() else: yield from self._finalize_recording(clock, payload) if self._policy_session: self._policy_session.close() self._policy_session = None - self._home(clock) + self._gate = None + self._home() self._running = False def _handle_directive(self, directive: Directive, clock: pimm.Clock) -> Generator[pimm.Command, None, None]: @@ -378,64 +402,69 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it return inputs - def _emit_commands(self, actions: list[dict[str, Any]]) -> None: - """Republish-all demux: emit every command channel from this action chunk. + def _step(self, clock: pimm.Clock) -> None: + """Keep one session call in flight and install the trajectory it returns. - Each channel gets the ``(ts_ns, value)`` waypoints the chunk carries for it; a channel the chunk - omits gets ``[]``, overwriting its last-value-wins signal so the driver holds. An empty chunk - therefore cancels every channel. + The call goes to the worker so the harness keeps playing while the model runs. The spin after + dispatch keeps a wrapper answering without inference — ``None``, a local decision — in the round it + was asked: it ends when the call is already done, or when the gate reports it reached the model. """ - for name, emitter in self.commands.items(): - # Wrappers do action-timing math in float seconds; every pimm channel client expects ns. This is - # the single explicit seconds->ns seam. - traj = [(int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in actions if name in a] - emitter.emit(traj) - - def _inference_delay(self, wall_start: float) -> float: - """The inference cost to simulate: measured wall time (``True``), a fixed float, or 0 (``False``).""" - if self._inference_latency is True: # bool is an int subclass — check identity first - return time.monotonic() - wall_start - return float(self._inference_latency) - - def _step(self, clock: pimm.Clock) -> Generator[pimm.Sleep, None, None]: - """Build obs, call the session, demux its chunk into per-channel emissions. - - The chunk already carries absolute timestamps, stamped by the outermost scheduling wrapper. + session, gate = self._policy_session, self._gate + assert session is not None and gate is not None # only a live episode steps + if self._future is None: + obs = self._build_obs(clock) + if obs is None: + return + if not self._rollout_started: + # The rollout begins at its first observation, not when the reset returned: a reset only asks + # the producer for a scene, and the turns spent delivering it are neither the trial's budget + # nor its duration. + self._rollout_started = True + self._telemetry.start_rollout(clock.now()) + if self._task is not None: + self._deadline = clock.now() + self._task.timeout + self._future = self._executor.submit(session, frozen_view(obs)) + while not (self._future.done() or gate.entered): + pass + self._collect(self._future, gate, clock) + + def _collect(self, future: Future[list[dict[str, Any]] | None], gate: InferenceGate, clock: pimm.Clock) -> None: + """Take the call's trajectory once the latency mode lets the harness have it. + + A mode that holds the world does it by not returning: the loop thread is what advances the sim clock, + so blocking here freezes it for exactly as long as the call is owed. + """ + if not future.done(): + hold = gate.hold() + if hold is not None and hold <= 0.0: + return + concurrent.futures.wait([future], timeout=hold) + if not future.done(): + return + self._future = None + actions = future.result() # taken on the loop thread, so a failing call still seals the episode + if actions is not None: + self._install(actions, clock) + + def _install(self, actions: 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 wrapper at the instant the gate released it. """ - obs = self._build_obs(clock) - if obs is None: - return - if not self._rollout_started: - # The rollout begins at its first observation, not when the reset returned: a reset only asks the - # producer for a scene, and the turns spent delivering it are neither the trial's budget nor its - # duration. - self._rollout_started = True - self._telemetry.start_rollout(clock.now()) - if self._task is not None: - self._deadline = clock.now() + self._task.timeout - - # Advance the sim clock by the inference cost so rollouts feel the model's latency, but only on - # cycles where inference ran — sleeping on blocked cycles would slow directive handling. The chunk - # was anchored pre-sleep, so it is post-shifted and the wrapper's gate bumped to match. - wall_start = time.monotonic() - actions = self._policy_session(frozen_view(obs)) - if actions is None: - return - delay = self._inference_delay(wall_start) - if delay > 0.0: - yield pimm.Sleep(delay) - actions = [{**a, keys.ACTION_TIMESTAMP: a[keys.ACTION_TIMESTAMP] + delay} for a in actions] - self._bump_schedule_end(delay) - - # The latency sleep (or a slow inference call on a real clock) may have crossed the deadline. Drop - # the chunk rather than emit past the advertised self-termination point; the run loop fires FINISH - # next cycle. - if self._deadline is not None and clock.now() >= self._deadline: - return - - self._telemetry.step() _assert_anchored(actions, clock.now()) - self._emit_commands(actions) + self._telemetry.step() + for name, player in self._players.items(): + # Wrappers do action-timing math in float seconds; the schedule and every pimm channel are in ns. + # This is the single explicit seconds->ns seam. + player.set([(int(a['timestamp'] * 1e9), a[name]) for a in actions 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.""" + now_ns = clock.now_ns() + for name, player in self._players.items(): + value = player.advance(now_ns) + if value is not None: + self.commands[name].emit(value) def _trial_terminal(self, clock: pimm.Clock) -> dict[str, Any] | None: """The terminal static payload if a self-driven trial has ended this round, else ``None``. @@ -456,7 +485,7 @@ def _trial_terminal(self, clock: pimm.Clock) -> dict[str, Any] | None: def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # Home the embodiment before the first episode; each ``_end_episode`` re-homes for the next one, so # every episode begins from the home pose (a real arm gets the inter-episode gap to reach it). - self._home(clock) + self._home() try: yield from self._run(should_stop, clock) @@ -466,6 +495,23 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p # losing their phases and charging the episode's wall to between_episodes. self._telemetry.seal(clock.now()) raise + finally: + self._shutdown() + + def _shutdown(self) -> None: + """Release the worker and the session. A call still in flight is unparked and its result dropped: the + run is over and nothing is left to install it. + + The harness does not own the policy's lifetime: the caller may run several harnesses over one policy + (a multi-eval sweep), so it closes the policy once, after the last run. + """ + if self._gate is not None: + self._gate.cancel() + self._future = None + self._executor.shutdown(wait=True, cancel_futures=True) + if self._policy_session is not None: + self._policy_session.close() + self._policy_session = None def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # noqa: C901 while not should_stop.value: @@ -483,7 +529,7 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ yield from self._handle_directive(directive_msg.data, clock) elif not self._running: if manual_msg.updated and manual_msg.data is not None: - self._emit_now(manual_msg.data, clock) + self._emit(manual_msg.data) elif self._trials is not None: trial = next(self._trials, None) if trial is None: # plan exhausted — let the recorder commit the final episode, then exit @@ -494,14 +540,11 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ yield from self._end_episode(clock, terminal) else: try: - yield from self._step(clock) + self._step(clock) except pimm.NoValueException: pass - yield self._pace() + self._play(clock) + yield self._pace(clock) if self._running: yield from self._finalize_recording(clock) - if self._policy_session: - self._policy_session.close() - # The harness does not own the policy's lifetime: the caller may run several harnesses over - # one policy (a multi-eval sweep), so it closes the policy once, after the last run. diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index ed6755cef..125849b74 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -35,8 +35,8 @@ Every field is *also* logged as ``rr.Scalars`` on a dedicated ``action_time`` timeline (each action stamped at the inference-request time plus its horizon offset), so a ``TimeSeriesView`` reads commanded values with real axes. That anchor is the pre-inference -``obs_time_ns``, so it precedes the harness's true execution time (stamped after -inference by ``ChunkedSchedule``) by the inference latency. Select ``action_time`` to see them. +``obs_time_ns``, so it precedes true execution — which the scheduling wrapper anchors at the instant +the inference gate releases its call — by the inference latency. Select ``action_time`` to see them. Entity paths are ``{tap_name}/{data_key}``. A tap's incoming observation keys and outgoing action keys share that namespace; in the rare case the same key appears on @@ -221,9 +221,8 @@ def _log_action_series(path: str, arr: np.ndarray, horizon: np.ndarray, base_ns: Each action is stamped at ``base_ns + horizon_i``, where ``base_ns`` is the inference-request time; successive chunks lay out along one clock so a ``TimeSeriesView`` - has real axes. This precedes true execution by the inference latency: the harness's - ``ChunkedSchedule`` anchors commands at ``clock.now()`` *after* inference, which a recorder - tap sitting inside it cannot observe. + has real axes. This precedes true execution by the inference latency: the scheduling wrapper + anchors commands at the instant the gate releases its call, which a recorder tap cannot observe. """ arr = np.asarray(arr, dtype=np.float64) if arr.ndim == 1: @@ -294,6 +293,9 @@ class Recorder: once per inference at the outermost tap and reused by inner taps so every tap stamps the inference identically. ``blueprint``, if given, is sent as the recording's layout instead of the auto-built one. + + The state shared across taps (the nesting depth, the timeline values) carries no lock: the harness + keeps one session call in flight, so a whole pipeline's taps run on one thread, one call at a time. """ def __init__( diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 142539a0d..967d2bf8b 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -26,9 +26,9 @@ def __init__(self, ws_session: InferenceSession, compress_images: bool = False): self._session = ws_session self._compress_images = compress_images - def _prepare_obs(self, obs: dict[str, Any]) -> dict[str, Any]: + def _prepare_obs(self, obs: cabc.Mapping[str, Any]) -> dict[str, Any]: if not self._compress_images: - return obs + return dict(obs) return {key: self._prepare_value(key, value) for key, value in obs.items()} def _prepare_value(self, key: str, value: Any) -> Any: @@ -42,7 +42,7 @@ def _prepare_value(self, key: str, value: Any) -> Any: return type(value)(self._prepare_value(key, v) for v in value) return value - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: cabc.Mapping[str, Any]) -> list[dict[str, Any]] | None: """Forwards the observation to the remote server and returns the action trajectory. Single-action server responses are wrapped into a 1-element list to honor @@ -102,7 +102,7 @@ def server_meta(self) -> dict[str, Any]: ws_session.close() return self._server_meta - def new_session(self, context=None, now=None) -> RemoteSession: + def new_session(self, context=None, now=None, gate=None) -> RemoteSession: compress = bool(self.server_meta().get(keys.COMPRESS_IMAGES)) ws_session = self._client.new_session() return RemoteSession(ws_session, compress_images=compress) @@ -166,8 +166,8 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, now=None) -> Session: - return self._policy().new_session(context, now) + def new_session(self, context=None, now=None, gate=None) -> Session: + return self._policy().new_session(context, now, gate) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/tests/golden_pipeline.json.gz b/positronic/policy/tests/golden_pipeline.json.gz index 43b3afddc64e5f1f8431b2b906cee2df6b07f924..34ad77dbe42719585b389f4c00379d741b280421 100644 GIT binary patch literal 7893 zcma)Bc{r497sqI_Jk}NyGL=M#5mCbEt1P3kWXm=)B#jgi*_Vc>NJL6$WNlFm8@6m`jcd&K|Cp*v68^b9>ujE>ezHdPbSOXv;*_!6H>qHa0u-7V1wear)5q z&<>U93_Kn5lJ|5uBW1~9+f^hzB-AMSY-{Np?P}2cbkO2_!$jcRlmI+%`41W;h=vPV z8lYcAE;D+g;r_HtXkP#$swe?Dak2c;*KN(pADbfJl8kM`O|QIl>XrB7V&)^Woz7}X z(zkUsz4Fm{ul%j2#!t?E%xjELciVcbv@X|rtgP<6^%#*a5D?-8ZRMAjxAe}3p$AB$ z>q!K7J9F^B(Z$92ja*Qk)kk!3@VLCc%(~}mxuz&JCn8Kh04kyNNb<(fz5d>FlamA3 zqYCx|H8tVkD-I2oae5X~LD;ob%XjYUKx%ljgqHdcr}Xi;uh>me_CE|iJqtd5UwW;D zc%hR1k&{ZK2Cg5ji@u^(`_T=zLvOv`-Ea=;-h+OyB@;RG|a=62*r&o&WgSvgil>ig z2Du2Ibd+$#f5=4O5CK^;_>N3x-0OfW8oodCCJqspHIM(AnT30;^TYd;|pMhiIVKLVR`=hMxBZ z@c^>0{HlHq5CH6leW*f#K>UC@Y_h5!1QG=P0~4%9ae{;a1PoT~`>}~L*&DP4Fo&V4 zlQ)1w06SQ4^$sv-8*mPmT&*L@f#R})xmEXX06_rz{FpoH$_I>r5I(RPq#tz$8^RAh z44FiIz=jBb;Se^oU}%US_$WjWZ4??J1lH|10lKbYUnM4wt_nS}3498ofo2OkvKedw zIf2#=JF*3A4zWQ8haK4pwt~2ztHO?m6t}p^M1@a+(U1r<>^l24up{I?+U+{KC>R5I ziq5~zE(X2`DM3$OXBP)whP*<-qe|R+%91P*K z-x0w~Cls>T9u-l&6C45AZ7rUf2z)=Rs6 z%vnF}@iAvaTJAAtgR~0H^~-(2zl&Xv_XhoIAb&1rY$ z>ALHUHQ-F(QfSuO_q<0^@}^qJF}EB2!qrmOR?I>+yZ+H|&Y3NPe@EZ12mTmQXyT_l zqc#O*F7*H2xc*fsNr?o!->u(_rC&ISoOp^jEby>3`ee%Y%mLwWgYxu7eqXZNjT$Nb zgo8F8|5>?&pLy&%!7PBd!sU`$*0f(K@q4zmyyak~{VzieH|@dH?>_k;?a*BnOzgRn z-wO2;)vB!%59yzl#!z>ke4N&KS0y0!i(v~VBJqg+X&DUvhW}Z=JX8OPid?ALzlMc1 zs0di5x(3Y+A}fPvGisE=JUhNl&8LUk!>(4iB8x_8tt}-AfsCxKiE*LZ6iP+jC&UuP zFi^Q0K7j8zOv$3jboD){OYL#JQojw!0qi#~U^_f*Kgw4Y5b4k&z^ zfwOCjw`qjK=Qp%n@5%6r{-k@Qi8yqoTSv;mT`CCNm`973hmU7pPN^SBf2DwIM@9|D z+jhn?-aKy_H0^QfRTIB|{JBJF*}G;#>v-!@q&D@9rkYnk6}{s9eJA&z4?_|TkBZG2 z;}?%hQMOmokJ`SVe=C{S9t@*e;c^RBsvzw_dAmo`1KJzq`NV3%xr|vD8jkLZU6SbT z?&&E_X1Gh}2hiKH$6hbFznE#+dx0itAl>EEl&RcC{8DX6T&cr;n&aG{McqiRPMmBI zMx-OmQ>$y1neywU#pg&IZLF8x(eAk8In>ds`SxwXi$ZNW~fkvbjIYmK}(afX@n=X10KDTP_0{k470O_1nMpUZH8 zt0^EwZ#w0mTen8tnZMsSQ%tS?415=tCLB&HF9gLuTX+k?u!d}kI4O`C_i6N z@%${YEJvSOOvj`;2Eq%OJHD7M72QDUNPX9r9pe&swPZ+ra3R#<67%pFe$na^v%fIz z6C;L?rhC>q#5ee`4wc9qKRdgW7hRTHpS5_-_;C&fCLIktmu@~Xq3o&IM?BwGMBO<0 zN|Sc5Zh^SF_kbz0in*U~xSnnXez^VKJ>7{?dVYg)wrXzWH>LxzBpw@1*)sd+(!I9j z;fe|2?)q}N5A*!Q(_2eD>^x`YULtj6ZaNcua3RwR_4N39W(AOz-_&X2JQmc#QFtu9 zfW9rlFVQUC3SZeuPW9O$plsdVp621yU2+OfkE(i($RQ&F?!Hv)_m#gGBo%$kp@!N= zVT?8Pihp0JOXnPja66A5OAn zBA2_uZ~>iW`OJe=8QkOfhX~xls;nYYng5|5axgD-+M9gI+{@8oR`G3LAY4U-_NXQF zV_TT#kf_Q`^FkIG!AQ!G@Y7If&qzAtz3~{*s{88n;1Gu@^TaR3zR2@G&}vC@(Af=d znTY4V{f@X2oRc(CdNAz#JNl$weOqJ2#Gc-^=y7yeEHgHE_@}bkg^q433(BXyE#_AT zU2}uF1ew=b?x$r{AJxNNaQ)`IE3r)hS5cLemUh;?RvBzlR7%%4+Dj17xm?!ioEdeQ zs>VM@YC~sdiHG@3#seOK?WQ`qlJDbZeE$KhnXPxr-7x^`}h${SYyx4-5+qc z9=&Z%IaDHZJd)S!FV@Edi&7S=V1nx{Ybbj~!5Uzgi2CErH&SYAn0&OHs7hhM`ok=I zU9$qWR}_iO!E2bD%zsmth0Jk;EuZ^tTCSn&pK$wqvdel*E}pf+ak55zu46Z1uM3%^ zqvT>*qjSiGT%&Wu#k5w{q*DrE>|$D{bH;^Sr-ODet=GZ03}JnjI(8?%*61^NVi>O# zzhF9gy1?)ct$4cW#OZ3oE%5k7)9KUAhGy{ifRPI0xv1B!ae=4p?{g}7r!I=kP}tb8 zYyc+=QWeaGFMb?pc*SOHRwV#6K* z!22O@QQFuj70xQbLR9~?YF;o5Qj3z@>3PmM1;QZ)T!K|q@7M_14tT-(tCKf^pgLw)A8IfF9q>@F$;pF zLNV$>-Uh@7a-?2w1o^feBZ7QKZ~F~$te)l#a-5#+4e~v`0Pgir|JCV-QZr{J!T&-6 z)~}!PI8hIOD^zX$E7sA6W|ANaS4$R2uR`LK7^}0Tx;mR(e+PT_3%?_?o!G8j{wt8a z2C2`d&z>7`{mA!Q*qa&G;RCLxXK>C(&6z%C0sb#O|2uS-J1L*`U;{JlzdP2U{KtDvv_$GB?$`I&gK63QF7XrQu8Qk5 zf{I6M#;KMb{X`A@7t)x8-S2Q|+;J+{*bJi;+zE3Vq(=(p>Az*1#3!iTTv9frU zEF8qQFIIibYsI~^z9P@nl>o6 zbH+%W2HMis^sPlUhb0SFtnvu8GY1HYw;X52n>TgpzHwO?mNGM8$;`YSku#hL+vWKk z6`@IfO)rY0Rzr}*Bh5x;I_p1YCLK5x>HB^E*qMslO3_w>XN3r+j82o)qdyijy@Zyl z@b1-&^Lm(yget7$E1Ql?{{HP)nXS--Up=oH8ETuMNe!zeAhOa}GLo7`J&dMWwvSb4 zNJNpn=O<^9atjdnIF<~n<$V9LG?m%n5a^Zk?B10kTSU4kOO{`@jJn78xwv{vZ_b!~ zQ)Dwj;{}{B@#ZWEP~S)FAv9AR9~kz$ow5n>>YeZ`MT(qRom=Uco&Y8wKlPPGX?q1d zI9J$5^{gZ$tW0IjX51^ZqlmSX(F*cXAXd>KO~s$z_*3ny(#=?c!?(Sm&(|k94cBhU zTkWc_T(MCzad4scHrRyr)TuG=WyyGJ_H&cfJeLD=FkaCirxMD!bC?9+0%m?D z$K#%pB6 zl^&BmQ^M<%qxmrQ!{D}=W@auq;b( z(pLh?sD7$1+ZbJ(aqoMri>d7*wiN4?Q)^~|i#BzoT|HGaTEvFv@mKkLyYABh6xP=5Hcy-u9TrFCYSqFryUQ=smw!(;RF*UP z@ktyv$nk*t`s_VfwsbTLT4oVNgti;6e6L2_j1x~%eb5CDatu^#D+y7f&#B`oXV!f6 zpU}|ks_3)#v4O_%cwg(N@-)kOnh$oF&H!Z3N;8zb~ z+p~M!jC!3S<%=gB0`3&ds<$Wb>_ z#IPIjqaI4vk3L&;L7=TOps)h*G?XP@op`qUm1{padDugoddX&mNjP;+t=&v}} zU)#$ConeXTuN6Db z1Aa(A?IOz~! z3CVXup{q&QFm7Ivt*T??j8f)cUg*c6Cb;%b67iuQs>tl~uYZQdo84GcVWp8Tt}l-G z(a492yEA9An8u2OWyBF0CRy$DRZ78R5Qn4>?j`9$AIUI>YW_(!n4YF?BBM-tY&{epLraMtgJNde^Lsjf_=`+ z_|G=~(`#|n0;B!RmT-E$F|Sx%zA>NJ=X_%Uu?+>rLfv|f2{OkEjJJq678r|&g%ucY z6H6~h5*V3xxO-l9+S?#Z%WT1PX zCR=H|aPkIPxVAs)n=3g8#0GG{1glVNAP#^B2CM34191YvFrz9II|vMj!rZF**+E=y zH$9O$RwHJFG71UV02YCuPJn?$*>5LHqS%M*d4HdH%Jc1g;kMu z@PhUM#jt)-GA~FGsDQE6>;OPYKs{_nO)>yd2HwN8Ys^Krr9Tns3+BVB0zEMA8a96H zLEsbYeoZhxRvj3FmDI2aU^Rg6u#TEw0jw6Vn4eT1fawoc6RbrEg8l>WD8Xv`1wqFE zVI`wlln_V<5LI%k?H2;+0lSp$)S`qzh_N)|i=u6T2Xg~lgdNX!o4AHl@B2TE?KRBI zr4W;D>?c}0_$$;Lw;+Vy8KW+Ymx}Qd#_tUwu(-nw@_oIe8{|Yi;tld6z1|z-$9jw# zl^DDo;D_LHvBf!9>L?{L=mek`~TNT*BLQg!}^k*Z`{iwWqz|ypNJQ^ zpBVCsxBpq7W<=6^@<^@NFDtnksMp9~hK|B-#&XTzib;Zc5gtBv8m@YyCF?q~H$Qdn zi{^QYe}UvIS4p~pzjF<9+kUg9KOk(8=%9BO`1juc8t;V)jeD%SapE{6D?>smEW!TR4UO{r7=@!u%|w ztdcPh9mH7bZKssK{kYUa8BACzj%FN%(!+yd;R_8*N?Nng^ZHAcxVdO`g^`k{gDJOf e{Y9B$8eEsBqL(??*p`<;Y~SIZBiWTW*!~OdG1933 literal 7724 zcma)>2{@G9`^QN=DI<~eXd78dL?c4UO!`ZqqKU{lLz0YaADBLj-+Hg-zFhZppYxn^?z22T=R45b+}vA^m|j^mZ*7mc zXm!=x!N$P~V{_F?_KLIpRgQt!!8>gl;dULzNiO)19aDVgy z^HIH6!y66F^gxenD}P<7jR~9b3$?{9t;FBV^cIn@^U#j@$yz>yH?!N#+>6nsG(SGe z*HOY8B@BDdbj6V7=tI>xm}iv)#+7=!bp?UnM`K#PfeWi(whp4vTpih@tY6C#WuRxt~7|&o#wISmC=eBbAMv_A>2ODULkhjCi$V$H+fc}Q zn-98kp@7Z~)_=B~Y!5b^{N9xm7`hYgG5P(`rbmKZNqlD_g8UZ)KTCQCAG+Be-{%!V zmn*O2hKemZodajm+$vTC7D~n@JRI3|Yv^*6N(fYb*>QF_8w_sA_$7M8KF5i1p4G&9 zsI*rmJ@=;rL^ZGm@oICF_SH$lE=&K=&G1HdPrFifYOa&bU6z5NJC+?)Ne+G9)fRZu3I}kldGu^>0X?2 zj39-T8YecUSr9U?QVe2y+D(F#x6~A|CoPeX;Vs3aVw33GB0D5!Q!8NWEa)0jO8q2G z4uBKnM!|eJIe|4Gh;s1dTm`HHB~ed&Ik^B{P#Hz{zzPi7q28Ny{v2%jt2sR zRh1p(!tp|sVc*NjxNv-sqcCpr(bYJ9$Vr$axokD<7Gjg{7hXM>A=!5gP6%QGyH4f+ za3JI&>;c&qz-@*Y!Y-4`*5I~4P_Q7f#?lI`lz28KcE3_N6Yhq4MCjtPy9GH0jEJ_G~JI=5I75-M8o_!Hv{^hA=<&8 zQwX>KUPeFh=iCCAHfvE56n2(x^FxaSkU`)g7=cCwkcENEU@|&BfV>sJfX~pK0px9f z16Y9;2_$cS{sFZgKS2lFz)tktz#BUOPmqp&7kC2-_<$T18-s3$u+rIRp&N8#7Z3=F zSlkVYRfsrk*}}IrFZ6U`h_&$*Nm$|VE@I74&p0U5fvq z?*O`cKGNmUZ?kn^=)TqW`WJ>JGOwuIPN>mE|`PeW5OvU*!}U zBH7+4K$f_vy*i3}G3;>b9|O@u&IgBmujBq$aj~iQT7$LJj*4yUj{c>^U;1>t(w;gb z^B>1@Q~Mgy_V9Qr0&@nTc&mXBpuJso&T0*G{*4U$e{V$TwBq8 z(sTxPq`!7(q^&TUJceZ^n)(oiUc2lVFyQpD4|SNCoyJNSjFJ#ju zOSM#CI3nk(!TQ@@Mru?;5qZ^7HPQu6(u`iG6B>CgafPXshniyuU6y)Rzs*fKo4+(= z)E=h{F?~d*!aRpmkTb`Wa|YhddF7e2{t4C&Rwqhd;=A(sL^Wo*FVf^q5WC7vRR?A-cq>4lgbH2CB~`Cq(>)^ z9i<1Qi=HAnFeTUCewuBi-kN9`;}3HoT}Nk{6-F8tB4U^aW=0-cj9(SgiX=Vb87Myt zwew_;)LNq6N)7rC>P4N4mmM8+NOhGbsx#h zE8JgOlh^pf=JxIyuNqNLW5(x6)~ZSWYV)mcLZgO^jBPilIGV@jFo|ytP9Vqfme#Ve zTc2{tHVcalD}wGpo3~v&&uUosE8Hv`V??W38(KfTlVOY+xY{_Vl|@C^^joH^D>@UY zcxUYJlc>NMSNrO;;d5_T3radNH8SBrv;AE7J){*&X=$yV%`{~-)mtAM>T(Q|IT#v) zAiOETKUS=jkNW%o;d0KMHBHuGvxdxyo`V-3I>$;XHC27nc;P8i>Y4EJZ8mG~WVM9q z_hYhDV|EvY#w1mCAOA8Ivaovp*ArJ%5&beXr%%t&oyK@PQI`Ajo?65yr1m{sY%Y&)7(@AJVSS5mLzBe^2){&qiBUOMpn4w{Iq8Nutv z_{p6U*IVxL4kJcV=UI!QH3z*{iA}pRH2!Ab@cuWdZuRB)JtfKm8LX0;@TQXH6udVc zsX>DnzoZ?0e11&sntmQ@@b$``9s9eVda65rj6Xd(N)s&FB4ke6I?dYP2BZ7WwA!_B8EvM-Wx_*NrW7Cm5c?->iSp0BULLAP1JZkY0ZM@*b}9?L6z4VX|e9 zl?3-kBC_{8r45C(He&NS*jpFkRFczjs#i*lHZAI2j~&`i7*&*FZQbqx+lu@cJk7ZX zYF66O`^vF*>&^TgZL`G%b8qE&F{EmGgfg6v3w{|U%Bae55w|=v&x(9 zOy?Oh7$%?K5oy1$T3AG^N8M~YPHt5Axu;5Xwz}!&D)x4pirQn(aLnd^|ENJv-28%f zYy>l|%{0=87&*OGvqUT1RNbci(Yllr=JuwmA9|t}3KYjfTuGGn*DZ0Wckk!O)v0b$ z$E z*dV+>#fQR`Bgz8LkXmyF8mr%OyGwszv-Cuy2Mb6~odo4`T%U5FbSaHeQHe*&3 zzXEf4D=_bXPsDSm?_&cqr83o3_ego=1}L*Q zOL76le+6c`BVfhKco~?xx^>s=#X%d%mAlVsp_7CZgV9M_72l$hb}Fv7ND_S>_Ez?6 zutk!T;#-R(S;h61Ne2|QELEV?vC7@=^?F9`ug7NmAX3we2*KEl31VxS2VuQ;#w77; zS~MZpJ7bzSaE^g5x7YJd`Vk{Q5f377`__>@4{(61Q6j!9o?Q(JqENnMPJkN}N5%V+ zR{;=E0oCbC<^nc8*mF!fC_pK`EPfRi4@4GLRwly5#S2k{b(Y0*aq&UmFb=ZFYA$|= z226w;znW`P=AOAz9KlMuWYiiiA&3#of!w)<3xrs}?vhae*H*}R7>XRfhHDGN4CYG~ z0k*9p%AecxYGX6X0mr!kI0Sx2J;8DE0jl66ijL#l2&jSFXqX=-KX4ob(GGsK3aFD} zKrvYWI0Y)8JN?L;09{ZWE#glW1oS{XG|Hd6888GdqT~I^LVyW)4K2+Z*l?ZOeq+E5 z5U>RAqICmq2m{t&Jo;|Hjjez!_yqkf;Knw<9xOv|47{-&aBe<>D^d7F+185o4OH6+ zcz^@wvOqN`;0?~8IfB$gfSaszjs~gi0s_G87QR7hqOJ|Pig+u;_8Pe2tpjm(#DuaS zp+8b;RfDc7{^JPn2eQ^W9Ypm%X%|EdtilITZR6#kG94vdPytdLaG;@9%(_#nShUhKYx!j!^Iq@6r7`8GA+x zYULgKJ)E;{2A$zK9VBLDfzK4vZMbx%|MA4Du{^hJuQAc>^8}Wk#PA6J%2)N5|TkTr8!zbB_2uQJ!qNPO)bYe+LATf|io9z#sVF)fyHFDkT)E3{|o zrROpeatOCo&}NQxmC;-JbC6OgY?~I7E)nFlX&+9;TVF6EGxITz1f&w!Hf?omKh@5= z1L{)0A%I$<$PbkzhD))nW6yXr)Nn@6V`^YvRATI}mhgP=3Ecy#85;K90ikELQuNqH zGP5bI+wQT=74b#@$(Q1cFz%Y&w?C^vify53;~s7%q>k4w_BHxGJE3zCDwEfj7Nw^d z$uer1*to!UrJJ@{_iNuP4rZj~$eIZDTBE(;|5(CDV{^OV3q~pWY=i1NCS|TG z8kW$oQ~PLnB)KutXL|3_1RvpcK*Z=n6YPDk3(K7j8aI+MOu zPk(X^*by<&HvPb}kdUvv)NYgX>9)|A&xW6&IZhs0IYb?8^U@5{^dA$9 zZL`~YP4*mp88ZajdiHu;-e`GtGmL0+$yxc z6sh9+8W`%tV6hl81jCWUgQV15ux_E37XOv6BceK6<)VRv zT3>5kcJ{$xLhYOn7vh?mJX9BZ2tr)*} z{2Qi&j7k`77!!mGBd@SVM_$GDf^F&7<=xwN5`WbV^OAZ|^KGS!jQQ@ERx?^u?yY#<_O+dziE?_Wd55Yab>@ z556uW%s*^go`9yk*x4MG>M`6gH%4N?e!UxFj$Ui-bDiouK6hRS*tgx#@ChA- zE-&}@=T))y$zgMB;N94?w%zJV*RDb2VIr{p~CPb=g1K>re4Eupk5q0@>>6g$|dEr{Y27Tqs7I|Y4NvigE@PTdfEf_*O_6m0&+ zz9Rf3iZ8-`T|9=B@lWM^dxnS=opN}RMId`I)N&0uj2=bf5Qa+8P{Zg|C?Z=Vw9aDm zUTIvwP+w`FF?z2xY%skzkGW>?=(k7OkqJ7H@W^RB+PPevkMKyQ-talH&Spg9jNaI} z_d5EBNNj)U*-89c`v~uI77td-T~C_XHr~O(f#ZOzhJj_i9Jo~w2u!(*gA=zJA_y}q z^X0?=knOPRWgM%<)Gx1lPX`1+JrwE&c?|%97g6yy$N;b%yoTz$L0$`pfPN^En`CZ4 z9E?DrZjM(F)*T%Iq`*v6*-bSbKo%@Uap2S-zyYuZbrh$@3;d5Nj#qCI&bjr4_aKT)QFBFo*KabzG8=MA+T(OIuZv$z1Ja>v7VMC$OsW zqdYhn$P3u_@-iOWen=UNn{pI_lY>;jBq?POoIK>+lSmB>o6CMmC`vppmlC8M=1URT zz;y`H1B<7`Z{Sja48h7MB79t`kO^2PC7zE9fk|^T-#XdFyIKyc@ZE?z27$;aS8(v- z)FFa$h84d2I1R{lx$6}i0=Sb9ak&Q-z5*&XdNBd{&^xx!7c5Zltwn{%)kb=)pZmH_`k(NFULI zlE^^Oe92#QYMmWXR44L()|);0_oM^}{9bQP3_BdTH0t*!+n%<{iUlv-1O6nKU+Q0z zdP(H>t*;QqpCQxa^AC>zf1Jrhkv=hY`UNrU_J15pr$8(<<&0fFACt=Zh2vQUTlX$b zap(K=@!wSPB?0{^k%#b!{wKO#c@2o{0UN}*elKC`IPb9;6|ap;?z^Yp6TVw_6rvS< zjBFgOu9+ukvsI!tDfw1}!oBcRWP^!T^eN4}y*7u}{7t1VPMbiL5(!4XB9QjCJj diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index e7aceab6f..b2733e930 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -37,7 +37,7 @@ from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import CartesianPosition, Reset, TrajectoryPlayer +from positronic.drivers.roboarm.command import CartesianPosition, CommandType, Reset from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation from positronic.geom import Rotation, Transform3D from positronic.policy.base import Policy, Session @@ -82,7 +82,7 @@ class ScriptedProportionalPolicy(Policy): clock, no images. Codec stamps/truncates; the harness anchors/schedules. """ - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _ScriptedSession() @@ -108,13 +108,11 @@ def ee_pose(self) -> Transform3D: class FakeRobot(pimm.ControlSystem): - """Deterministic closed-loop arm: applies the latest command immediately. + """Deterministic closed-loop arm: applies each command as it arrives. - Mirrors ``MujocoSim``'s arm loop (read latest command, apply, emit - state). ``ee_pose`` becomes the applied ``CartesianPosition`` target and the - first three joints track it, so recorded state is a lossless re-expression - of applied commands. Closed loop: the policy's next chunk evolves with this - feedback. + Mirrors ``MujocoSim``'s arm loop (execute on updated, emit state). ``ee_pose`` becomes the applied + ``CartesianPosition`` target and the first three joints track it, so recorded state is a lossless + re-expression of applied commands. Closed loop: the policy's next chunk evolves with this feedback. """ def __init__(self): @@ -122,7 +120,7 @@ def __init__(self): self._q = INITIAL_Q.copy() self._status = RobotStatus.AVAILABLE self._error_pending = False - self.commands = pimm.ControlSystemReceiver(self, default=[]) + self.commands = pimm.ControlSystemReceiver[CommandType](self) self.state = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) @@ -142,19 +140,12 @@ def _apply(self, cmd): def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): self.robot_meta.emit({}) - player = TrajectoryPlayer() while not should_stop.value: cmd_msg = self.commands.read() - if cmd_msg.updated: - player.set(cmd_msg.data) if self._status == RobotStatus.ERROR: - self._status = RobotStatus.AVAILABLE - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale waypoint. - player.set([]) - else: - cmd = player.advance(clock.now_ns()) - if cmd is not None: - self._apply(cmd) + self._status = RobotStatus.AVAILABLE # the command that arrived with the error is skipped + elif cmd_msg is not None and cmd_msg.updated: + self._apply(cmd_msg.data) if self._error_pending: self._status = RobotStatus.ERROR self._error_pending = False @@ -167,18 +158,14 @@ class FakeGripper(pimm.ControlSystem): def __init__(self): self._grip = 0.0 - self.target_grip = pimm.ControlSystemReceiver(self, default=[]) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.grip = pimm.ControlSystemEmitter(self) def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): - player = TrajectoryPlayer() while not should_stop.value: msg = self.target_grip.read() - if msg.updated: - player.set(msg.data) - grip = player.advance(clock.now_ns()) - if grip is not None: - self._grip = float(grip) + if msg is not None and msg.updated: + self._grip = float(msg.data) self.grip.emit(self._grip) yield pimm.Sleep(CONTROL_PERIOD_S) @@ -202,6 +189,9 @@ def _run_pipeline(tmp_path: Path) -> dict: }, static_meta=dict(ROBOT_STATIC_META), meta_source=robot.robot_meta, + # ``inference_latency`` is a sim-only knob, and the fake robot's control-period sleep is this + # world's sole time-master — the shape a sim eval runs in. + simulated=True, ) harness = Harness(ChunkedSchedule().wrap(policy), 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 68ddd9008..de0bebfdb 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,3 +1,4 @@ +import time from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -15,20 +16,16 @@ from positronic.drivers.roboarm.command import ( CartesianDelta, CartesianPosition, - JointDelta, - JointPosition, Reset, TrajectoryPlayer, - _compose_delta, from_wire, - reduce, to_wire, ) from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model 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 Policy, Session +from positronic.policy.base import DelegatingSession, Policy, SchedulingWrapper, Session from positronic.policy.codec import ActionTimestamp from positronic.policy.harness import Directive, DirectiveType, Harness, _assert_anchored from positronic.policy.remote import RemoteSession @@ -52,12 +49,13 @@ def _eval_pass(run_id: str): CAM = 'image.cam' -def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None) -> Embodiment: +def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simulated=False) -> Embodiment: """Minimal Franka-shaped embodiment for harness unit tests. The sources/dests are no-ops: these tests pair the harness ports directly (never via ``wire_embodiment``), so only the spec — names, serializers, - home values, descriptor — is read by the Harness. + home values, descriptor — is read by the Harness. ``simulated`` is what makes + ``inference_latency`` bite, since the knob is sim-only. """ observations = { 'robot_state': Observation(pimm.NoOpEmitter(), Serializers.robot_state), @@ -69,7 +67,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None) -> E keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command), 'target_grip': Command(pimm.NoOpReceiver(), 0.0, None), } - return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter()) + return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter(), simulated=simulated) class _SpySession(Session): @@ -92,7 +90,7 @@ 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, now=None): + def new_session(self, context=None, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _SpySession(self) @@ -137,7 +135,7 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -167,7 +165,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _ChunkSession(self) @@ -202,7 +200,7 @@ def __init__(self, command: roboarm.command.CommandType | None = None, target_gr self.command = command self.target_grip = float(target_grip) - def new_session(self, context=None, now=None) -> RemoteSession: + def new_session(self, context=None, now=None, gate=None) -> RemoteSession: action = [{'robot_command': self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] return RemoteSession(_FakeInferenceSession(action)) @@ -235,6 +233,17 @@ def emit_ready_payload(frame_emitter, robot_emitter, grip_emitter, robot_state): grip_emitter.emit(0.25) +class _Pacer(pimm.ControlSystem): + """Stands in for the simulator: the sole time-master, sleeping one control period every turn.""" + + def __init__(self, period: float = 0.005): + self._period = period + + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): + while not should_stop.value: + yield pimm.Sleep(self._period) + + def _pair_all(world, harness): """Pair all harness signals and return a dict of test handles.""" ds_recorder = RecordingEmitter() @@ -260,39 +269,27 @@ def _ds_types(p) -> list[DsWriterCommandType]: def _last_command(p): - """Extract the last robot command from the trajectory signal.""" + """The latest robot command the harness put on the channel.""" msg = p['command_rx'].read() - if msg is None or msg.data is None: - return None - traj = msg.data # list[tuple[float, CommandType]] - return traj[-1][1] if traj else None + assert msg is not None, 'no robot command was emitted' + return msg.data def _last_grip(p): - """Extract the last grip value from the grip trajectory signal.""" + """The latest grip target the harness put on the channel.""" msg = p['grip_rx'].read() - if msg is None or msg.data is None: - return None - traj = msg.data # list[tuple[float, float]] - return traj[-1][1] if traj else None - - -def _all_grips(p): - """Extract all grip values from the grip trajectory signal.""" - msg = p['grip_rx'].read() - if msg is None or msg.data is None: - return [] - return [g for _, g in msg.data] + assert msg is not None, 'no grip target was emitted' + return msg.data def _emitted_commands(recorder): - """All robot commands across a recorder's non-empty emitted trajectories.""" - return [cmd for _ts, traj in recorder.emitted if traj for _t, cmd in traj] + """Every robot command a recorder saw, in emission order.""" + return [cmd for _ts, cmd in recorder.emitted] def _emitted_grips(recorder): - """All grip targets across a recorder's non-empty emitted trajectories.""" - return [g for _ts, traj in recorder.emitted if traj for _t, g in traj] + """Every grip target a recorder saw, in emission order.""" + return [grip for _ts, grip in recorder.emitted] @pytest.mark.timeout(3.0) @@ -344,7 +341,6 @@ def test_harness_emits_cartesian_move(world): 'descriptor', } - # Last non-empty command (a trailing ``[]`` cancel is emitted on shutdown). cmds = _emitted_commands(cmd_recorder) assert cmds, 'no robot command emitted' cmd = cmds[-1] @@ -548,7 +544,7 @@ 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, now=None): + def new_session(self, context=None, now=None, gate=None): return _StaticMetaSession(self._command) # Session.meta defaults to {} @property @@ -856,18 +852,15 @@ def run(self, should_stop, clock): @pytest.mark.timeout(3.0) def test_done_after_deadline_is_a_timeout(world): - """The deadline is hard: a ``done`` delivered past it (here during the latency sleep) records as a - timeout — ``eval.terminated`` False, payload dropped — not a late stop-signal success.""" + """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(), task=Task(instruction='t', timeout=0.05), trials=[{'inference_latency': 0.2}] - ) + harness = Harness(policy, make_embodiment(), task=Task(instruction='t', timeout=0.05), trials=[{}]) p = _pair_all(world, harness) done_em = world.pair(harness.done) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) - # Obs starts inference + the 0.2s latency sleep; the 0.05s deadline lapses during it, and done is - # delivered at ~0.1s — past the deadline but before the harness next polls. The timeout must win. + # The 0.05s deadline lapses first; done lands at ~0.1s, after the trial has already timed out. driver = ManualDriver([ (partial(emit_ready_payload, p['frame_em'], p['robot_em'], p['grip_em'], robot_state), 0.1), (partial(done_em.emit, {keys.EVAL_SUCCESS: True}), 0.3), @@ -929,12 +922,15 @@ def test_trial_plan_self_drives(world): @pytest.mark.timeout(3.0) -def test_timeout_crossed_during_latency_sleep_drops_chunk(world): - """A chunk whose latency sleep crosses the deadline is dropped, never emitted.""" +def test_timeout_during_inference_drops_the_chunk(world): + """A trial whose deadline lapses while the model is still owed its latency ends with the call in flight: + the trajectory it eventually returns is discarded, never emitted past the advertised termination point.""" policy = StubPolicy() - # The 0.2s latency sleep crosses the 0.05s deadline before the chunk is emitted. harness = Harness( - policy, make_embodiment(), task=Task(instruction='test', timeout=0.05), trials=[{'inference_latency': 0.2}] + ChunkedSchedule().wrap(policy), + make_embodiment(simulated=True), + task=Task(instruction='test', timeout=0.05), + trials=[{keys.INFERENCE_LATENCY: 0.2}], # the gate holds the answer well past the deadline ) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() @@ -951,14 +947,13 @@ def test_timeout_crossed_during_latency_sleep_drops_chunk(world): driver = ManualDriver([(partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), (None, 0.3)]) - scheduler = world.start([harness, driver]) - drive_scheduler(scheduler, steps=200) + scheduler = world.start([harness, driver, _Pacer()]) + drive_scheduler(scheduler, steps=2000) stops = [data for _, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] assert len(stops) == 1 assert stops[0].static_data[keys.EVAL_TERMINATED] is False - # The post-deadline chunk must not reach the drivers: the only non-empty emissions are the homing - # Reset / home grip from the startup home and the timeout FINISH. + # The only commands are the homing Reset / home grip from the startup home and the timeout FINISH. assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) assert _emitted_grips(grip_recorder) == [0.0, 0.0] @@ -1046,28 +1041,25 @@ def reset(_context): assert policy.last_reset_context[keys.TASK] == 'resolved-on-reset' -@pytest.mark.timeout(3.0) -def test_finish_cancels_buffered_trajectory_before_stop_episode(world): - """FINISH must cancel the recording's trajectory tail *before* `STOP_EPISODE`. +class _LabeledRecorder(pimm.SignalEmitter): + """Records emissions from several channels into one shared list, so their order is comparable.""" - `STOP_EPISODE` calls `flush()` on `TrajectoryOverrideSerializer`, which - commits whatever is still buffered. The harness must emit `[]` on - `robot_command`/`target_grip` first, so the serializer drops its tail and - canceled waypoints are not recorded. - """ + def __init__(self, label, events): + self._label = label + self._events = events - class _LabeledRecorder(pimm.SignalEmitter): - def __init__(self, label, events): - self._label = label - self._events = events + def emit(self, data, ts: int = -1): + self._events.append((self._label, data)) - def emit(self, data, ts: int = -1): - self._events.append((self._label, data)) - events: list[tuple[str, object]] = [] +@pytest.mark.timeout(3.0) +def test_finish_stops_playing_the_live_chunk(world): + """FINISH drops the schedule the harness is playing: the chunk's remaining waypoints never reach the + devices, and the only command after the recorder's STOP is the home the close emits.""" policy = ChunkPolicy() wrapped = ActionTimestamp(fps=5.0).wrap(policy) # 1.8 s chunk — won't drain before FINISH harness = Harness(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)) harness.ds_command._bind(_LabeledRecorder('ds_command', events)) @@ -1083,40 +1075,28 @@ def emit(self, data, ts: int = -1): (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), (None, 0.1), (partial(directive_em.emit, Directive.FINISH()), 0.0), - (None, 0.1), + (None, 0.5), ] scheduler = world.start([harness, ManualDriver(script)]) - drive_scheduler(scheduler, steps=200) + drive_scheduler(scheduler, steps=400) - cancels = [i for i, (lbl, data) in enumerate(events) if lbl == keys.ROBOT_COMMAND and data == []] - stops = [ - i - for i, (lbl, data) in enumerate(events) - if lbl == 'ds_command' and getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE - ] - assert cancels, 'FINISH did not emit a cancel on robot_command' + stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'FINISH did not emit STOP_EPISODE' - assert cancels[0] < stops[0], ( - f'cancel ({cancels[0]}) must precede STOP_EPISODE ({stops[0]}); otherwise flush() commits canceled waypoints' - ) + grips_after = [data for lbl, data in events[stops[0] :] if lbl == 'target_grip'] + assert grips_after == [0.0], f'the cancelled chunk kept playing past FINISH: {grips_after}' @pytest.mark.timeout(3.0) -def test_empty_chunk_cancels_both_robot_and_grip(world): - """A session returning ``[]`` must cancel *both* driver buffers. - - Empty action chunk is the session-level cancel signal (per the - ``Session.__call__`` contract). If only ``robot_command`` gets ``[]`` while - ``target_grip`` is skipped, the gripper ``TrajectoryPlayer`` keeps draining - stale waypoints — a partial cancel that's worse than no cancel. - """ +def test_empty_trajectory_leaves_every_channel_holding(world): + """A trajectory with no waypoints schedules nothing on any channel, so every device holds where the + startup home left it rather than one channel draining on while another stops.""" class _EmptyChunkSession(Session): def __call__(self, obs): return [] class EmptyChunkPolicy(Policy): - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _EmptyChunkSession() harness = Harness(EmptyChunkPolicy(), make_embodiment()) @@ -1140,10 +1120,8 @@ def new_session(self, context=None, now=None): scheduler = world.start([harness, ManualDriver(script)]) drive_scheduler(scheduler, steps=200) - cmd_emits = [data for _ts, data in cmd_recorder.emitted] - grip_emits = [data for _ts, data in grip_recorder.emitted] - assert [] in cmd_emits, 'empty chunk did not cancel robot_command buffer' - assert [] in grip_emits, 'empty chunk did not cancel target_grip buffer' + assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) # only the startup home + assert _emitted_grips(grip_recorder) == [0.0] @pytest.mark.timeout(3.0) @@ -1162,8 +1140,7 @@ def test_harness_clears_trajectory_on_home(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=5) - grips = _all_grips(p) - assert grips[0] >= 100.0, f'Expected chunk 1, got {grips}' + assert _last_grip(p) >= 100.0, 'Expected chunk 1' p['directive_em'].emit(Directive.ABORT()) drive_scheduler(scheduler, steps=2) @@ -1174,8 +1151,7 @@ def test_harness_clears_trajectory_on_home(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=4) - grips = _all_grips(p) - assert grips[0] >= 200.0, f'Expected chunk 2 (>= 200.0), got {grips}. Trajectory clearing failed!' + assert _last_grip(p) >= 200.0, 'Expected chunk 2; trajectory clearing failed' @pytest.mark.timeout(3.0) @@ -1194,8 +1170,7 @@ def test_harness_clears_trajectory_on_run(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=5) - grips = _all_grips(p) - assert grips[0] >= 100.0 + assert _last_grip(p) >= 100.0 p['directive_em'].emit(Directive.RUN(task='test-restart')) drive_scheduler(scheduler, steps=1) @@ -1203,8 +1178,7 @@ def test_harness_clears_trajectory_on_run(world): emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], robot_state) drive_scheduler(scheduler, steps=4) - grips = _all_grips(p) - assert grips[0] >= 200.0, f'Expected chunk 2 (>= 200.0), got {grips}. Trajectory clearing on RUN failed!' + assert _last_grip(p) >= 200.0, 'Expected chunk 2; trajectory clearing on RUN failed' @pytest.mark.timeout(3.0) @@ -1224,8 +1198,7 @@ def test_harness_skips_inference_on_error(world): drive_scheduler(scheduler, steps=1) emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], state_ok) drive_scheduler(scheduler, steps=3) - grips = _all_grips(p) - assert grips[0] >= 100.0 + assert _last_grip(p) >= 100.0 obs_before = len(policy.observations) p['robot_em'].emit(state_err) @@ -1233,9 +1206,8 @@ def test_harness_skips_inference_on_error(world): assert len(policy.observations) == obs_before # the errored state is never fed to the policy emit_ready_payload(p['frame_em'], p['robot_em'], p['grip_em'], state_ok) - drive_scheduler(scheduler, steps=3) - grips = _all_grips(p) - assert grips[0] >= 200.0 + drive_scheduler(scheduler, steps=20) # long enough for the first chunk to play out and the next to land + assert _last_grip(p) >= 200.0 def test_directive_preserves_payload(): @@ -1273,6 +1245,18 @@ def test_cartesian_delta_without_a_frame_is_rejected(): from_wire(wire) +def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): + player = TrajectoryPlayer() + player.set([(10, 'a'), (20, 'b'), (30, 'c')]) + assert player.next_due() == 10 + assert player.advance(5) is None + assert player.advance(25) == 'b' # a late round overtakes 'a'; the trailing setpoint is the live one + assert player.next_due() == 30 + assert player.advance(30) == 'c' + assert player.next_due() is None + assert player.advance(40) is None + + def test_cartesian_delta_applies_in_world_frame(): current = Transform3D(np.array([0.5, 0.1, 0.3]), Rotation.from_rotvec(np.array([0.2, 0.1, 0.4]))) delta = Transform3D(np.array([0.02, -0.01, 0.05]), Rotation.from_rotvec(np.array([0.1, 0.0, 0.0]))) @@ -1284,64 +1268,6 @@ def test_cartesian_delta_applies_in_world_frame(): assert not np.allclose(target.translation, (current * delta).translation) # guards against body-frame compose -def test_reduce_accumulates_due_cartesian_deltas(): - # Rotations about different axes so the world-frame compose is non-commutative -- this pins the fold order - # (apply d0 then d1), not just that a fold happened. - d0 = Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.array([0.3, 0.0, 0.0]))) - d1 = Transform3D(np.array([0.02, 0.01, 0.0]), Rotation.from_rotvec(np.array([0.0, 0.0, 0.2]))) - out = reduce([(10, CartesianDelta(d0)), (20, CartesianDelta(d1))]) - assert isinstance(out, CartesianDelta) - expected = _compose_delta(d0, d1) # two due deltas catch up as their world-frame compose, not last-wins - np.testing.assert_allclose(out.delta.translation, expected.translation) - np.testing.assert_allclose(out.delta.rotation.as_quat, expected.rotation.as_quat, atol=1e-12) - assert not np.allclose(out.delta.rotation.as_quat, _compose_delta(d1, d0).rotation.as_quat) - - -def test_reduce_sums_due_joint_deltas(): - out = reduce([(10, JointDelta(np.array([0.1, -0.2, 0.3]))), (20, JointDelta(np.array([0.0, 0.2, -0.1])))]) - assert isinstance(out, JointDelta) - np.testing.assert_allclose(out.velocities, [0.1, 0.0, 0.2]) - - -def test_reduce_absolute_run_keeps_last(): - p0 = CartesianPosition(Transform3D(np.array([0.1, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - p1 = JointPosition(np.array([0.2, 0.0, 0.0])) - assert reduce([(10, p0), (20, p1)]) is p1 - - -def test_reduce_raises_on_absolute_delta_mix(): - cart_pos = CartesianPosition(Transform3D(np.zeros(3), Rotation.from_rotvec(np.zeros(3)))) - cart_delta = CartesianDelta(Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - joint_pos = JointPosition(np.zeros(3)) - joint_delta = JointDelta(np.array([0.1, 0.0, 0.0])) - with pytest.raises(ValueError): - reduce([(10, cart_pos), (20, cart_delta)]) - with pytest.raises(ValueError): - reduce([(10, cart_delta), (20, cart_pos)]) - with pytest.raises(ValueError): # JointPosition then JointDelta: the delta has no faithful anchor to fold onto - reduce([(10, joint_pos), (20, joint_delta)]) - - -def test_reduce_raises_on_mixed_delta_spaces(): - cart_delta = CartesianDelta(Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3)))) - joint_delta = JointDelta(np.array([0.1, 0.0, 0.0])) - with pytest.raises(ValueError): - reduce([(10, cart_delta), (20, joint_delta)]) - with pytest.raises(ValueError): - reduce([(10, joint_delta), (20, cart_delta)]) - - -def test_trajectory_player_accumulates_missed_deltas(): - d0 = Transform3D(np.array([0.01, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3))) - d1 = Transform3D(np.array([0.02, 0.0, 0.0]), Rotation.from_rotvec(np.zeros(3))) - player = TrajectoryPlayer(reduce=reduce) - player.set([(10, CartesianDelta(d0)), (20, CartesianDelta(d1))]) - out = player.advance(20) # both waypoints due in one tick -> summed, not dropped to the last - assert isinstance(out, CartesianDelta) - np.testing.assert_allclose(out.delta.translation, [0.03, 0.0, 0.0]) - assert player.advance(30) is None - - @pytest.mark.parametrize('status', [RobotStatus.RESETTING, RobotStatus.ERROR]) def test_robot_state_serializer_drops_not_ready(status): state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=status) @@ -1354,27 +1280,15 @@ def test_robot_state_serializer_available_has_no_error_key(): @pytest.mark.timeout(3.0) -def test_shutdown_cancels_trajectory_before_stop(world): - """Shutdown while recording must cancel buffered trajectories before STOP_EPISODE. - - ``STOP_EPISODE`` flushes ``TrajectoryOverrideSerializer``; without a prior - cancel it would commit the unexecuted tail of an in-flight chunk (the - FINISH/RUN paths already cancel first). - """ +def test_shutdown_stops_playing_the_live_chunk(world): + """Shutdown while recording drops the schedule too: the unplayed tail of the live chunk never reaches + the devices after the recorder's STOP.""" events: list[tuple[str, object]] = [] - - class _LabeledRecorder(pimm.SignalEmitter): - def __init__(self, label): - self._label = label - - def emit(self, data, ts: int = -1): - events.append((self._label, data)) - wrapped = ActionTimestamp(fps=5.0).wrap(ChunkPolicy()) # 1.8 s chunk — won't drain before shutdown harness = Harness(wrapped, make_embodiment()) - harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND)) - harness.commands['target_grip']._bind(_LabeledRecorder('target_grip')) - harness.ds_command._bind(_LabeledRecorder('ds_command')) + harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) + harness.commands['target_grip']._bind(_LabeledRecorder('target_grip', events)) + harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) robot_em = world.pair(harness.observations['robot_state']) @@ -1382,7 +1296,7 @@ def emit(self, data, ts: int = -1): directive_em = world.pair(harness.directive) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) - # RUN + a complete obs buffers a chunk; the driver then ends, which makes the + # RUN + a complete obs schedules a chunk; the driver then ends, which makes the # world signal shutdown while still recording — exercising the run() finalizer. driver = ManualDriver([ (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), @@ -1392,15 +1306,9 @@ def emit(self, data, ts: int = -1): scheduler = world.start([harness, driver]) drive_scheduler(scheduler, steps=200) - cancels = [i for i, (lbl, data) in enumerate(events) if lbl == keys.ROBOT_COMMAND and data == []] - stops = [ - i - for i, (lbl, data) in enumerate(events) - if lbl == 'ds_command' and getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE - ] - assert cancels, 'shutdown did not cancel robot_command' + stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'shutdown did not emit STOP_EPISODE' - assert cancels[0] < stops[0], 'cancel must precede STOP_EPISODE on shutdown' + assert not [lbl for lbl, _ in events[stops[0] :] if lbl == 'target_grip'] @pytest.mark.timeout(5.0) @@ -1650,3 +1558,230 @@ def test_doubly_anchored_chunk_is_refused(): def test_anchored_chunk_passes(): """A real chunk spans seconds around now, and a late action sits just behind it.""" _assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) + + +class _SlowSession(Session): + """A session whose inference costs ``wall_sec`` of real time and returns a fixed-length chunk.""" + + def __init__(self, wall_sec: float, span_sec: float, steps: int): + self._wall_sec = wall_sec + self._span_sec = span_sec + self._steps = steps + + def __call__(self, obs): + time.sleep(self._wall_sec) + dt = self._span_sec / self._steps + pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) + return [ + {keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': float(i), 'timestamp': i * dt} + for i in range(self._steps) + ] + + +class SlowPolicy(Policy): + def __init__(self, wall_sec: float = 0.0, span_sec: float = 0.2, steps: int = 10): + self._wall_sec = wall_sec + self._span_sec = span_sec + self._steps = steps + + def new_session(self, context=None, now=None, gate=None): + return _SlowSession(self._wall_sec, self._span_sec, self._steps) + + +class _ReplanEarly(SchedulingWrapper): + """Infers on the first observation and again halfway through the chunk it returned. + + 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. + """ + + class _Session(DelegatingSession): + def __init__(self, inner: Session, now): + super().__init__(inner) + self._now = now + self._replan_at: float | None = None + + def __call__(self, obs): + if self._replan_at is not None and self._now() < self._replan_at: + return None + result = self._inner(obs) + assert result is not None, 'the inner policy of this test wrapper always returns a chunk' + now = self._now() + result = [{**action, 'timestamp': now + action['timestamp']} for action in result] + self._replan_at = now + (result[-1]['timestamp'] - now) / 2 + return result + + def wrap_session(self, inner: Session, context, now): + return _ReplanEarly._Session(inner, now) + + +class _TimedRecorder(pimm.SignalEmitter): + """Records each emission against the world clock, so a test can read when a command actually went out.""" + + def __init__(self, clock: pimm.Clock): + self._clock = clock + self.emitted: list[tuple[float, Any]] = [] + + def emit(self, data, ts: int = -1): + self.emitted.append((self._clock.now(), data)) + + +def _run_sim_episode(world, policy, wrapper, *, latency, steps=4000, run_sec=1.5) -> list[tuple[float, Any]]: + """One sim trial under ``latency``; returns the grip commands with the world time each went out at.""" + harness = Harness(wrapper.wrap(policy), make_embodiment(simulated=True)) + grip_recorder = _TimedRecorder(world.clock) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands['target_grip']._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=latency)), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, run_sec), + ]) + drive_scheduler(world.start([harness, driver, _Pacer()]), steps=steps) + return grip_recorder.emitted[1:] # drop the startup home + + +@pytest.mark.timeout(20.0) +def test_default_latency_pauses_the_world_for_the_call(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 call really took.""" + played = _run_sim_episode(world, SlowPolicy(wall_sec=0.05), ChunkedSchedule(), latency=False) + + assert played, 'no command was played' + assert played[0][0] < 0.01, f'the world advanced during the call: first command at {played[0][0]}s' + + +@pytest.mark.timeout(20.0) +@pytest.mark.parametrize('wall_sec', [0.0, 0.05]) +def test_declared_latency_ignores_what_the_call_really_took(world, wall_sec): + """The reproducible mode: the wrapper is released a fixed delay after the call started, so the played + trace is the same against a fast server and a slow one.""" + played = _run_sim_episode(world, SlowPolicy(wall_sec=wall_sec), ChunkedSchedule(), latency=0.3) + + assert played, 'no command was played' + assert played[0][0] == pytest.approx(0.3, abs=0.02), f'first command at {played[0][0]}s, expected the 0.3s delay' + + +@pytest.mark.timeout(20.0) +def test_measured_latency_charges_the_calls_own_wall_duration(world): + """``inference_latency=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.""" + played = _run_sim_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=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 call took' + + +@pytest.mark.timeout(20.0) +def test_harness_keeps_playing_while_a_call_is_in_flight(world): + """A wrapper that replans before its chunk is exhausted leaves waypoints due during inference, and the + harness emits them on time instead of standing still until the model answers.""" + played = _run_sim_episode(world, SlowPolicy(span_sec=0.4, steps=20), _ReplanEarly(), latency=0.15) + + # The second call starts halfway through the first chunk (0.2s in) and is owed 0.15s; the waypoints due + # in that window have to keep going out. + during = [t for t, _ in played if 0.2 <= t < 0.35] + assert len(during) >= 3, f'the harness stopped playing during inference: {[t for t, _ in played]}' + + +@pytest.mark.timeout(3.0) +def test_installed_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): + """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): + 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 [{keys.ROBOT_COMMAND: command, 'target_grip': 0.5, 'timestamp': i * 0.01} for i in range(10)] + return [{keys.ROBOT_COMMAND: command, 'timestamp': i * 0.01} for i in range(10)] + + class _GripThenArmPolicy(Policy): + def new_session(self, context=None, now=None, gate=None): + return _GripThenArm() + + harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) + grip_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands['target_grip']._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, 0.5), + ]) + drive_scheduler(world.start([harness, driver]), steps=1000) + + grips = _emitted_grips(grip_recorder) + assert grips[0] == 0.0 # the startup home + assert set(grips[1:]) == {0.5}, f'the second chunk kept the gripper playing: {grips}' + + +@pytest.mark.timeout(3.0) +def test_home_and_manual_commands_are_emitted_as_plain_values(world): + """Homing and operator commands bypass the schedule: they are the command, not a plan to play.""" + harness = Harness(StubPolicy(), make_embodiment()) + cmd_recorder = RecordingEmitter() + grip_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands['target_grip']._bind(grip_recorder) + harness.ds_command._bind(RecordingEmitter()) + manual_em = world.pair(harness.manual_command) + + pose = Transform3D(translation=np.array([0.1, 0.1, 0.1], dtype=np.float32), rotation=Rotation.identity) + manual = CartesianPosition(pose=pose) + driver = ManualDriver([(partial(manual_em.emit, {keys.ROBOT_COMMAND: manual}), 0.01), (None, 0.02)]) + drive_scheduler(world.start([harness, driver]), steps=50) + + assert _emitted_commands(cmd_recorder) == [Reset(), manual] + assert _emitted_grips(grip_recorder) == [0.0] + + +@pytest.mark.timeout(20.0) +def test_abort_discards_a_call_that_is_still_in_flight(world): + """An ABORT while the gate is still holding the model's answer throws that answer away: the trajectory + it carries never reaches the devices.""" + harness = Harness(ChunkedSchedule().wrap(SlowPolicy()), make_embodiment(simulated=True)) + cmd_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands['target_grip']._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=1.0)), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (None, 0.05), # well inside the 1.0s the gate owes the call + (partial(directive_em.emit, Directive.ABORT()), 0.0), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver, _Pacer()]), steps=2000) + + assert all(isinstance(c, Reset) for c in _emitted_commands(cmd_recorder)) diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index 452637d3b..89dc306ea 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -164,12 +164,12 @@ class _ChunkPolicy(Policy): def __init__(self, actions: list[dict]): self._actions = actions - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _FixedSession(list(self._actions)) class _SinglePolicy(Policy): - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _FixedSession({'v': 42}) @@ -183,7 +183,7 @@ def encode(self, data): class _MetaPolicy(Policy): - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _FixedSession({}) @property diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index e88acba1a..7448c8cb1 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -27,7 +27,7 @@ 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 - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self.session_count += 1 return _TrackingSession(self._actions, {'policy_key': 'policy_value'}) @@ -57,7 +57,7 @@ def __init__(self, rec, actions): self._actions = actions self.last_session = None - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self.last_session = _CapturingSession(self._rec, self._actions) return self.last_session @@ -146,7 +146,7 @@ def __call__(self, obs): return None class _NonePolicy(Policy): - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _NoneSession() session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 6a19ab5fd..19f08f0ae 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -1,5 +1,10 @@ -"""Unit tests for PolicyWrapper composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" +"""Unit tests for PolicyWrapper composition, the InferenceGate, ChunkedSchedule, TemporalStack, and the +policy-pipeline algebra.""" +import concurrent.futures +import threading +import time +from concurrent.futures import ThreadPoolExecutor from typing import Any import numpy as np @@ -15,7 +20,7 @@ JointDeltaAction, RelativePositionAction, ) -from positronic.policy.base import Policy, PolicyWrapper, Session +from positronic.policy.base import InferenceGate, LatencyMode, Policy, PolicyWrapper, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -58,7 +63,7 @@ def __init__(self, actions): self._actions = actions self._session: _ConstSession | None = None - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self._session = _ConstSession(self._actions) return self._session @@ -225,7 +230,7 @@ class _CapturePolicy(Policy): def __init__(self): self.session = _CaptureSession() - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return self.session @@ -571,3 +576,115 @@ def test_survives_a_wire_round_trip(self): rebuilt = spec.from_spec(RestrictImageSize(64, 48).to_spec()) assert isinstance(rebuilt, RestrictImageSize) assert rebuilt.encode({'cam': _image(480, 640)})['cam'].shape == (48, 64, 3) + + +class _BlockingSession(Session): + """Holds the call inside the gate until the test releases it, and reports when it was entered.""" + + def __init__(self): + self.entered = threading.Event() + self.release = threading.Event() + + def __call__(self, obs): + self.entered.set() + self.release.wait(timeout=5.0) + return [{'v': 1, 'timestamp': 0.0}] + + +class _SlowSession(Session): + """Spends ``wall_sec`` of real time in the model, which is what the measured mode charges.""" + + def __init__(self, wall_sec: float): + self._wall_sec = wall_sec + + def __call__(self, obs): + time.sleep(self._wall_sec) + return [{'v': 1, 'timestamp': 0.0}] + + +def _run_gated(gate: InferenceGate, inner: Session, executor: ThreadPoolExecutor): + return executor.submit(gate.wrap(inner), _obs()) + + +class TestInferenceGate: + def test_declared_mode_parks_the_call_until_its_delay_has_passed(self): + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) + with ThreadPoolExecutor(max_workers=1) as executor: + future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + concurrent.futures.wait([future], timeout=0.05) + assert not future.done(), 'the gate let the call through before its declared delay' + clock.t = 1.5 + assert future.result(timeout=5.0) is not None + assert gate.t0 == 1.0 + + def test_measured_mode_parks_the_call_for_its_own_wall_duration(self): + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.MEASURED) + with ThreadPoolExecutor(max_workers=1) as executor: + future = _run_gated(gate, _SlowSession(0.1), executor) + concurrent.futures.wait([future], timeout=0.3) + assert not future.done(), 'the gate charged nothing for a call that took 0.1s' + clock.t = 1.5 # past t0 + the wall duration, whatever it measured + assert future.result(timeout=5.0) is not None + + @pytest.mark.parametrize( + 'mode,delay_sec', [(LatencyMode.LIVE, 0.0), (LatencyMode.DECLARED, 0.0)], ids=['live', 'declared-zero'] + ) + def test_a_call_owed_nothing_is_never_parked(self, mode, delay_sec): + """Hardware pays what the model took, and sim's default charges nothing: neither owes the wrapper a + wait, so the call returns with the clock standing still.""" + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, mode, delay_sec) + with ThreadPoolExecutor(max_workers=1) as executor: + future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + assert future.result(timeout=1.0) is not None # the clock never moves and the call still returns + assert gate.entered is False + + def test_entered_is_visible_while_the_call_is_at_the_model(self): + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.LIVE) + inner = _BlockingSession() + with ThreadPoolExecutor(max_workers=1) as executor: + future = _run_gated(gate, inner, executor) + assert inner.entered.wait(timeout=5.0) + assert gate.entered is True + inner.release.set() + future.result(timeout=5.0) + assert gate.entered is False + + def test_cancel_unparks_a_call_the_harness_no_longer_wants(self): + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) + with ThreadPoolExecutor(max_workers=1) as executor: + future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + concurrent.futures.wait([future], timeout=0.05) + assert not future.done() + gate.cancel() + assert future.result(timeout=5.0) is not None # released without waiting out the 60s delay + + def test_gate_wraps_the_inner_session_of_a_scheduling_wrapper(self): + """The cost lands below the wrapper: by the time ``ChunkedSchedule`` anchors its chunk, the delay + has already been paid, so the anchor is the release instant.""" + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) + policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + session = policy.new_session(now=clock.now, gate=gate) + with ThreadPoolExecutor(max_workers=1) as executor: + future = executor.submit(session, _obs()) + concurrent.futures.wait([future], timeout=0.05) + assert not future.done() + clock.t = 1.5 + chunk = future.result(timeout=5.0) + assert chunk is not None and chunk[0]['timestamp'] == 1.5 + + def test_a_stack_without_a_scheduling_wrapper_gets_no_gate(self): + """``TemporalStack`` owns no plan, so it is not what the platform charges; its call runs straight + through and the gate is never entered.""" + clock = _FakeClock(t=1.0) + gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) + policy = TemporalStack(keys=('v',), offsets_sec=(0.0,)).wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + session = policy.new_session(now=clock.now, gate=gate) + + assert session({**_obs(), 'v': np.array([1.0])}) is not None + assert gate.entered is False diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index 7b0e1d7d2..219c82bd5 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -12,7 +12,7 @@ import numpy as np from positronic import keys -from positronic.policy.base import DelegatingSession, Now, PolicyWrapper, Session +from positronic.policy.base import DelegatingSession, Now, PolicyWrapper, SchedulingWrapper, Session def _obs_time(obs) -> float: @@ -20,7 +20,7 @@ def _obs_time(obs) -> float: return obs[keys.OBS_TIME_NS] / 1e9 -class ChunkedSchedule(PolicyWrapper): +class ChunkedSchedule(SchedulingWrapper): """Wait for the current trajectory to finish before calling the inner policy again. Owns relative→absolute time conversion: inner layers (codecs, models) emit relative timestamps; diff --git a/positronic/replay_record.py b/positronic/replay_record.py index 1525e561b..a490b06ca 100644 --- a/positronic/replay_record.py +++ b/positronic/replay_record.py @@ -25,13 +25,6 @@ from positronic.utils.logging import init_logging -class _TrajectoryEmitter(pimm.ControlSystemEmitter): - """Sends each replayed value as the one-waypoint trajectory the command channels carry, at its playback time.""" - - def emit(self, data: Any, ts: int = -1) -> None: - super().emit([(ts, data)], ts) - - class Replay(DsPlayerAgent): """Adapts `DsPlayerAgent` to be used as a policy control system.""" @@ -42,8 +35,8 @@ def __init__(self, poll_hz: float = 100.0): self.gripper_state = pimm.FakeReceiver(self) self.robot_meta_in = pimm.FakeReceiver(self) self.frames = pimm.ReceiverDict(self, fake=True) - self.outputs['robot_commands'] = _TrajectoryEmitter(self) - self.outputs[keys.TARGET_GRIP] = _TrajectoryEmitter(self) + self.outputs['robot_commands'] = pimm.ControlSystemEmitter(self) + self.outputs[keys.TARGET_GRIP] = pimm.ControlSystemEmitter(self) @property def robot_commands(self) -> pimm.ControlSystemEmitter: diff --git a/positronic/robot_controller.py b/positronic/robot_controller.py index 07fbc67f2..4df95f6f8 100644 --- a/positronic/robot_controller.py +++ b/positronic/robot_controller.py @@ -1,4 +1,5 @@ import configuronic as cfn +import numpy as np import pimm import positronic.cfg.hardware.roboarm @@ -18,15 +19,15 @@ def main(robot): match text_command.split(' '): case ['reset']: - command_emmiter.emit([(world.clock.now_ns(), command.Reset())]) + command_emmiter.emit(command.Reset()) case ['move', x, y, z, qw, qx, qy, qz]: pos = [float(x) for x in [x, y, z]] quat = geom.Rotation.from_quat([float(qw), float(qx), float(qy), float(qz)]) move = command.CartesianPosition(geom.Transform3D(translation=pos, rotation=quat)) - command_emmiter.emit([(world.clock.now_ns(), move)]) + command_emmiter.emit(move) case ['joint_move', *args]: args = [float(x) for x in args] - command_emmiter.emit([(world.clock.now_ns(), command.JointPosition(positions=args))]) + command_emmiter.emit(command.JointPosition(positions=np.asarray(args))) case ['info']: print('Q', state_receiver.value.q) print('DQ', state_receiver.value.dq) diff --git a/positronic/simulator/env_server/adapter.py b/positronic/simulator/env_server/adapter.py index eed79fee5..6fc09c0db 100644 --- a/positronic/simulator/env_server/adapter.py +++ b/positronic/simulator/env_server/adapter.py @@ -1,14 +1,13 @@ """The ``EnvAdapter`` interface: the per-benchmark canonical<->raw mappings, on the client side. ``RemoteEnvControlSystem`` is a dumb translator — it moves data between pimm signals and this adapter. -The adapter is the smart half: it turns the Harness's command trajectories into the env's raw action -(owning trajectory playing and how to hold between waypoints), maps raw observations back to canonical -signals — policy-facing and privileged ground-truth kept separate — and reads the terminal. Each -benchmark ships one adapter (``vendors/``-style); the native ``MujocoSim`` fixture is the reference. +The adapter is the smart half: it turns the Harness's commands into the env's raw action (owning what to +hold between them), maps raw observations back to canonical signals — policy-facing and privileged +ground-truth kept separate — and reads the terminal. Each benchmark ships one adapter (``vendors/``-style); +the native ``MujocoSim`` fixture is the reference. """ from abc import ABC, abstractmethod -from collections import defaultdict from typing import Any, final import numpy as np @@ -18,14 +17,6 @@ from positronic.drivers.roboarm import command as roboarm_command -def fresh_command_players() -> defaultdict[str, roboarm_command.TrajectoryPlayer]: - """A trajectory player per command channel: ``robot_command`` accumulates the deltas due in one tick (a - missed tick catches up instead of dropping motion), every other channel keeps the last value due.""" - players = defaultdict(roboarm_command.TrajectoryPlayer) - players[keys.ROBOT_COMMAND] = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - return players - - class EnvAdapter(ABC): """The mappings between the canonical embodiment contract and an env's raw wire payloads.""" @@ -38,11 +29,11 @@ def reset_token(self, context: dict[str, Any]) -> Any: """ @abstractmethod - def action(self, commands: dict[str, pimm.Message], now_ns: int) -> dict[str, Any]: - """The latest per-channel command messages + the clock -> the raw action the env steps. + def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: + """The latest per-channel command messages -> the raw action the env steps. - The adapter owns trajectory playing (sampling each channel's waypoints down to ``now_ns``) and - what to do between waypoints — e.g. hold the last commanded value, the absolute-mode invariant. + A channel delivers a command only when one comes due, so the adapter owns what happens in between — + e.g. hold the last commanded value, the absolute-mode invariant. """ @abstractmethod @@ -102,10 +93,9 @@ def _wire_command(cmd: Any) -> dict[str, Any]: class WireCommandAdapter(EnvAdapter): """An adapter whose action is the shared wire payload ``{'command': , 'grip': float}``. - The command side of every remote benchmark adapter: it plays each command channel's trajectory down to - the clock — holding an absolute setpoint between waypoints, firing a relative delta once — and flattens - the held arm command (a pose as ``[t(3), R(9)]``, joint positions, or per-step joint deltas) plus the - gripper closure into one payload. + The command side of every remote benchmark adapter: it holds an absolute setpoint until the next command + arrives and fires a relative delta once, and flattens the held arm command (a pose as ``[t(3), R(9)]``, + joint positions, or per-step joint deltas) plus the gripper closure into one payload. All action *encoding* — how the tagged command becomes the env's native action — stays server-side with the env's own model. Subclasses implement ``_reset_token`` (the base clears the per-trial command state around it) and keep the observation and terminal mappings to themselves. @@ -118,10 +108,9 @@ def __init__(self, env_control_frame: geom.Transform3D | None = None): self._reset_command_state() def _reset_command_state(self) -> None: - self._players = fresh_command_players() - self._held: dict[str, Any] = {} # last sampled waypoint per channel — re-sent until it changes - # Last commanded gripper closure, held across a cancelled grip trajectory: grip is an absolute [0, 1] - # value with no 'hold' command to fall back on (unlike the arm), so cancelling must freeze it, not reopen. + self._held: dict[str, Any] = {} # last command per channel — re-sent until the next one arrives + # Last commanded gripper closure, held across an episode boundary: grip is an absolute [0, 1] value + # with no 'hold' command to fall back on (unlike the arm), so a fresh trial must freeze it, not reopen. self._grip = 0.0 @final @@ -133,21 +122,14 @@ def reset_token(self, context: dict[str, Any]) -> Any: def _reset_token(self, context: dict[str, Any]) -> Any: """The per-trial RUN context -> the env's opaque reset token; the command state is already cleared.""" - def action(self, commands: dict[str, pimm.Message], now_ns: int) -> dict[str, Any]: + def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: for name, msg in commands.items(): - player = self._players[name] if msg.updated: - player.set(msg.data) - if not msg.data: # an empty trajectory cancels: stop replaying the held waypoint - self._held.pop(name, None) - value = player.advance(now_ns) - if value is not None: - self._held[name] = value + self._held[name] = msg.data # The server maps the held command into its controller's action. Reset has no env-side action, so it # forwards as a hold; a delta — Cartesian or joint — is a one-shot relative motion, forwarded once then # dropped. Re-sending a stale delta would re-compose it against the moving arm every tick (the eef - # drifts, or the joints walk toward their limits), so once a delta's trajectory is exhausted the arm - # holds its measured pose. + # drifts, or the joints walk toward their limits), so after one step the arm holds its measured pose. cmd = self._held.get(keys.ROBOT_COMMAND) match cmd: case roboarm_command.Reset(): diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index e8a39aae7..e801d1b18 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -5,7 +5,7 @@ to, but no command logic. Each control period it hands the latest command messages to the ``EnvAdapter``, round-trips the raw action it returns over the wire, and re-emits the canonical signals the adapter maps back — so only raw arrays cross the boundary and the World's virtual clock advances by the env's -``control_dt`` per step. The adapter owns trajectory playing, holding, and the canonical<->raw mappings; +``control_dt`` per step. The adapter owns holding and the canonical<->raw mappings; ``control_dt`` is whatever the latest observation reports (``reset`` and every ``step``). """ @@ -36,7 +36,7 @@ def __init__(self, adapter: EnvAdapter, serve: AbstractContextManager[tuple[str, self._cleanup = ExitStack() self._conn: EnvConnection | None = None - self.commands: pimm.ReceiverDict = pimm.ReceiverDict(self, default=[]) + self.commands: pimm.ReceiverDict = pimm.ReceiverDict(self) self.observations: pimm.EmitterDict = pimm.EmitterDict(self) self.privileged: pimm.EmitterDict = pimm.EmitterDict(self) self.robot_meta: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) @@ -119,7 +119,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p # observation assembly (shared-memory image allocation + camera copies) inside it, so the # reduce can split materialisation out of the wire cost. with telemetry.span(telemetry_keys.SPAN_ENV_STEP): - self._frame = self._step_env(clock) + self._frame = self._step_env() with telemetry.span(telemetry_keys.SPAN_MATERIALIZE): self._emit_payload(self._frame['obs']) finally: @@ -127,9 +127,9 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p # ever connected. self._cleanup.close() - def _step_env(self, clock: pimm.Clock) -> dict[str, Any]: + def _step_env(self) -> dict[str, Any]: commands = {name: receiver.read() for name, receiver in self.commands.items()} - result = self._conn.step(self._adapter.action(commands, clock.now_ns())) + result = self._conn.step(self._adapter.action(commands)) payload = self._adapter.terminal(result) if payload: # truthy-valued done: a non-empty payload ends the trial, an empty/``None`` one continues self.done.emit(payload) diff --git a/positronic/simulator/env_server/tests/mujoco_env.py b/positronic/simulator/env_server/tests/mujoco_env.py index fdff9ec66..416ce45c8 100644 --- a/positronic/simulator/env_server/tests/mujoco_env.py +++ b/positronic/simulator/env_server/tests/mujoco_env.py @@ -133,10 +133,9 @@ def step(self, action: dict[str, Any]) -> dict[str, Any]: cmd = roboarm_command.CartesianDelta(geom.Transform3D.from_vector(command['delta'], _ROTMAT)) case other: raise ValueError(f'MujocoEnv got unsupported command type {other!r}') - now_ns = self._clock.now_ns() if cmd is not None: - self._cmd_emit.emit([(now_ns, cmd)]) - self._grip_emit.emit([(now_ns, float(action['grip']))]) + self._cmd_emit.emit(cmd) + self._grip_emit.emit(float(action['grip'])) self._advance(self._timestep) return {'obs': self._read_obs(), 'done': False, 'control_dt': self._timestep} diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 76bc0818f..b8d30edbf 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -211,7 +211,7 @@ class _CountdownAdapter(EnvAdapter): def reset_token(self, context): return context.get('eval.seed') - def action(self, commands, now_ns): + def action(self, commands): return {} def observations(self, raw_obs): @@ -312,7 +312,7 @@ def __init__(self, command: roboarm_command.CommandType, chunk_len: int): self.chunk_len = chunk_len self.chunks = 0 - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): return _JointposChunkSession(self) diff --git a/positronic/simulator/mujoco/sim.py b/positronic/simulator/mujoco/sim.py index 8ab1852a6..2fd3996f8 100644 --- a/positronic/simulator/mujoco/sim.py +++ b/positronic/simulator/mujoco/sim.py @@ -99,9 +99,9 @@ class MujocoSim(pimm.ControlSystem): ``reset`` rebuilds the scene and flags frame-0 publication; the run loop publishes that post-reset scene on its next turn — in sequence, before any step — so the first inference reads it and the - recorder logs it. Every other turn applies the due command waypoints, steps once, and emits the due - streams (post-step, Gym-style). The sim sleeps one control period each turn, so it is the eval's sole - time-master. Each stream has an independent rate (``*_fps``, ``None`` = every physics tick). + recorder logs it. Every other turn applies whatever command has just arrived, steps once, and emits + the due streams (post-step, Gym-style). The sim sleeps one control period each turn, so it is the + eval's sole time-master. Each stream has an independent rate (``*_fps``, ``None`` = every physics tick). """ def __init__( @@ -140,20 +140,14 @@ def __init__( self._home() self._error = False self._adapters: dict[str, pimm.shared_memory.NumpySMAdapter] | None = None - self._arm_player = roboarm_command.TrajectoryPlayer(reduce=roboarm_command.reduce) - self._grip_player = roboarm_command.TrajectoryPlayer() self._last_grip = 0.0 # Set by ``reset``; the run loop publishes frame-0 (instead of stepping) on its next turn and clears it. self._reset_pending = False - self.commands: pimm.SignalReceiver[roboarm_command.Trajectory[roboarm_command.CommandType]] = ( - pimm.ControlSystemReceiver(self, default=[]) - ) + self.commands = pimm.ControlSystemReceiver[roboarm_command.CommandType](self) self.state: pimm.SignalEmitter[MujocoFrankaState] = pimm.ControlSystemEmitter(self) self.robot_meta = pimm.ControlSystemEmitter(self) - self.target_grip: pimm.SignalReceiver[roboarm_command.Trajectory[float]] = pimm.ControlSystemReceiver( - self, default=[] - ) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.grip: pimm.SignalEmitter = pimm.ControlSystemEmitter(self) self.cameras: pimm.EmitterDict = pimm.EmitterDict(self) # Privileged ground truth: the full ``save_state`` dict, spec keys prefixed with '.' so the @@ -183,23 +177,13 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p continue now = clock.now() cmd_msg = self.commands.read() - if cmd_msg.updated: - self._arm_player.set(cmd_msg.data) if self._error: - self._error = False - # Drop the in-flight trajectory so the arm holds position rather than resuming a stale - # waypoint once the error clears. - self._arm_player.set([]) - else: - cmd = self._arm_player.advance(clock.now_ns()) - if cmd is not None: - self._apply_command(cmd) + self._error = False # the command that arrived with the error is skipped, not queued + elif cmd_msg is not None and cmd_msg.updated: + self._apply_command(cmd_msg.data) grip_msg = self.target_grip.read() - if grip_msg.updated: - self._grip_player.set(grip_msg.data) - grip = self._grip_player.advance(clock.now_ns()) - if grip is not None: - self._last_grip = grip + if grip_msg is not None and grip_msg.updated: + self._last_grip = grip_msg.data self._apply_grip(self._last_grip) # An env step is the sim advance plus the observations it produces, rendering included @@ -228,17 +212,14 @@ def reset(self, seed: int | None = None): colors, cameras) re-randomize too; the renderer and IK physics rebind lazily. The run loop publishes the prepared scene as frame-0 on its next turn — in sequence, before any step — so the first inference reads the reset state and the recorder logs it. Stale commands queued while idle - (e.g. the inter-episode home) are dropped and the run-loop's trajectory players and held grip are - cleared, so the first step neither applies a queued command nor replays the previous episode's - trajectory on the freshly reset scene. + (e.g. the inter-episode home) are dropped and the held grip is cleared, so the first step does not + apply a queued command on the freshly reset scene. """ self._load_scene(seed) self._home() self._error = False self.commands.read() self.target_grip.read() - self._arm_player.set([]) - self._grip_player.set([]) self._last_grip = 0.0 self._reset_pending = True diff --git a/positronic/tests/test_data_collection.py b/positronic/tests/test_data_collection.py index f114fad74..56a1c1ded 100644 --- a/positronic/tests/test_data_collection.py +++ b/positronic/tests/test_data_collection.py @@ -7,12 +7,7 @@ import pimm from positronic import keys, wire from positronic.data_collection import DataCollectionController, OperatorPosition, controller_positions_serializer -from positronic.dataset.ds_writer_agent import ( - DsWriterAgent, - DsWriterCommand, - DsWriterCommandType, - TrajectoryOverrideSerializer, -) +from positronic.dataset.ds_writer_agent import DsWriterAgent, DsWriterCommand, DsWriterCommandType from positronic.dataset.episode import Episode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers @@ -162,8 +157,8 @@ def test_data_collection_with_mujoco_robot_gripper(tmp_path): writer_cm = LocalDatasetWriter(tmp_path) agent = DsWriterAgent(writer_cm.__enter__()) - agent.add_signal('target_grip', TrajectoryOverrideSerializer(None)) - agent.add_signal(keys.ROBOT_COMMAND, TrajectoryOverrideSerializer(Serializers.robot_command)) + agent.add_signal('target_grip') + agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) agent.add_signal('controller_positions', controller_positions_serializer) agent.add_signal('robot_state', Serializers.robot_state) agent.add_signal(keys.GRIP) @@ -258,9 +253,9 @@ def finger_qpos() -> float: grip = world.pair(sim.grip) driver = ManualDriver([ - (lambda: target_grip.emit([(world.clock.now_ns(), 1.0)]), 0.5), + (lambda: target_grip.emit(1.0), 0.5), (lambda: snapshots.update(closed=finger_qpos(), closed_grip=grip.read().data), 0.0), - (lambda: target_grip.emit([(world.clock.now_ns(), 0.0)]), 0.5), + (lambda: target_grip.emit(0.0), 0.5), (lambda: snapshots.update(opened=finger_qpos(), opened_grip=grip.read().data), 0.0), ]) diff --git a/positronic/utils/registration.py b/positronic/utils/registration.py index 5cfcb30ab..cc2c8ad25 100644 --- a/positronic/utils/registration.py +++ b/positronic/utils/registration.py @@ -186,7 +186,7 @@ def perform_registration(webxr, robot_arm): while not w.should_stop: if move_throttler.wait_time() <= 0: waypoint = roboarm_command.CartesianPosition(WAYPOINTS[current_point]) - commands.emit([(w.clock.now_ns(), waypoint)]) + commands.emit(waypoint) current_point += 1 if current_point >= len(WAYPOINTS): break diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index 01c8c05f4..d4aac37c1 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -286,7 +286,7 @@ class DreamZeroPolicy(Policy): def __init__(self, sp: DreamZeroSubprocess): self._subprocess = sp - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): client = RoboarenaClient(port=self._subprocess.roboarena_port) client.connect() return _DreamZeroSession(client, str(uuid.uuid4())) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 6071fb46e..0258e735f 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -220,7 +220,7 @@ def __init__(self, groot: Gr00tSubprocess, checkpoint_path: str): self._groot = groot self._checkpoint_path = checkpoint_path - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self._groot.client.reset() return _Gr00tSession(self._groot.client) diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index fa83c9922..62746c864 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -102,7 +102,7 @@ def config(self) -> PreTrainedConfig: """The checkpoint's own declaration of what this policy takes.""" return self._policy.config - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self._policy.reset() return _LerobotSession(self._policy, self._preprocessor, self._postprocessor, self._device, self._meta) diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 5e608c1b1..c20e03a33 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -92,7 +92,7 @@ def __init__(self, policy: PreTrainedPolicy, device: str | None = None, extra_me self._policy = policy.to(self._device) self._meta = extra_meta or {} - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): self._policy.reset() return _LerobotSession(self._policy, self._device, self._meta) diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index b61a2e0e2..b18853c63 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -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, now=None) -> Session: + def new_session(self, context=None, now=None, gate=None) -> Session: 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 519b0ad37..ff9e0c174 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -140,7 +140,7 @@ class OpenpiPolicy(Policy): def __init__(self, subproc: OpenpiSubprocess): self._subproc = subproc - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None, gate=None): client = self._subproc.client client.reset() return _OpenpiSession(client) diff --git a/positronic/wire.py b/positronic/wire.py index cdc15af7d..6fa0ef1c2 100644 --- a/positronic/wire.py +++ b/positronic/wire.py @@ -3,7 +3,7 @@ import pimm from positronic import keys, telemetry, telemetry_keys from positronic.dataset import DatasetWriter -from positronic.dataset.ds_writer_agent import DsWriterAgent, TimeMode, TrajectoryOverrideSerializer +from positronic.dataset.ds_writer_agent import DsWriterAgent, TimeMode from positronic.dataset.serializers import Serializers, StatefulSerializer from positronic.eval import ROBOT_STATIC_META, Embodiment, Observation @@ -44,12 +44,10 @@ def wire( # noqa: C901 for signal_name in cameras.keys(): ds_agent.add_signal(signal_name, Serializers.camera_images) if robot_arm is not None: - # Command channels carry whole trajectories; flatten with last-writer-wins so the - # recording is a dense per-command stream. See TrajectoryOverrideSerializer. - ds_agent.add_signal(keys.ROBOT_COMMAND, TrajectoryOverrideSerializer(Serializers.robot_command)) + ds_agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) ds_agent.add_signal('robot_state', Serializers.robot_state) if gripper is not None: - ds_agent.add_signal(keys.TARGET_GRIP, TrajectoryOverrideSerializer(None)) + ds_agent.add_signal(keys.TARGET_GRIP) ds_agent.add_signal(keys.GRIP) for signal_name, emitter in cameras.items(): @@ -109,9 +107,7 @@ def wire_embodiment( ds_agent.add_signal(name, obs.serializer) world.connect(obs.source, ds_agent.inputs[name]) for name, cmd in embodiment.commands.items(): - # Command channels carry whole trajectories; flatten with last-writer-wins so the - # recording is a dense per-command stream. See TrajectoryOverrideSerializer. - ds_agent.add_signal(name, TrajectoryOverrideSerializer(cmd.serializer)) + ds_agent.add_signal(name, cmd.serializer) world.connect(harness.commands[name], ds_agent.inputs[name]) for name, priv in privileged.items(): ds_agent.add_signal(name, priv.serializer) diff --git a/utilities/fake_dataset_generator.py b/utilities/fake_dataset_generator.py index 3523927c8..aa0aba0ab 100644 --- a/utilities/fake_dataset_generator.py +++ b/utilities/fake_dataset_generator.py @@ -138,7 +138,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): 'eval.tote_placement': random.choice(['left', 'right']), 'eval.external_camera': random.choices(['left', 'right', 'NA'], [5, 5, 1])[0], 'inference.policy_fps': self.fps, - 'inference_latency': False, + keys.INFERENCE_LATENCY: False, **self.policy_meta, } From 78f1027fb894cdbabc21a7940b8d599283e8085b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Mon, 10 Aug 2026 12:44:18 +0300 Subject: [PATCH 02/64] Polish per-tick command wire: trim comments, move `TrajectoryPlayer` to the harness Move `TrajectoryPlayer` and `Trajectory` out of `command.py` into `harness.py`, so the wire module is only the command types and their serialization. Make `new_session`'s `now`/`gate` keyword-only. Filter empty command reads in the env proxy rather than in the adapter. --- .basedpyright/baseline.json | 8 --- ARCHITECTURE.md | 6 +- positronic/drivers/roboarm/command.py | 36 +---------- positronic/drivers/roboarm/franka.py | 4 +- positronic/drivers/roboarm/yam.py | 6 +- positronic/keys.py | 8 +-- .../offboard/tests/test_remote_policy.py | 2 +- positronic/offboard/tests/test_server.py | 2 +- positronic/policy/base.py | 19 +++--- positronic/policy/harness.py | 63 +++++++++++++------ positronic/policy/recording.py | 4 +- positronic/policy/remote.py | 6 +- .../policy/tests/test_golden_pipeline.py | 4 +- positronic/policy/tests/test_harness.py | 36 ++++------- positronic/policy/tests/test_policy_io.py | 6 +- positronic/policy/tests/test_recording.py | 6 +- positronic/policy/tests/test_wrappers.py | 42 ++++++------- positronic/policy/wrappers.py | 13 ++-- positronic/probe.py | 2 +- positronic/simulator/env_server/adapter.py | 8 +-- positronic/simulator/env_server/proxy.py | 3 +- .../env_server/tests/test_remote_env.py | 2 +- positronic/simulator/mujoco/sim.py | 2 +- positronic/vendors/dreamzero/server.py | 2 +- positronic/vendors/gr00t/server.py | 2 +- positronic/vendors/lerobot/policy.py | 2 +- positronic/vendors/lerobot_0_3_3/policy.py | 2 +- positronic/vendors/molmoact2/policy.py | 2 +- positronic/vendors/openpi/server.py | 2 +- 29 files changed, 130 insertions(+), 170 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index efb280b64..4094a569c 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -6936,14 +6936,6 @@ "endColumn": 32, "lineCount": 1 } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 54, - "endColumn": 62, - "lineCount": 1 - } } ], "./positronic/simulator/env_server/server.py": [ diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index b9b764d36..9e3ae3180 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -36,9 +36,9 @@ library must supply the tools that make late binding possible — codecs, per-tr projections over raw recordings. **Every decision lives with the party that has the information.** Only a driver knows its motion -capabilities, so drivers own how they reach a commanded setpoint. Only a sensor knows its own cadence, so sensors -run at their own rate instead of a rate the loop imposes. Only a policy knows what its model was -trained on, so translation to model I/O ships with the policy. +capabilities, so drivers own how they reach a commanded setpoint. Only a sensor knows its own +cadence, so sensors run at their own rate instead of a rate the loop imposes. Only a policy knows +what its model was trained on, so translation to model I/O ships with the policy. **Components are functions over flowing data.** A component sees nothing but its inputs and touches nothing but its outputs. Whatever varies enters as data — time is an observation field, hardware diff --git a/positronic/drivers/roboarm/command.py b/positronic/drivers/roboarm/command.py index db136733e..f7c62e016 100644 --- a/positronic/drivers/roboarm/command.py +++ b/positronic/drivers/roboarm/command.py @@ -1,7 +1,7 @@ """Collection of commands that can be sent to the robot.""" from dataclasses import dataclass, field -from typing import Any, TypeAlias, TypeVar +from typing import Any import numpy as np @@ -80,12 +80,6 @@ def apply(self, current: geom.Transform3D) -> geom.Transform3D: CommandType = Reset | CartesianPosition | JointPosition | JointDelta | CartesianDelta -_T = TypeVar('_T') - -# A schedule the harness plays: waypoints stamped with absolute clock ns, ascending. Command channels -# themselves carry one value -- the command due at the moment it is emitted. -Trajectory: TypeAlias = list[tuple[int, _T]] - def to_wire(command: CommandType) -> dict[str, Any]: match command: @@ -105,34 +99,6 @@ def to_wire(command: CommandType) -> dict[str, Any]: } -class TrajectoryPlayer: - """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value - to emit.""" - - def __init__(self): - self._trajectory: Trajectory[Any] = [] - self._index: int = 0 - - def set(self, trajectory: Trajectory[Any]): - self._trajectory = trajectory - self._index = 0 - - def next_due(self) -> int | None: - """Timestamp of the earliest waypoint not yet played, or ``None`` once the schedule is exhausted.""" - return self._trajectory[self._index][0] if self._index < len(self._trajectory) else None - - def advance(self, current_time: int): - """The single value due at ``current_time``, or ``None`` when no waypoint has come due since the last - call. Several waypoints due at once collapse to the last: an absolute setpoint supersedes the ones it - overtook, and only a late round makes it happen. - """ - value = None - while self._index < len(self._trajectory) and self._trajectory[self._index][0] <= current_time: - value = self._trajectory[self._index][1] - self._index += 1 - return value - - def from_wire(wire: dict[str, Any]) -> CommandType: match wire['type']: case 'reset': diff --git a/positronic/drivers/roboarm/franka.py b/positronic/drivers/roboarm/franka.py index af4cd1a42..e8c31effb 100644 --- a/positronic/drivers/roboarm/franka.py +++ b/positronic/drivers/roboarm/franka.py @@ -316,7 +316,7 @@ def _park(self, robot) -> Iterator[pimm.Sleep]: except Exception: logging.exception('Parking failed, the arm stays where it stands') - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: with self._desk_session(): robot = self._ensure_robot() try: @@ -345,7 +345,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p if in_error: # The driver always clears a recoverable error itself; making it optional (hold in # ERROR for out-of-band recovery instead) is a config knob to add when an embodiment - # needs it. A command that arrived with the error is skipped, not queued. + # needs it. robot.recover_from_errors() yield rate_limiter.wait() continue diff --git a/positronic/drivers/roboarm/yam.py b/positronic/drivers/roboarm/yam.py index 3db22f89b..f4abd423e 100644 --- a/positronic/drivers/roboarm/yam.py +++ b/positronic/drivers/roboarm/yam.py @@ -184,13 +184,13 @@ def __init__( self._sim = sim self._connect = connect - self.commands = pimm.ControlSystemReceiver[command.CommandType](self, default=None) - self.target_grip = pimm.ControlSystemReceiver[float](self, default=None) + self.commands = pimm.ControlSystemReceiver[command.CommandType](self) + self.target_grip = pimm.ControlSystemReceiver[float](self) self.state = pimm.ControlSystemEmitter[YamState](self) self.grip = pimm.ControlSystemEmitter[float](self) self.robot_meta = pimm.ControlSystemEmitter[dict[str, Any]](self) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: arm = self._connect(self._channel, self._sim) try: kin = _Kinematics() diff --git a/positronic/keys.py b/positronic/keys.py index 09dbba2b7..9fddef819 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -79,11 +79,6 @@ def is_robot_command(name: str) -> bool: COMPRESS_IMAGES = 'compress_images' POSITRONIC_VERSION = 'positronic_version' -# The sim-only knob a trial's RUN context carries: what the platform charges the scheduling wrapper's -# model call. ``False`` (or absent) holds the world still for it, a float charges that fixed delay, ``True`` -# charges the call's own wall duration. -INFERENCE_LATENCY = 'inference_latency' - POLICY_META = 'inference.policy' SERVER_META = f'{POLICY_META}.{SERVER}' @@ -93,3 +88,6 @@ def is_robot_command(name: str) -> bool: # leaves it absent on failure — a reader defaults it rather than assuming a False. EVAL_SUCCESS = 'eval.success' EVAL_TERMINATED = 'eval.terminated' + +# The sim-only inference cost a trial's RUN context asks the platform to charge the model call. +INFERENCE_LATENCY = 'inference_latency' diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 65af1fbb5..8f0bec5d5 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -302,7 +302,7 @@ def test_no_truncation_without_horizon(self): def test_remote_session_normalizes_single_dict(): - """Server returning a single action dict (legacy shape) is wrapped into a 1-element list.""" + """Server returning a single action dict is wrapped into a 1-element list.""" endpoint, _ = _mock_endpoint(infer_return={keys.ROBOT_COMMAND: 'X', 'timestamp': 0.0}) session = endpoint.new_session() diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 0105b1f61..647619fdd 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -244,7 +244,7 @@ def __call__(self, obs): class _ScriptedPolicy(Policy): """Deterministic base policy: every session returns the same untimestamped chunk.""" - def new_session(self, context=None, now=None, gate=None) -> Session: + def new_session(self, context=None, *, now=None, gate=None) -> Session: return _ScriptedSession() diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 4645c3bda..6d5b3dbe4 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -116,11 +116,10 @@ def hold(self) -> float | None: case LatencyMode.LIVE: return 0.0 case LatencyMode.DECLARED: - return None if self._now() >= self.t0 + self._delay_sec else 0.0 + return None if self._now() >= self._release_at() else 0.0 case LatencyMode.MEASURED: - # The world may run no further ahead of the call's start than wall time has: measured - # charging only means anything with the world at or below real time during the call. - return max(0.0, (self._now() - self.t0) - (time.monotonic() - self._wall_t0)) + # The world may run no further ahead of the call's start than wall time has. + return max(0.0, self._now() - self._release_at()) def _release_at(self) -> float: match self._mode: @@ -132,7 +131,7 @@ def _release_at(self) -> float: return self.t0 + (time.monotonic() - self._wall_t0) class _Session(DelegatingSession): - """Charges the inner call, on whatever thread the harness dispatched it to.""" + """Charges the inner call.""" def __init__(self, inner: Session, gate: InferenceGate): super().__init__(inner) @@ -163,7 +162,7 @@ class Policy(ABC): @abstractmethod def new_session( - self, context: dict[str, Any] | None = None, now: Now | None = None, gate: InferenceGate | None = None + self, context: dict[str, Any] | None = None, *, now: Now | None = None, gate: InferenceGate | None = None ) -> Session: """Create a new inference session for an episode. @@ -191,8 +190,8 @@ class DelegatingPolicy(Policy): def __init__(self, inner: Policy): self._inner = inner - def new_session(self, context=None, now=None, gate=None): - return self._inner.new_session(context, now, gate) + def new_session(self, context=None, *, now=None, gate=None): + return self._inner.new_session(context, now=now, gate=gate) @property def meta(self): @@ -276,8 +275,8 @@ def __init__(self, inner: Policy, wrapper: PolicyWrapper): super().__init__(inner) self._wrapper = wrapper - def new_session(self, context=None, now=None, gate=None): - inner = self._inner.new_session(context, now, gate) + def new_session(self, context=None, *, now=None, gate=None): + inner = self._inner.new_session(context, now=now, gate=gate) if gate is not None and isinstance(self._wrapper, SchedulingWrapper): inner = gate.wrap(inner) return self._wrapper.wrap_session(inner, context, now) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 7b3f4b1a4..849065acb 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -4,7 +4,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum -from typing import Any +from typing import Any, TypeAlias from opentelemetry.trace import Span @@ -12,7 +12,6 @@ from positronic import keys, telemetry, telemetry_keys from positronic.dataset.ds_writer_agent import DsWriterCommand from positronic.dataset.serializers import expand_suffixed -from positronic.drivers.roboarm.command import TrajectoryPlayer from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Task from positronic.policy.base import InferenceGate, LatencyMode, Policy, Session @@ -27,6 +26,39 @@ # noticed, and with it the granularity every command timestamp is quantized to. POLL_PERIOD_SEC = 0.01 +# One channel's schedule: waypoints stamped with absolute clock ns, ascending. +Trajectory: TypeAlias = list[tuple[int, Any]] + + +class TrajectoryPlayer: + """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value + to emit.""" + + def __init__(self): + self._trajectory: Trajectory = [] + self._index: int = 0 + + def set(self, trajectory: Trajectory): + self._trajectory = trajectory + self._index = 0 + + def next_due(self) -> int | None: + """Timestamp of the earliest waypoint not yet played, or ``None`` once the schedule is exhausted.""" + return self._trajectory[self._index][0] if self._index < len(self._trajectory) else None + + def advance(self, current_time: int): + """The single value due at ``current_time``, or ``None`` when no waypoint has come due since the + last call. + + TODO: several waypoints due at once collapse to the last, which drops the motion of every delta but + the final one. Revisit if pacing turns out not to hold one waypoint due per round. + """ + value = None + while self._index < len(self._trajectory) and self._trajectory[self._index][0] <= current_time: + value = self._trajectory[self._index][1] + self._index += 1 + return value + def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: """Reject a chunk whose timestamps are not times on the harness clock.""" @@ -147,14 +179,10 @@ class Harness(pimm.ControlSystem): Handles directives (RUN/FINISH/ABORT) and dataset recording. Inference intelligence — scheduling, error recovery, blending, absolute time stamping — lives in the policy/session layer: the wrapper owns - the plan, the harness plays it. Each round emits at most one command per channel, the one due now, so a - command channel carries execution rather than intent and its granularity is the round — - ``POLL_PERIOD_SEC`` on real hardware, one control period in sim. The RUN context is handed whole to the - task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). - - The session call runs on a worker thread, so playing continues while the model does; the - ``InferenceGate`` installed below the scheduling wrapper is what charges that call the trial's - inference latency. + the plan, the harness plays it, one command per channel per round. The session call runs on a worker + thread so playing continues while the model does, and the ``InferenceGate`` installed below the + scheduling wrapper charges that call the trial's inference latency. The RUN context is handed whole to + the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's @@ -211,8 +239,6 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] - # TODO: a late round collapses several due waypoints to the last, which drops the motion of every - # delta but the final one. Revisit if pacing turns out not to hold one waypoint due per round. self._players = {name: TrajectoryPlayer() for name in embodiment.commands} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) @@ -327,7 +353,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} self._gate = self._new_gate(clock) - self._policy_session = self.policy.new_session(self.context, clock.now, self._gate) + self._policy_session = self.policy.new_session(self.context, now=clock.now, gate=self._gate) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -405,9 +431,9 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: def _step(self, clock: pimm.Clock) -> None: """Keep one session call in flight and install the trajectory it returns. - The call goes to the worker so the harness keeps playing while the model runs. The spin after - dispatch keeps a wrapper answering without inference — ``None``, a local decision — in the round it - was asked: it ends when the call is already done, or when the gate reports it reached the model. + The call goes to the worker so the harness keeps playing while the model runs; a wrapper that + answers without inference still resolves in the round it was asked. A stack with no + ``SchedulingWrapper`` never enters the gate, so its call blocks the round like a direct one. """ session, gate = self._policy_session, self._gate assert session is not None and gate is not None # only a live episode steps @@ -424,8 +450,9 @@ def _step(self, clock: pimm.Clock) -> None: if self._task is not None: self._deadline = clock.now() + self._task.timeout self._future = self._executor.submit(session, frozen_view(obs)) + # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. while not (self._future.done() or gate.entered): - pass + time.sleep(0) self._collect(self._future, gate, clock) def _collect(self, future: Future[list[dict[str, Any]] | None], gate: InferenceGate, clock: pimm.Clock) -> None: @@ -513,7 +540,7 @@ def _shutdown(self) -> None: self._policy_session.close() self._policy_session = None - def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: # noqa: C901 + def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: while not should_stop.value: # One action per round, mutually exclusive: handle a directive, start the next trial (or exit # when the plan is done), finish a self-driven trial that is out of budget or done, or step the diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index 125849b74..a124f6986 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -293,9 +293,6 @@ class Recorder: once per inference at the outermost tap and reused by inner taps so every tap stamps the inference identically. ``blueprint``, if given, is sent as the recording's layout instead of the auto-built one. - - The state shared across taps (the nesting depth, the timeline values) carries no lock: the harness - keeps one session call in flight, so a whole pipeline's taps run on one thread, one call at a time. """ def __init__( @@ -308,6 +305,7 @@ def __init__( self._stream: rr.RecordingStream | None = None self._rrd_path: Path | None = None self._live = 0 + # Shared across a pipeline's taps and unlocked: the harness keeps one session call in flight. self._depth = 0 self._timeline_values: dict[str, Any] = {} self._image_paths: list[str] = [] diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 967d2bf8b..6bb529ec8 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -102,7 +102,7 @@ def server_meta(self) -> dict[str, Any]: ws_session.close() return self._server_meta - def new_session(self, context=None, now=None, gate=None) -> RemoteSession: + def new_session(self, context=None, *, now=None, gate=None) -> RemoteSession: compress = bool(self.server_meta().get(keys.COMPRESS_IMAGES)) ws_session = self._client.new_session() return RemoteSession(ws_session, compress_images=compress) @@ -166,8 +166,8 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, now=None, gate=None) -> Session: - return self._policy().new_session(context, now, gate) + def new_session(self, context=None, *, now=None, gate=None) -> Session: + return self._policy().new_session(context, now=now, gate=gate) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index b2733e930..1b25a5a71 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -82,7 +82,7 @@ class ScriptedProportionalPolicy(Policy): clock, no images. Codec stamps/truncates; the harness anchors/schedules. """ - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _ScriptedSession() @@ -143,7 +143,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): while not should_stop.value: cmd_msg = self.commands.read() if self._status == RobotStatus.ERROR: - self._status = RobotStatus.AVAILABLE # the command that arrived with the error is skipped + self._status = RobotStatus.AVAILABLE elif cmd_msg is not None and cmd_msg.updated: self._apply(cmd_msg.data) if self._error_pending: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index de0bebfdb..345145f21 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -13,21 +13,14 @@ from positronic.dataset.serializers import Serializers from positronic.drivers import roboarm from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import ( - CartesianDelta, - CartesianPosition, - Reset, - TrajectoryPlayer, - from_wire, - to_wire, -) +from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model 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, Policy, SchedulingWrapper, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import Directive, DirectiveType, Harness, _assert_anchored +from positronic.policy.harness import Directive, DirectiveType, Harness, TrajectoryPlayer, _assert_anchored from positronic.policy.remote import RemoteSession from positronic.policy.wrappers import ChunkedSchedule from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -54,8 +47,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simu The sources/dests are no-ops: these tests pair the harness ports directly (never via ``wire_embodiment``), so only the spec — names, serializers, - home values, descriptor — is read by the Harness. ``simulated`` is what makes - ``inference_latency`` bite, since the knob is sim-only. + home values, descriptor — is read by the Harness. """ observations = { 'robot_state': Observation(pimm.NoOpEmitter(), Serializers.robot_state), @@ -90,7 +82,7 @@ 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, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _SpySession(self) @@ -135,7 +127,7 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -165,7 +157,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self.reset_calls += 1 self.last_reset_context = context return _ChunkSession(self) @@ -200,7 +192,7 @@ def __init__(self, command: roboarm.command.CommandType | None = None, target_gr self.command = command self.target_grip = float(target_grip) - def new_session(self, context=None, now=None, gate=None) -> RemoteSession: + def new_session(self, context=None, *, now=None, gate=None) -> RemoteSession: action = [{'robot_command': self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] return RemoteSession(_FakeInferenceSession(action)) @@ -544,7 +536,7 @@ 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, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _StaticMetaSession(self._command) # Session.meta defaults to {} @property @@ -1096,7 +1088,7 @@ def __call__(self, obs): return [] class EmptyChunkPolicy(Policy): - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _EmptyChunkSession() harness = Harness(EmptyChunkPolicy(), make_embodiment()) @@ -1125,8 +1117,7 @@ def new_session(self, context=None, now=None, gate=None): @pytest.mark.timeout(3.0) -def test_harness_clears_trajectory_on_home(world): - """Verify that HOME resets trajectory state so next RUN gets a fresh chunk.""" +def test_harness_clears_trajectory_on_abort(world): policy = ChunkPolicy() harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) @@ -1156,7 +1147,6 @@ def test_harness_clears_trajectory_on_home(world): @pytest.mark.timeout(3.0) def test_harness_clears_trajectory_on_run(world): - """Verify that RUN resets trajectory state so a fresh chunk is emitted.""" policy = ChunkPolicy() harness = Harness(policy, make_embodiment()) p = _pair_all(world, harness) @@ -1584,7 +1574,7 @@ def __init__(self, wall_sec: float = 0.0, span_sec: float = 0.2, steps: int = 10 self._span_sec = span_sec self._steps = steps - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _SlowSession(self._wall_sec, self._span_sec, self._steps) @@ -1616,7 +1606,7 @@ def wrap_session(self, inner: Session, context, now): class _TimedRecorder(pimm.SignalEmitter): - """Records each emission against the world clock, so a test can read when a command actually went out.""" + """Records each emission against the world clock, so a test can read when a command went out.""" def __init__(self, clock: pimm.Clock): self._clock = clock @@ -1712,7 +1702,7 @@ def __call__(self, obs): return [{keys.ROBOT_COMMAND: command, 'timestamp': i * 0.01} for i in range(10)] class _GripThenArmPolicy(Policy): - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _GripThenArm() harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index 89dc306ea..c743e1583 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -164,12 +164,12 @@ class _ChunkPolicy(Policy): def __init__(self, actions: list[dict]): self._actions = actions - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _FixedSession(list(self._actions)) class _SinglePolicy(Policy): - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _FixedSession({'v': 42}) @@ -183,7 +183,7 @@ def encode(self, data): class _MetaPolicy(Policy): - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _FixedSession({}) @property diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index 7448c8cb1..e6704a314 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -27,7 +27,7 @@ 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 - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self.session_count += 1 return _TrackingSession(self._actions, {'policy_key': 'policy_value'}) @@ -57,7 +57,7 @@ def __init__(self, rec, actions): self._actions = actions self.last_session = None - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self.last_session = _CapturingSession(self._rec, self._actions) return self.last_session @@ -146,7 +146,7 @@ def __call__(self, obs): return None class _NonePolicy(Policy): - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _NoneSession() session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 19f08f0ae..ee5f58f44 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -49,12 +49,14 @@ def now_ns(self) -> int: class _ConstSession(Session): - def __init__(self, actions): + def __init__(self, actions, wall_sec: float = 0.0): self._actions = actions + self._wall_sec = wall_sec self.call_count = 0 def __call__(self, obs): self.call_count += 1 + time.sleep(self._wall_sec) return self._actions @@ -63,11 +65,14 @@ def __init__(self, actions): self._actions = actions self._session: _ConstSession | None = None - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self._session = _ConstSession(self._actions) return self._session +_ONE_ACTION = [{'v': 1, 'timestamp': 0.0}] + + def _obs(now_sec=0.0): return {keys.OBS_TIME_NS: int(now_sec * 1e9)} @@ -113,7 +118,7 @@ def test_re_infers_after_trajectory_consumed(self): def test_single_action_refires_immediately_after(self): """Single action at ts=0 → trajectory_end = now → next tick re-infers.""" clock = _FakeClock(t=1.0) - policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)) session = policy.new_session(now=clock.now) session(_obs()) clock.t = 1.01 @@ -128,7 +133,7 @@ def test_wrapper_pipe_wrapper(self): clock = _FakeClock(t=1.0) pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = pipeline.wrap(_ConstPolicy(_ONE_ACTION)) session = policy.new_session(now=clock.now) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None @@ -230,7 +235,7 @@ class _CapturePolicy(Policy): def __init__(self): self.session = _CaptureSession() - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return self.session @@ -473,7 +478,7 @@ def test_inline_full_pipe(self): def test_inline_tolerates_marker_less_pipe(self): clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstPolicy(_ONE_ACTION) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) session = policy.new_session(now=clock.now) result = session(_obs()) @@ -588,18 +593,7 @@ def __init__(self): def __call__(self, obs): self.entered.set() self.release.wait(timeout=5.0) - return [{'v': 1, 'timestamp': 0.0}] - - -class _SlowSession(Session): - """Spends ``wall_sec`` of real time in the model, which is what the measured mode charges.""" - - def __init__(self, wall_sec: float): - self._wall_sec = wall_sec - - def __call__(self, obs): - time.sleep(self._wall_sec) - return [{'v': 1, 'timestamp': 0.0}] + return _ONE_ACTION def _run_gated(gate: InferenceGate, inner: Session, executor: ThreadPoolExecutor): @@ -611,7 +605,7 @@ def test_declared_mode_parks_the_call_until_its_delay_has_passed(self): clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) concurrent.futures.wait([future], timeout=0.05) assert not future.done(), 'the gate let the call through before its declared delay' clock.t = 1.5 @@ -622,7 +616,7 @@ def test_measured_mode_parks_the_call_for_its_own_wall_duration(self): clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, LatencyMode.MEASURED) with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _SlowSession(0.1), executor) + future = _run_gated(gate, _ConstSession(_ONE_ACTION, wall_sec=0.1), executor) concurrent.futures.wait([future], timeout=0.3) assert not future.done(), 'the gate charged nothing for a call that took 0.1s' clock.t = 1.5 # past t0 + the wall duration, whatever it measured @@ -637,7 +631,7 @@ def test_a_call_owed_nothing_is_never_parked(self, mode, delay_sec): clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, mode, delay_sec) with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) assert future.result(timeout=1.0) is not None # the clock never moves and the call still returns assert gate.entered is False @@ -657,7 +651,7 @@ def test_cancel_unparks_a_call_the_harness_no_longer_wants(self): clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession([{'v': 1, 'timestamp': 0.0}]), executor) + future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) concurrent.futures.wait([future], timeout=0.05) assert not future.done() gate.cancel() @@ -668,7 +662,7 @@ def test_gate_wraps_the_inner_session_of_a_scheduling_wrapper(self): has already been paid, so the anchor is the release instant.""" clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) - policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)) session = policy.new_session(now=clock.now, gate=gate) with ThreadPoolExecutor(max_workers=1) as executor: future = executor.submit(session, _obs()) @@ -683,7 +677,7 @@ def test_a_stack_without_a_scheduling_wrapper_gets_no_gate(self): through and the gate is never entered.""" clock = _FakeClock(t=1.0) gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) - policy = TemporalStack(keys=('v',), offsets_sec=(0.0,)).wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = TemporalStack(keys=('v',), offsets_sec=(0.0,)).wrap(_ConstPolicy(_ONE_ACTION)) session = policy.new_session(now=clock.now, gate=gate) assert session({**_obs(), 'v': np.array([1.0])}) is not None diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index 219c82bd5..a6e7eb006 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -2,9 +2,9 @@ Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only -``ChunkedSchedule`` needs the live clock — it anchors a chunk to inference *completion*, which the -pre-inference observation stamp cannot give — so the harness passes ``now`` (a ``Callable[[], float]`` -in seconds) to ``new_session`` and it reaches that one session. +``ChunkedSchedule`` needs the live clock — it anchors a chunk to the instant the inference gate releases +its call, which the pre-inference observation stamp cannot give — so the harness passes ``now`` (a +``Callable[[], float]`` in seconds) to ``new_session``, and it reaches every session in the stack. """ from collections import deque @@ -24,9 +24,9 @@ class ChunkedSchedule(SchedulingWrapper): """Wait for the current trajectory to finish before calling the inner policy again. Owns relative→absolute time conversion: inner layers (codecs, models) emit relative timestamps; - this wrapper anchors them to ``now()`` *after* inner inference returns, so execution aligns to - inference-finish (not inference-start). Returns ``None`` ("keep executing the current trajectory") - until the last action's timestamp is reached, then calls the inner policy. + this wrapper anchors them to ``now()`` *after* the inner call returns, which is the instant the + inference gate releases it. Returns ``None`` ("keep executing the current trajectory") until the last + action's timestamp is reached, then calls the inner policy. """ class _Session(DelegatingSession): @@ -53,7 +53,6 @@ def __call__(self, obs): # immediate action executes instead of raising. if isinstance(result, dict): result = [result] - # Anchor to post-inference time so execution starts when inference *finished*. # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). now = self._now() result = [{**r, keys.ACTION_TIMESTAMP: now + r.get(keys.ACTION_TIMESTAMP, 0.0)} for r in result] diff --git a/positronic/probe.py b/positronic/probe.py index 76c6d9cb9..12cc70642 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -134,7 +134,7 @@ def main( image_keys = [k for k in obs if k.startswith(keys.IMAGE_PREFIX)] rec = Recorder(pos3.sync(output_dir)) - session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None, time.time) + session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None, now=time.time) meta = dict(session.meta) name = label or _recording_name(meta) try: diff --git a/positronic/simulator/env_server/adapter.py b/positronic/simulator/env_server/adapter.py index 6fc09c0db..14d096e73 100644 --- a/positronic/simulator/env_server/adapter.py +++ b/positronic/simulator/env_server/adapter.py @@ -109,9 +109,6 @@ def __init__(self, env_control_frame: geom.Transform3D | None = None): def _reset_command_state(self) -> None: self._held: dict[str, Any] = {} # last command per channel — re-sent until the next one arrives - # Last commanded gripper closure, held across an episode boundary: grip is an absolute [0, 1] value - # with no 'hold' command to fall back on (unlike the arm), so a fresh trial must freeze it, not reopen. - self._grip = 0.0 @final def reset_token(self, context: dict[str, Any]) -> Any: @@ -137,6 +134,5 @@ def action(self, commands: dict[str, pimm.Message]) -> dict[str, Any]: cmd = None case roboarm_command.CartesianDelta() | roboarm_command.JointDelta(): self._held.pop(keys.ROBOT_COMMAND) - if keys.TARGET_GRIP in self._held: - self._grip = float(self._held[keys.TARGET_GRIP]) - return {'command': _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), 'grip': self._grip} + grip = float(self._held.get(keys.TARGET_GRIP, 0.0)) + return {'command': _wire_command(_in_env_control_frame(cmd, self.env_control_frame)), 'grip': grip} diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index e801d1b18..55ed9d624 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -128,7 +128,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self._cleanup.close() def _step_env(self) -> dict[str, Any]: - commands = {name: receiver.read() for name, receiver in self.commands.items()} + reads = ((name, receiver.read()) for name, receiver in self.commands.items()) + commands = {name: msg for name, msg in reads if msg is not None} result = self._conn.step(self._adapter.action(commands)) payload = self._adapter.terminal(result) if payload: # truthy-valued done: a non-empty payload ends the trial, an empty/``None`` one continues diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index b8d30edbf..3633e180a 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -312,7 +312,7 @@ def __init__(self, command: roboarm_command.CommandType, chunk_len: int): self.chunk_len = chunk_len self.chunks = 0 - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): return _JointposChunkSession(self) diff --git a/positronic/simulator/mujoco/sim.py b/positronic/simulator/mujoco/sim.py index 2fd3996f8..b92dd02ed 100644 --- a/positronic/simulator/mujoco/sim.py +++ b/positronic/simulator/mujoco/sim.py @@ -178,7 +178,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p now = clock.now() cmd_msg = self.commands.read() if self._error: - self._error = False # the command that arrived with the error is skipped, not queued + self._error = False elif cmd_msg is not None and cmd_msg.updated: self._apply_command(cmd_msg.data) grip_msg = self.target_grip.read() diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index d4aac37c1..45ebd9bf5 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -286,7 +286,7 @@ class DreamZeroPolicy(Policy): def __init__(self, sp: DreamZeroSubprocess): self._subprocess = sp - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): client = RoboarenaClient(port=self._subprocess.roboarena_port) client.connect() return _DreamZeroSession(client, str(uuid.uuid4())) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 0258e735f..0ae1f7458 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -220,7 +220,7 @@ def __init__(self, groot: Gr00tSubprocess, checkpoint_path: str): self._groot = groot self._checkpoint_path = checkpoint_path - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self._groot.client.reset() return _Gr00tSession(self._groot.client) diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index 62746c864..53c21a7ef 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -102,7 +102,7 @@ def config(self) -> PreTrainedConfig: """The checkpoint's own declaration of what this policy takes.""" return self._policy.config - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self._policy.reset() return _LerobotSession(self._policy, self._preprocessor, self._postprocessor, self._device, self._meta) diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index c20e03a33..8c41eb11b 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -92,7 +92,7 @@ def __init__(self, policy: PreTrainedPolicy, device: str | None = None, extra_me self._policy = policy.to(self._device) self._meta = extra_meta or {} - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): self._policy.reset() return _LerobotSession(self._policy, self._device, self._meta) diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index b18853c63..76516a148 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -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, now=None, gate=None) -> Session: + def new_session(self, context=None, *, now=None, gate=None) -> Session: 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 ff9e0c174..57f20f047 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -140,7 +140,7 @@ class OpenpiPolicy(Policy): def __init__(self, subproc: OpenpiSubprocess): self._subproc = subproc - def new_session(self, context=None, now=None, gate=None): + def new_session(self, context=None, *, now=None, gate=None): client = self._subproc.client client.reset() return _OpenpiSession(client) From c42ef5d64980a4fe84c7892c3398eae21e946a54 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 22:09:12 +0300 Subject: [PATCH 03/64] Charge inference from the trial context instead of an `InferenceGate` `ChunkedSchedule` stamps its chunk at the observation instant plus the trial's `inference_latency` charge (a fixed value, or the call's wall duration); the harness withholds the trajectory, and the world clock, until that instant. Both sides read the same context key and the same observation timestamp, so nothing crosses between the wrapper and the harness: `InferenceGate`, `LatencyMode`, the `SchedulingWrapper` marker and the `now`/`gate` parameters of `new_session` are deleted. Release instants compare in integer ns, so a fixed charge lands on the exact instant instead of one round late. --- .basedpyright/baseline.json | 10 - ARCHITECTURE.md | 11 +- positronic/keys.py | 4 +- .../offboard/tests/test_remote_policy.py | 13 +- positronic/offboard/tests/test_server.py | 17 +- positronic/policy/base.py | 118 +-------- positronic/policy/codec.py | 2 +- positronic/policy/harness.py | 144 ++++++----- positronic/policy/recording.py | 2 +- positronic/policy/remote.py | 6 +- .../policy/tests/golden_pipeline.json.gz | Bin 7893 -> 7893 bytes .../policy/tests/test_golden_pipeline.py | 2 +- positronic/policy/tests/test_harness.py | 40 +-- positronic/policy/tests/test_policy_io.py | 6 +- positronic/policy/tests/test_recording.py | 6 +- positronic/policy/tests/test_wrappers.py | 235 +++++------------- positronic/policy/wrappers.py | 45 ++-- positronic/probe.py | 2 +- .../env_server/tests/test_remote_env.py | 2 +- positronic/vendors/dreamzero/server.py | 2 +- positronic/vendors/gr00t/server.py | 2 +- positronic/vendors/lerobot/policy.py | 2 +- positronic/vendors/lerobot_0_3_3/policy.py | 2 +- positronic/vendors/molmoact2/policy.py | 2 +- positronic/vendors/openpi/server.py | 2 +- 25 files changed, 233 insertions(+), 444 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 4094a569c..100110818 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -6210,16 +6210,6 @@ } } ], - "./positronic/policy/wrappers.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 20, - "lineCount": 1 - } - } - ], "./positronic/probe.py": [ { "code": "reportAttributeAccessIssue", diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 9e3ae3180..16633701c 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -132,11 +132,12 @@ episode lifecycle — nothing else. Scheduling, blending, history stacking and e the wrapper stack around the policy; a session returning `None` means "keep executing the current trajectory". -**The platform charges inference, not the wrapper.** A scheduling wrapper may be someone else's -submission, so what a model call costs cannot be left to the wrapper's own bookkeeping. The harness -runs the call on a worker thread and installs an `InferenceGate` around the wrapper's path to the -model: the gate holds the answer until the cost is paid, so the wrapper resumes — and anchors its -chunk — at the release instant, and cannot see a result early however it is written. +**Inference cost is a fact of the trial, charged twice from the same source.** The trial context +carries `inference_latency` — a fixed charge in seconds, or the call's own wall duration — and both +sides read it independently: the scheduling wrapper stamps its chunk for the observation instant +plus the charge, and the harness withholds the trajectory, and the world clock, until that instant +arrives. Neither side signals the other; the observation's own timestamp is the shared anchor, so +the charge is deterministic wherever a fixed one is asked for. **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every diff --git a/positronic/keys.py b/positronic/keys.py index 9fddef819..66305bd34 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -89,5 +89,7 @@ def is_robot_command(name: str) -> bool: EVAL_SUCCESS = 'eval.success' EVAL_TERMINATED = 'eval.terminated' -# The sim-only inference cost a trial's RUN context asks the platform to charge the model call. +# The inference cost each model call charges the world clock: seconds (fixed, reproducible), or ``True`` +# for the call's own wall duration. A sim trial without the key charges nothing (the world holds still per +# call); hardware always pays wall. INFERENCE_LATENCY = 'inference_latency' diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 8f0bec5d5..388c31fc1 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -390,10 +390,9 @@ def test_empty_declaration_fails_before_motion(): def test_declared_stack_built_at_session_open(): """The server-declared local stack runs in front of the connection.""" - clock = [1.0] policy, mock_ws = _mock_remote_policy(CHUNKED_STACK, infer_return=[{'a': 1, 'timestamp': 0.0}]) - session = policy.new_session(now=lambda: clock[0]) - actions = session({keys.OBS_TIME_NS: 0}) + session = policy.new_session({keys.INFERENCE_LATENCY: 0.0}) + actions = session({keys.OBS_TIME_NS: int(1e9)}) assert actions == [{'a': 1, 'timestamp': 1.0}] @@ -512,13 +511,13 @@ def test_remote_policy_lifecycle(inference_server, mock_policy): host, port = inference_server policy = RemotePolicy(f'{host}:{port}') - session = policy.new_session(now=lambda: 0.0) + session = policy.new_session({keys.INFERENCE_LATENCY: 0.0}) meta = session.meta assert meta['server.model_name'] == 'test_model' assert meta['type'] == 'remote' - obs = {'dataset': 'test'} + obs = {'dataset': 'test', 'obs_time_ns': 0} action = session(obs) # Single-dict server response is normalized to a 1-element list (Session contract) and # anchored to absolute time by the declared ChunkedSchedule. @@ -527,7 +526,7 @@ def test_remote_policy_lifecycle(inference_server, mock_policy): session.close() # New session - session2 = policy.new_session(now=lambda: 0.0) + session2 = policy.new_session() session2.close() @@ -535,7 +534,7 @@ def test_remote_session_meta(inference_server): """Session meta must include server metadata.""" host, port = inference_server policy = RemotePolicy(f'{host}:{port}') - session = policy.new_session(now=lambda: 0.0) + session = policy.new_session() meta = session.meta assert meta['type'] == 'remote' diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 647619fdd..8710bbf48 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -244,7 +244,7 @@ def __call__(self, obs): class _ScriptedPolicy(Policy): """Deterministic base policy: every session returns the same untimestamped chunk.""" - def new_session(self, context=None, *, now=None, gate=None) -> Session: + def new_session(self, context=None) -> Session: return _ScriptedSession() @@ -254,15 +254,15 @@ def test_in_process_equals_remote_for_same_pipeline(start_server): def pipeline(): return ChunkedSchedule() | remote | ActionTimestamp(fps=10.0) | PolicySource(_ScriptedPolicy()) - clock = [100.0] + context = {keys.INFERENCE_LATENCY: 0.0} host, port, _server = start_server(pipeline()) - remote_session = RemotePolicy(f'{host}:{port}').new_session(now=lambda: clock[0]) + remote_session = RemotePolicy(f'{host}:{port}').new_session(context) - local_session = inline(pipeline()).new_session(now=lambda: clock[0]) + local_session = inline(pipeline()).new_session(context) - remote_actions = remote_session({keys.OBS_TIME_NS: 0}) - local_actions = local_session({keys.OBS_TIME_NS: 0}) + remote_actions = remote_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) + local_actions = local_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) assert remote_actions == local_actions # Three scripted actions plus the chunk-closing validity sentinel ActionTimestamp appends. assert local_actions == [ @@ -273,9 +273,8 @@ def pipeline(): ] # Both gate identically while the chunk plays out. - clock[0] = 100.15 - assert remote_session({keys.OBS_TIME_NS: 0}) is None - assert local_session({keys.OBS_TIME_NS: 0}) is None + assert remote_session({keys.OBS_TIME_NS: int(100.15 * 1e9)}) is None + assert local_session({keys.OBS_TIME_NS: int(100.15 * 1e9)}) is None remote_session.close() diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 6d5b3dbe4..3cd6b435f 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,13 +1,9 @@ from __future__ import annotations -import time from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping -from enum import Enum +from collections.abc import Mapping from typing import Any -Now = Callable[[], float] - # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. SEQ = 'seq' PAR = 'par' @@ -73,85 +69,6 @@ def close(self): self._inner.close() -class LatencyMode(Enum): - """When the platform lets a scheduling wrapper have the model's answer.""" - - # At completion, and the world runs on meanwhile: the cost is whatever the model really took. - LIVE = 'live' - # A fixed delay after the call started, whatever the model really took. A delay of zero holds the world - # still for the whole call. - DECLARED = 'declared' - # The call's own wall duration after it started, charged on the world clock. - MEASURED = 'measured' - - -class InferenceGate: - """The platform's hold on a scheduling wrapper's path to the model. - - Installed around the wrapper's inner session, so no wrapper can reach a result before the mode's cost - has been paid. The wrapper resumes at the release instant and anchors there. - """ - - def __init__(self, now: Now, mode: LatencyMode, delay_sec: float = 0.0): - self._now = now - self._mode = mode - self._delay_sec = delay_sec - self._wall_t0 = 0.0 - self._cancelled = False - # True while a call is inside the model — a wrapper that answered on its own never sets it. - # ``t0`` is the world instant that call started, valid once ``entered``. - self.t0 = 0.0 - self.entered = False - - def wrap(self, inner: Session) -> Session: - return InferenceGate._Session(inner, self) - - def cancel(self) -> None: - """Release a parked call, whose result is on its way to a harness that no longer wants it.""" - self._cancelled = True - - def hold(self) -> float | None: - """Wall seconds the world must not advance for, or ``None`` to hold until the call completes.""" - match self._mode: - case LatencyMode.LIVE: - return 0.0 - case LatencyMode.DECLARED: - return None if self._now() >= self._release_at() else 0.0 - case LatencyMode.MEASURED: - # The world may run no further ahead of the call's start than wall time has. - return max(0.0, self._now() - self._release_at()) - - def _release_at(self) -> float: - match self._mode: - case LatencyMode.LIVE: - return self.t0 - case LatencyMode.DECLARED: - return self.t0 + self._delay_sec - case LatencyMode.MEASURED: - return self.t0 + (time.monotonic() - self._wall_t0) - - class _Session(DelegatingSession): - """Charges the inner call.""" - - def __init__(self, inner: Session, gate: InferenceGate): - super().__init__(inner) - self._gate = gate - - def __call__(self, obs): - gate = self._gate - gate.t0 = gate._now() - gate._wall_t0 = time.monotonic() - gate.entered = True - result = self._inner(obs) - release = gate._release_at() - # The world clock is advanced by the harness's thread, so the park has to poll it; sleeping - # zero hands over the GIL without adding a wake-up granularity to the release instant. - while not gate._cancelled and gate._now() < release: - time.sleep(0) - gate.entered = False - return result - - class Policy(ABC): """Factory for inference sessions. @@ -161,18 +78,11 @@ class Policy(ABC): """ @abstractmethod - def new_session( - self, context: dict[str, Any] | None = None, *, now: Now | None = None, gate: InferenceGate | None = None - ) -> Session: + def new_session(self, context: dict[str, Any] | None = None) -> Session: """Create a new inference session for an episode. Args: context: Episode context (task description, eval metadata, etc.). - now: The runtime clock (current time in seconds), supplied by the harness and passed down - to every wrapped session. ``None`` where no runtime clock exists (server-side, warmup). - gate: The platform's hold on the path to the model, supplied by the harness and installed - around the inner session of every ``SchedulingWrapper`` in the stack. ``None`` where no - runtime imposes inference cost (server-side, warmup). """ @property @@ -190,8 +100,8 @@ class DelegatingPolicy(Policy): def __init__(self, inner: Policy): self._inner = inner - def new_session(self, context=None, *, now=None, gate=None): - return self._inner.new_session(context, now=now, gate=gate) + def new_session(self, context=None): + return self._inner.new_session(context) @property def meta(self): @@ -223,12 +133,12 @@ class PolicyWrapper: def wrap(self, policy: Policy) -> Policy: """Apply this wrapper to a policy. Default: wrap every session it creates via ``wrap_session``. - Composition happens at config time; the runtime clock reaches the wrapped + Composition happens at config time; the trial context reaches the wrapped sessions through ``new_session``. """ return _WrapperPolicy(policy, self) - def wrap_session(self, inner: Session, context: dict[str, Any] | None, now: Now | None) -> Session: + def wrap_session(self, inner: Session, context: dict[str, Any] | None) -> Session: """Wrap a single session. Subclasses override this for per-session wrapping.""" raise NotImplementedError('Override wrap_session or wrap') @@ -256,15 +166,6 @@ def _wrappers(self) -> tuple: return (self,) -class SchedulingWrapper(PolicyWrapper): - """A wrapper that owns the plan: it decides when to call the model and returns the trajectory the - harness plays, rather than one action for the moment. - - Being one is what earns the wrapper an ``InferenceGate`` around its inner session, so the inference - cost is imposed below it instead of trusted to it. - """ - - class _WrapperPolicy(DelegatingPolicy): """Generic policy wrapper produced by ``PolicyWrapper.wrap()``. @@ -275,11 +176,8 @@ def __init__(self, inner: Policy, wrapper: PolicyWrapper): super().__init__(inner) self._wrapper = wrapper - def new_session(self, context=None, *, now=None, gate=None): - inner = self._inner.new_session(context, now=now, gate=gate) - if gate is not None and isinstance(self._wrapper, SchedulingWrapper): - inner = gate.wrap(inner) - return self._wrapper.wrap_session(inner, context, now) + def new_session(self, context=None): + return self._wrapper.wrap_session(self._inner.new_session(context), context) @property def meta(self): diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 16ae3a000..38ae68021 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -80,7 +80,7 @@ def training_encoder(self) -> EpisodeTransform: def meta(self) -> dict: return {} - def wrap_session(self, inner: Session, context, now): + def wrap_session(self, inner: Session, context): return _CodecSession(inner, self) @final diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 849065acb..867194c8c 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -14,7 +14,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 InferenceGate, LatencyMode, Policy, Session +from positronic.policy.base import Policy, Session from positronic.utils import flatten_dict, frozen_view # How far from now an action may be scheduled. A chunk spans seconds, so this is loose enough that no real @@ -26,6 +26,11 @@ # noticed, and with it the granularity every command timestamp is quantized to. POLL_PERIOD_SEC = 0.01 +# How long a submitted session call may take to answer and still resolve within its round. A wrapper that +# skips inference answers in microseconds; a real model call runs far past this and is then paced by +# ``_take`` across rounds. +SKIP_REPLY_SEC = 0.001 + # One channel's schedule: waypoints stamped with absolute clock ns, ascending. Trajectory: TypeAlias = list[tuple[int, Any]] @@ -180,9 +185,9 @@ class Harness(pimm.ControlSystem): Handles directives (RUN/FINISH/ABORT) and dataset recording. Inference intelligence — scheduling, error recovery, blending, absolute time stamping — lives in the policy/session layer: the wrapper owns the plan, the harness plays it, one command per channel per round. The session call runs on a worker - thread so playing continues while the model does, and the ``InferenceGate`` installed below the - scheduling wrapper charges that call the trial's inference latency. The RUN context is handed whole to - the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). + thread so playing continues while the model does; the harness withholds the trajectory, and the world + clock, until the trial's inference charge (``inference_latency``) is paid. The RUN context is handed + whole to the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's @@ -192,8 +197,7 @@ class Harness(pimm.ControlSystem): timeout records False. A task-less session has neither deadline nor budget and ends only on directives. The ``Embodiment`` supplies the observation serializers (which own the canonical key names), the command - channels and the home action. The policy owns its wrapper stack; the harness runs what it is given, - passing ``new_session`` the clock the scheduling wrapper anchors chunks to. + channels and the home action. The policy owns its wrapper stack; the harness runs what it is given. """ def __init__( @@ -216,11 +220,16 @@ def __init__( self._policy_session: Session | None = None # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. self._running = False - # One session call at a time, on a worker so the harness keeps playing while the model runs, and the - # gate that charges it the trial's inference latency. The gate lives for one episode. + # One session call at a time, on a worker so the harness keeps playing while the model runs. self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') self._future: Future[list[dict[str, Any]] | None] | None = None - self._gate: InferenceGate | None = None + # The in-flight call's start: the world instant its observation was built, and the wall instant it + # was submitted. + self._t0_ns = 0 + self._wall_t0 = 0.0 + # Seconds each model call costs the world clock this episode, or ``None`` to charge the call's own + # wall duration (hardware pace, and the sim's ``inference_latency=True``). + self._charge: float | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None # Whether this episode's first observation has landed. Until it does the deadline stands where the @@ -293,8 +302,6 @@ def _cancel_session(self) -> None: """ for player in self._players.values(): player.set([]) - if self._gate is not None: - self._gate.cancel() if self._future is not None: future, self._future = self._future, None future.result() # nothing may close or re-enter the session while the worker is still inside it @@ -317,20 +324,6 @@ def _finalize_recording( # period) to the closing episode — the cooperative scheduler cannot give the recorder a turn alone. self._telemetry.end(virtual_now) - def _new_gate(self, clock: pimm.Clock) -> InferenceGate: - """The inference cost this trial imposes on the scheduling wrapper's call. - - Hardware pays whatever the model takes. A sim trial pays what its ``inference_latency`` asks for: the - call's own wall duration (``True``), or a fixed delay — the reproducible mode, and by default zero, - which holds the world still for the whole call as sim-only harnesses do. - """ - if not self._embodiment.simulated: - return InferenceGate(clock.now, LatencyMode.LIVE) - latency = self.context.get(keys.INFERENCE_LATENCY, False) - if latency is True: # bool is an int subclass — check identity first - return InferenceGate(clock.now, LatencyMode.MEASURED) - return InferenceGate(clock.now, LatencyMode.DECLARED, float(latency)) - def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. @@ -340,7 +333,13 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: armed here and moved to that first observation once it lands, so an episode that never gets one is still bounded. """ - self.context = context + self.context = dict(context) + if self._embodiment.simulated: + # A sim trial that doesn't ask for latency simulation runs free of it: the world holds still for + # every model call. Hardware (no key) pays what the call really takes. + self.context.setdefault(keys.INFERENCE_LATENCY, False) + latency = self.context.get(keys.INFERENCE_LATENCY, True) + self._charge = None if latency is True else float(latency) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -352,8 +351,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._gate = self._new_gate(clock) - self._policy_session = self.policy.new_session(self.context, now=clock.now, gate=self._gate) + self._policy_session = self.policy.new_session(self.context) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -377,7 +375,6 @@ def _end_episode( if self._policy_session: self._policy_session.close() self._policy_session = None - self._gate = None self._home() self._running = False @@ -432,51 +429,66 @@ def _step(self, clock: pimm.Clock) -> None: """Keep one session call in flight and install the trajectory it returns. The call goes to the worker so the harness keeps playing while the model runs; a wrapper that - answers without inference still resolves in the round it was asked. A stack with no - ``SchedulingWrapper`` never enters the gate, so its call blocks the round like a direct one. + answers without inference still resolves in the round it was asked. """ - session, gate = self._policy_session, self._gate - assert session is not None and gate is not None # only a live episode steps - if self._future is None: - obs = self._build_obs(clock) - if obs is None: - return - if not self._rollout_started: - # The rollout begins at its first observation, not when the reset returned: a reset only asks - # the producer for a scene, and the turns spent delivering it are neither the trial's budget - # nor its duration. - self._rollout_started = True - self._telemetry.start_rollout(clock.now()) - if self._task is not None: - self._deadline = clock.now() + self._task.timeout - self._future = self._executor.submit(session, frozen_view(obs)) - # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. - while not (self._future.done() or gate.entered): - time.sleep(0) - self._collect(self._future, gate, clock) - - def _collect(self, future: Future[list[dict[str, Any]] | None], gate: InferenceGate, clock: pimm.Clock) -> None: - """Take the call's trajectory once the latency mode lets the harness have it. - - A mode that holds the world does it by not returning: the loop thread is what advances the sim clock, - so blocking here freezes it for exactly as long as the call is owed. + session = self._policy_session + assert session is not None # only a live episode steps + if self._future is not None and not self._take(self._future, clock): + return + obs = self._build_obs(clock) + if obs is None: + return + if not self._rollout_started: + # The rollout begins at its first observation, not when the reset returned: a reset only asks + # the producer for a scene, and the turns spent delivering it are neither the trial's budget + # nor its duration. + self._rollout_started = True + self._telemetry.start_rollout(clock.now()) + if self._task is not None: + self._deadline = clock.now() + self._task.timeout + self._t0_ns = clock.now_ns() + self._wall_t0 = time.monotonic() + self._future = self._executor.submit(session, frozen_view(obs)) + # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. + while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: + time.sleep(0) + self._take(self._future, clock) + + def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> bool: + """Install the call's trajectory once the world has paid for it; True once the future is consumed. + + A skip (``None``) costs nothing and is consumed on sight. A model call's trajectory is stamped for + ``t0`` plus the trial's charge, so it is withheld until the world clock reaches that instant — and + the world is withheld from running past it: blocking here blocks the loop thread, which is what + advances a virtual clock. """ - if not future.done(): - hold = gate.hold() - if hold is not None and hold <= 0.0: - return - concurrent.futures.wait([future], timeout=hold) + if future.done() and future.result() is None: + self._future = None + return True + if self._charge is not None: + # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP + # and slips the install a full round. + if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): + return False + concurrent.futures.wait([future]) + elif not future.done(): + # The world may run no further ahead of the call's start than wall time has. + ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) + if ahead <= 0.0: + return False + concurrent.futures.wait([future], timeout=ahead) if not future.done(): - return + return False self._future = None actions = future.result() # taken on the loop thread, so a failing call still seals the episode if actions is not None: self._install(actions, clock) + return True def _install(self, actions: 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 wrapper at the instant the gate released it. + by the scheduling wrapper for the instant its charge is paid. """ _assert_anchored(actions, clock.now()) self._telemetry.step() @@ -526,14 +538,12 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self._shutdown() def _shutdown(self) -> None: - """Release the worker and the session. A call still in flight is unparked and its result dropped: the - run is over and nothing is left to install it. + """Release the worker and the session. A call still in flight runs to completion and its result is + dropped: the run is over and nothing is left to install it. The harness does not own the policy's lifetime: the caller may run several harnesses over one policy (a multi-eval sweep), so it closes the policy once, after the last run. """ - if self._gate is not None: - self._gate.cancel() self._future = None self._executor.shutdown(wait=True, cancel_futures=True) if self._policy_session is not None: diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index a124f6986..d9ed32bf8 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -342,7 +342,7 @@ def __init__(self, rec: Recorder, name: str): self._rec = rec self._name = name - def wrap_session(self, inner: Session, context, now) -> Session: + def wrap_session(self, inner: Session, context) -> Session: stream = self._rec._open_stream() return _RecordingTapSession(inner, self._rec, self._name, stream) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 6bb529ec8..438e66af6 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -102,7 +102,7 @@ def server_meta(self) -> dict[str, Any]: ws_session.close() return self._server_meta - def new_session(self, context=None, *, now=None, gate=None) -> RemoteSession: + def new_session(self, context=None) -> RemoteSession: compress = bool(self.server_meta().get(keys.COMPRESS_IMAGES)) ws_session = self._client.new_session() return RemoteSession(ws_session, compress_images=compress) @@ -166,8 +166,8 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None, *, now=None, gate=None) -> Session: - return self._policy().new_session(context, now=now, gate=gate) + def new_session(self, context=None) -> Session: + return self._policy().new_session(context) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/tests/golden_pipeline.json.gz b/positronic/policy/tests/golden_pipeline.json.gz index 34ad77dbe42719585b389f4c00379d741b280421..0df3117bead5ec06c2b34a4a9dd0a606eefa8808 100644 GIT binary patch literal 7893 zcma)>2{=^iAIE8GQjV=fBQi>h5}gT=V%!!})D0nJ8)TO?OBmZ^%Q6*;L^LV;Qe++b zmPljGzGP?YV;P43jHTPX_kX+hIq&oMJ@0wOTBI|%gsqD@h(l5*$43ha zmKNVltSr~(rWB-2FOWN1KYK-TNb#@uFNi3;dP4C}H+&<-Pvsh}t#ptr7didv$)bOv z;aw^ITB&w*4!umh0wKG0-C7~L4&5#xyH4GSq$xcgHCU=8TBoQO>$!t@ALkh}!bK;c`GeUPWJSaaC>W%7JGZM zv8}yupl6Pf5qX)1IHgS%Vtepjxc$;62Dcu&ybv3(`-tv(#*a56j+g-wt$Q4&^qmEx z9{fh|vZG@7K-wg+7{v@1a!y|l?h*2}#^_~vr5zWBdhp(xfw-_7$M*|+|CBpQRx!7@ zeA3x}5>Z^i1n1uj@K(nI)D~bGR@r@CI1Xvq1z3+>VsrBVVQ<{(U3e3!9sR}I2{izCU%-hlwIONY&2`fI*jXuFyr5i2Gxl=FtG^d>2ZFs|s zQ_zs&z?9t{3cd~=vl2^2r6J>~isLv(*FDSlUigvK*hLf(8NZ@97q=waBYvVlk#c=| z2@g9~ysAKPm4yLhDwJWsG6GDXKw%*R76Pz?R|;hqvD*L`XjE9ph}{nGfk-BjjA$O? zHt1fM8rssG5ensk8KG0%8Mi|Zz#Pz%?u<;(gD@X7=)t%HdIT1VCU@6DWTv1eV2S8L zkKdS~!mvCv1NJundKy-ZmceG|F}WQyDk|KD-3OpSpQ2XxhGRE(m$83CTVfemp=V)3 z=u|8t8&nFmfS$xMvO|A^G3kI089AVrVK5!bhh)7%OMloTR2~NEe0}J?8+rpKqr>j$ zz6W{}rl6zX>An|w3x?Kt?CH)4Rfk#XeDw^_leq8iJqgu;`RIUNj9gHCSg4Mr7vp}Y zF)UUm)r*lEdJmSRGwH?11GR!x>VV!gda3sQbCXbeShvn&Z+Bj(6KqoFtG7EG>I!4f zW%qICgJNOqx(Yt-2ch0DKHbMYNqPcK{`yFs>dx>>T(!^e(S<&oFb?n#`ck2f8wVr* z7TK--On=R&42-5AR4cgurPz?s1s+WwJd4%r!7s9 z|8U^5Qeu$Ai(RN&y+-mpm>oa&TZudKnP^z_Z{7SD-U#uMWX=@_$p0HMt&k8}C1K{{LutZG@vn*WrX=w($D-ql%PzOW%?sT7jQf9UJHo>vIW> z-v1w%qeiJigzqfER#~MzyS(Lj{xxVuq*+#DBdxik*FT zJ@CQqyaOV~DjGf^?5(dw@_aZ}G5ZPOYJDw=M?|1vPZh#@%io5ng}WN9r^-ZIv|mhO z`*Y`|m5#3b)p~7Lhk3GiS8$3J(8Al0G+Nx&^*p7Dyks1wxv#)cg*0#H{e3=jLgZ|a z>&@xU2!{af5nJPkCrc7Zc>yj5zf`tY*R*SY#wZ`w)Qh!cfgj3<7#yf-xJEsF6MnhU zxi)gVQ|B-)Ct15WYI^?xw$FXnT+~0E+QV7)JuG)zn^5IlEh26sGT&_7=4+rQWpjQ0 zteRiSK$5=sfXiYwPPj4ucx+3)#j&Ro?HLn?rVH#$`Befd_kxksjQV%pE>HE+-RGV1@}Xmhx7G+4k*~i%3;tmVX4qHmdPB93I%(l1J(}-0~sH$$SrcT>vkEgx5 z=Qtd6Zl5T8@QL!f8|3OQW0~PesCS8D?*diRYa`jz{0U~iaf{i|R>4(l9syz&|Uc6Pwv_ z(~DA{&dom+H9e*Ml(rV-PvN6VfS@hw#NB2W!)7H^*at&9gMyARidu~B_+>UT1*bc? zLEg1FBUNi#Sc4JW(+sQ0VcCC>8u^KQdd|_Bws;moi(2NN=R8`q)D9JScV%bH);ZA@ z$na?!y@^w|A7E+#apCUfORrU`o`Oee!A@f=LxM!OBaE98!7U#^0B{dS|p`kf#U?~Hg6D(;6 zgM!L{h9!encBeR6ZhdaA#z3-1r*=gswx+7pCC86*_%GSAE{=q~h|@M0$?75)X^G9< ze9)IH^?O=Q^M_oo6C0<=wB^-^%ROU<*~kR(T)10?9W}&|K9#}8?asHpWyH6nIVg~` z<;{#S^E0Q{jN8fkMZP|Vzp>3!fV{)%LP4!gG z5+#OPf_8yL77OI_zuwar?0ktIzL0Ze{MHe5J+_gCmZ% zvPEUmmipQX=~Hv}mwTjRhvbG6lT@4p_5B}+v$)3Th0;!0|54oSawWouG6@xz<5MrZ zoLG|V5kk8UQ0od6ZD&p6eoADOL}KlOfP!aD)L6~hG1?923X{ShMH>t!l)QV(J#9z# zi#Do`RMLJkF5GtH`wu?GDGml)IY+QAnCwFyGm+S3aSB5*9yVM>mJ6e$RFV^}*a=6Ma zbH&wO!N|6X)I~m7@t-xkd!GM0S6FSOvuu;d`^YDQyfl(dYfJ0mJZGLZ$eC-NHo7YoRb!sUk=_f7Gi>bAdp4eig$h5l1%;oUuZ|c| zI`ppb^-F%=i^rYU7(3viux?LD$lFUR zPT_np4a#i@$pTJ)0_+Xa4tP)F_CF{*)+4xxZM2(kHCRb4*3(m0E!NA^S}oSw(?c!R zC+2|ubAiig3VWIL(-ihGd!{LHF~_7OvW+YfuFeX~IGo!_RUF1Cxi(Pn-{6fa0`n*W za_5TT6z-{O!!mvduCgky)PTWHa))#YP~T#8Sbd7P`@3RSp$xFCXg;?h2IyAU9yHpm zh!M&JJB*HUE7}4DV4~=5w;~9XJrw@Xj&Fab+ZOCj;3Sw;$N<3tfGGH_&<%oR1CM>Nq*b807%j3#fwVqNr_TK|Z~+EjR8cEpYax3& zy4$0O84AKm(IlCXpg3P$CNnUrh+#XH2e=P@D{|Y8S4wV$LVUYj>f~i!7ohBog5CWxz?AVI{JK{=d6Czlu{5_o2VB-2~ zf5Xp21@tB6zjO8vvRxk8}j!KUFs$cUFNCr z!E;-s^9Q<=5+6%Q?n24;4qfHB?u!?+NT;lqoHaN2Aq?>tg(R5v91o~xp}nWBei|&7e{*;P0UlD>2#eX>otb9IZb=;zBbvOl?w?qlIsJd0l zHrDBPTb!_p>`!<%j}gm#Jfqb!U6xHEcI700#kfuy=V?B#q&UE0NDF?SnD>s_o(Z5v zdgPC5+aP&Wy*g$}k4N^>(rQ9)Z7X>?I}1+QhNq+X+?5z*IeJ3AufrQ%AtV#Ow+!y< z#vJDs^i&d{CztTA*qoYq%vC2<;Wu+_qKlSm<(1_yFB{}t&{H>ur&_f7FAl^7%hA&; zO4S7JJf?cOLj}*NU863)(Y~uzTtPLDqw&XAd^_UQjY)?4^OAewLhjHLc{p8bTrL5B z1I{Vf*4D8*Yz@|rhbYk#dGU&r$8~c}ktW0U6{w|`m;hs;7^M?yX{74}dD@|8Ql-gA zP*%zjaPq>9rreotkV~Q`%K3~Y&GQ5Mms(IC9@iy_%3F{rmg^NyP+4OtF8149&ileU zk*7e&z4n2Wk^OG7S^S&12=JV0V*(ov#ox#a&(E)Kl-;Zu9ir8NTv5&$G_1%mxvCPj zIq_OSLgF`xw%NL4H*`#-llzrACfLXRnBi=e1%52#$*uP4DUGdmW498<#Lv84@30Pj zdzjcg=htqU!C5(uGnP%zKShv|%1{zqOQX0qm^f9l&x<^Fy%^IMfg7N5J_+Qbd!P%lQ?{J%gv2; z>x~;L+N1_uZogrP&B;(V(O67KUOx@Wl#7B=(plw3P4~c%wuwRSJb%A$@|zo<*f28B zU0eNVpVq125{sBSj>A%Gdo|3gWHK+vcRk~yysdRbS%Z-6@h@*T$Kn5mlGk;QT1 z7w#(FQJI=MUiPmx_Bcez0%bBj72=rd@;t_?ffvoPsN(Ss^J1Ik<`QV9G}g>nnyP&A z`x<}wyVl3gDHbHDjbk9KR5d>n{X$rmed8e1qT-)XBRW!qSa7*h@#4AKm3kzK;Ej#9 z!IXJ9u`azxnl6~KUUKZ$#75Rp@@tQ*ANG)yZ*l}}{)qk(m2lXA@@THXth_cw^uw0* zg|fI~l1Lruk7*!L*y3``PATNVn-zD)D9{sCHaf;dCoL>Mem2YHhH0~Y$LAC8_sB)l z69v1Yrfv>3n)RlYIBN89BC6#xWPX2~;iiDN0i#vgYE* z=J9>q|52MY?tM;xc_Ld~InOAiZUBRM(I2{bVYK;GOb^`32(=LnsOqq%9-@eOh4fyG zT)V*ZJ3J9$qLllv>OT2kaNd15zBxmCp=q=3%#$4Hp^tNqKE%wHNuDBvoG#C^AS{Z` zcf46Ui-db|X{*}gUUDnj-HC*j&|3|{mNzzjRJx)?g{|joeS_sQ>dGc;k_B&!x{xm1 zSiAjlNwvzR6s?im+Rm8Yad(4F*~(vJgh;Ocs3{GJ?-a__j#9`;>g2Q+XAKzsurMul z@4@Dsl1q}mUe-{q#lsV)ZdP!_`t~dS&9&l=OK_eQ?9Eto`Mr|ZDtPsJ!l+zk*iSQU z#qwI%+=+|d9V0bB%s==RcjNw&&c@ZD`n_>X(|3w#b@V-$F=uj`OwkbBzxe}lE?Ak8 zh!InH)ftKiEt9<+lOg{;;KiL)yCjR%sg;gZ>gUy^T6^tPYDX;+Ge7vM6GIIkICm_K zx;dB)aS=s&YkiRXwGj^5ub#}&y?i9i%jcqd`2%z>pG@=e`?9@h*gIx7mRT04K9*gU ztUi`o_Cg$};33o9D(OwlYQI@lAn2 z|C_D_XQ9LNntB+#(yOp@-Mat&n5I3}JpD~s;Y^VBqN@~|me`OE{V%)UFu^mdUegGD zLuG8d(M+k$qj&#j&pzP8k_3VO&*pv)KB>9cZ>{VPEwK@lUp}s@5nH^eFzC4$LVnfS z7lIA5{)ea#6^@~a9ju`6^Y$MeKbg|H_T}LZEwL^tS_BwqW!J3zxC47|bc0_bMX-vl zK4@nB57&Mzz?;!N8~m{LPsa1XSASP_^*Wgx4liW4{wSLJ1V zA6#rzfP!j7R!Q4-2`tO@i^~-ZqeCIOld_BABL$2MP(~ON8gyf1gl>bu(3WnDTcA5& z0_apXMhKJzCV`%GW86BhJ^Fx;mv~fR)D{RcAPg23@FQ5kgWZBC=v0Q)=7)uis`wO|d(T}n2tk83?Ni->H zba#qxFN+7rUW{gf90DGJU~xAS#Zkn-K>!Fz%=-^1kDDK1}NYpO;n7ZgSt3CXdfUDe5FzBkc$8- zLZ+0D19BPQM4(HfI3Os17vWaQw+nI=IEsj(iHdkBuMBvgDNR%W_dXNV(U#2rinVtu zVyUo?G(Y&CZ0~jYy+6wz=v7L5B5`Mz;f-D+xo#5OuJ^?wtnMETY$_#&OT5^PvhOvz z!9(@M|7LYxAdqw0&wSs9s^~Y$rkTjU?&b&KlZTESU^((H)K3~x-t@~!L(d{er(gd^ zcRm(QD8zI5{AZC6fUkd$F8SX@!h=~WoUHK1XNDhh`o}K;>4Cyz;@g$f8yMRCQ(>WN zg#OU^D9yF#7go64{T|$z^R#aMi;cfl=70@QaTxOJ8pwg)13v}n&Gm&!pB=anxFgsG z+8=tMs14z3UHzJ;UZCQ972=6?^&6g>$160d5MkEUZ+YsE|9zO`P`Kn#pI-|Y2g^^H z6J1o&5|T>xw%2x@2`!@v6|_xI@4=V6rRUgI54$$tsCdj6oEo%@ke-4+(CD*zQgNAg l-vi6OqV<55Gc_0JNBd`lTiwFIu(}FiAQNJL6$WNlFm8@6m`jcd&K|Cp*v68^b9>ujE>ezHdPbSOXv;*_!6H>qHa0u-7V1wear)5q z&<>U93_Kn5lJ|5uBW1~9+f^hzB-AMSY-{Np?P}2cbkO2_!$jcRlmI+%`41W;h=vPV z8lYcAE;D+g;r_HtXkP#$swe?Dak2c;*KN(pADbfJl8kM`O|QIl>XrB7V&)^Woz7}X z(zkUsz4Fm{ul%j2#!t?E%xjELciVcbv@X|rtgP<6^%#*a5D?-8ZRMAjxAe}3p$AB$ z>q!K7J9F^B(Z$92ja*Qk)kk!3@VLCc%(~}mxuz&JCn8Kh04kyNNb<(fz5d>FlamA3 zqYCx|H8tVkD-I2oae5X~LD;ob%XjYUKx%ljgqHdcr}Xi;uh>me_CE|iJqtd5UwW;D zc%hR1k&{ZK2Cg5ji@u^(`_T=zLvOv`-Ea=;-h+OyB@;RG|a=62*r&o&WgSvgil>ig z2Du2Ibd+$#f5=4O5CK^;_>N3x-0OfW8oodCCJqspHIM(AnT30;^TYd;|pMhiIVKLVR`=hMxBZ z@c^>0{HlHq5CH6leW*f#K>UC@Y_h5!1QG=P0~4%9ae{;a1PoT~`>}~L*&DP4Fo&V4 zlQ)1w06SQ4^$sv-8*mPmT&*L@f#R})xmEXX06_rz{FpoH$_I>r5I(RPq#tz$8^RAh z44FiIz=jBb;Se^oU}%US_$WjWZ4??J1lH|10lKbYUnM4wt_nS}3498ofo2OkvKedw zIf2#=JF*3A4zWQ8haK4pwt~2ztHO?m6t}p^M1@a+(U1r<>^l24up{I?+U+{KC>R5I ziq5~zE(X2`DM3$OXBP)whP*<-qe|R+%91P*K z-x0w~Cls>T9u-l&6C45AZ7rUf2z)=Rs6 z%vnF}@iAvaTJAAtgR~0H^~-(2zl&Xv_XhoIAb&1rY$ z>ALHUHQ-F(QfSuO_q<0^@}^qJF}EB2!qrmOR?I>+yZ+H|&Y3NPe@EZ12mTmQXyT_l zqc#O*F7*H2xc*fsNr?o!->u(_rC&ISoOp^jEby>3`ee%Y%mLwWgYxu7eqXZNjT$Nb zgo8F8|5>?&pLy&%!7PBd!sU`$*0f(K@q4zmyyak~{VzieH|@dH?>_k;?a*BnOzgRn z-wO2;)vB!%59yzl#!z>ke4N&KS0y0!i(v~VBJqg+X&DUvhW}Z=JX8OPid?ALzlMc1 zs0di5x(3Y+A}fPvGisE=JUhNl&8LUk!>(4iB8x_8tt}-AfsCxKiE*LZ6iP+jC&UuP zFi^Q0K7j8zOv$3jboD){OYL#JQojw!0qi#~U^_f*Kgw4Y5b4k&z^ zfwOCjw`qjK=Qp%n@5%6r{-k@Qi8yqoTSv;mT`CCNm`973hmU7pPN^SBf2DwIM@9|D z+jhn?-aKy_H0^QfRTIB|{JBJF*}G;#>v-!@q&D@9rkYnk6}{s9eJA&z4?_|TkBZG2 z;}?%hQMOmokJ`SVe=C{S9t@*e;c^RBsvzw_dAmo`1KJzq`NV3%xr|vD8jkLZU6SbT z?&&E_X1Gh}2hiKH$6hbFznE#+dx0itAl>EEl&RcC{8DX6T&cr;n&aG{McqiRPMmBI zMx-OmQ>$y1neywU#pg&IZLF8x(eAk8In>ds`SxwXi$ZNW~fkvbjIYmK}(afX@n=X10KDTP_0{k470O_1nMpUZH8 zt0^EwZ#w0mTen8tnZMsSQ%tS?415=tCLB&HF9gLuTX+k?u!d}kI4O`C_i6N z@%${YEJvSOOvj`;2Eq%OJHD7M72QDUNPX9r9pe&swPZ+ra3R#<67%pFe$na^v%fIz z6C;L?rhC>q#5ee`4wc9qKRdgW7hRTHpS5_-_;C&fCLIktmu@~Xq3o&IM?BwGMBO<0 zN|Sc5Zh^SF_kbz0in*U~xSnnXez^VKJ>7{?dVYg)wrXzWH>LxzBpw@1*)sd+(!I9j z;fe|2?)q}N5A*!Q(_2eD>^x`YULtj6ZaNcua3RwR_4N39W(AOz-_&X2JQmc#QFtu9 zfW9rlFVQUC3SZeuPW9O$plsdVp621yU2+OfkE(i($RQ&F?!Hv)_m#gGBo%$kp@!N= zVT?8Pihp0JOXnPja66A5OAn zBA2_uZ~>iW`OJe=8QkOfhX~xls;nYYng5|5axgD-+M9gI+{@8oR`G3LAY4U-_NXQF zV_TT#kf_Q`^FkIG!AQ!G@Y7If&qzAtz3~{*s{88n;1Gu@^TaR3zR2@G&}vC@(Af=d znTY4V{f@X2oRc(CdNAz#JNl$weOqJ2#Gc-^=y7yeEHgHE_@}bkg^q433(BXyE#_AT zU2}uF1ew=b?x$r{AJxNNaQ)`IE3r)hS5cLemUh;?RvBzlR7%%4+Dj17xm?!ioEdeQ zs>VM@YC~sdiHG@3#seOK?WQ`qlJDbZeE$KhnXPxr-7x^`}h${SYyx4-5+qc z9=&Z%IaDHZJd)S!FV@Edi&7S=V1nx{Ybbj~!5Uzgi2CErH&SYAn0&OHs7hhM`ok=I zU9$qWR}_iO!E2bD%zsmth0Jk;EuZ^tTCSn&pK$wqvdel*E}pf+ak55zu46Z1uM3%^ zqvT>*qjSiGT%&Wu#k5w{q*DrE>|$D{bH;^Sr-ODet=GZ03}JnjI(8?%*61^NVi>O# zzhF9gy1?)ct$4cW#OZ3oE%5k7)9KUAhGy{ifRPI0xv1B!ae=4p?{g}7r!I=kP}tb8 zYyc+=QWeaGFMb?pc*SOHRwV#6K* z!22O@QQFuj70xQbLR9~?YF;o5Qj3z@>3PmM1;QZ)T!K|q@7M_14tT-(tCKf^pgLw)A8IfF9q>@F$;pF zLNV$>-Uh@7a-?2w1o^feBZ7QKZ~F~$te)l#a-5#+4e~v`0Pgir|JCV-QZr{J!T&-6 z)~}!PI8hIOD^zX$E7sA6W|ANaS4$R2uR`LK7^}0Tx;mR(e+PT_3%?_?o!G8j{wt8a z2C2`d&z>7`{mA!Q*qa&G;RCLxXK>C(&6z%C0sb#O|2uS-J1L*`U;{JlzdP2U{KtDvv_$GB?$`I&gK63QF7XrQu8Qk5 zf{I6M#;KMb{X`A@7t)x8-S2Q|+;J+{*bJi;+zE3Vq(=(p>Az*1#3!iTTv9frU zEF8qQFIIibYsI~^z9P@nl>o6 zbH+%W2HMis^sPlUhb0SFtnvu8GY1HYw;X52n>TgpzHwO?mNGM8$;`YSku#hL+vWKk z6`@IfO)rY0Rzr}*Bh5x;I_p1YCLK5x>HB^E*qMslO3_w>XN3r+j82o)qdyijy@Zyl z@b1-&^Lm(yget7$E1Ql?{{HP)nXS--Up=oH8ETuMNe!zeAhOa}GLo7`J&dMWwvSb4 zNJNpn=O<^9atjdnIF<~n<$V9LG?m%n5a^Zk?B10kTSU4kOO{`@jJn78xwv{vZ_b!~ zQ)Dwj;{}{B@#ZWEP~S)FAv9AR9~kz$ow5n>>YeZ`MT(qRom=Uco&Y8wKlPPGX?q1d zI9J$5^{gZ$tW0IjX51^ZqlmSX(F*cXAXd>KO~s$z_*3ny(#=?c!?(Sm&(|k94cBhU zTkWc_T(MCzad4scHrRyr)TuG=WyyGJ_H&cfJeLD=FkaCirxMD!bC?9+0%m?D z$K#%pB6 zl^&BmQ^M<%qxmrQ!{D}=W@auq;b( z(pLh?sD7$1+ZbJ(aqoMri>d7*wiN4?Q)^~|i#BzoT|HGaTEvFv@mKkLyYABh6xP=5Hcy-u9TrFCYSqFryUQ=smw!(;RF*UP z@ktyv$nk*t`s_VfwsbTLT4oVNgti;6e6L2_j1x~%eb5CDatu^#D+y7f&#B`oXV!f6 zpU}|ks_3)#v4O_%cwg(N@-)kOnh$oF&H!Z3N;8zb~ z+p~M!jC!3S<%=gB0`3&ds<$Wb>_ z#IPIjqaI4vk3L&;L7=TOps)h*G?XP@op`qUm1{padDugoddX&mNjP;+t=&v}} zU)#$ConeXTuN6Db z1Aa(A?IOz~! z3CVXup{q&QFm7Ivt*T??j8f)cUg*c6Cb;%b67iuQs>tl~uYZQdo84GcVWp8Tt}l-G z(a492yEA9An8u2OWyBF0CRy$DRZ78R5Qn4>?j`9$AIUI>YW_(!n4YF?BBM-tY&{epLraMtgJNde^Lsjf_=`+ z_|G=~(`#|n0;B!RmT-E$F|Sx%zA>NJ=X_%Uu?+>rLfv|f2{OkEjJJq678r|&g%ucY z6H6~h5*V3xxO-l9+S?#Z%WT1PX zCR=H|aPkIPxVAs)n=3g8#0GG{1glVNAP#^B2CM34191YvFrz9II|vMj!rZF**+E=y zH$9O$RwHJFG71UV02YCuPJn?$*>5LHqS%M*d4HdH%Jc1g;kMu z@PhUM#jt)-GA~FGsDQE6>;OPYKs{_nO)>yd2HwN8Ys^Krr9Tns3+BVB0zEMA8a96H zLEsbYeoZhxRvj3FmDI2aU^Rg6u#TEw0jw6Vn4eT1fawoc6RbrEg8l>WD8Xv`1wqFE zVI`wlln_V<5LI%k?H2;+0lSp$)S`qzh_N)|i=u6T2Xg~lgdNX!o4AHl@B2TE?KRBI zr4W;D>?c}0_$$;Lw;+Vy8KW+Ymx}Qd#_tUwu(-nw@_oIe8{|Yi;tld6z1|z-$9jw# zl^DDo;D_LHvBf!9>L?{L=mek`~TNT*BLQg!}^k*Z`{iwWqz|ypNJQ^ zpBVCsxBpq7W<=6^@<^@NFDtnksMp9~hK|B-#&XTzib;Zc5gtBv8m@YyCF?q~H$Qdn zi{^QYe}UvIS4p~pzjF<9+kUg9KOk(8=%9BO`1juc8t;V)jeD%SapE{6D?>smEW!TR4UO{r7=@!u%|w ztdcPh9mH7bZKssK{kYUa8BACzj%FN%(!+yd;R_8*N?Nng^ZHAcxVdO`g^`k{gDJOf e{Y9B$8eEsBqL(??*p`<;Y~SIZBiWTW*!~OdG1933 diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index 1b25a5a71..e14bf5b1b 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -82,7 +82,7 @@ class ScriptedProportionalPolicy(Policy): clock, no images. Codec stamps/truncates; the harness anchors/schedules. """ - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _ScriptedSession() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 345145f21..c3aeb6216 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -18,7 +18,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 DelegatingSession, Policy, SchedulingWrapper, Session +from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session from positronic.policy.codec import ActionTimestamp from positronic.policy.harness import Directive, DirectiveType, Harness, TrajectoryPlayer, _assert_anchored from positronic.policy.remote import RemoteSession @@ -82,7 +82,7 @@ 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, *, now=None, gate=None): + def new_session(self, context=None): self.reset_calls += 1 self.last_reset_context = context return _SpySession(self) @@ -127,7 +127,7 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -157,7 +157,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self.reset_calls += 1 self.last_reset_context = context return _ChunkSession(self) @@ -192,7 +192,7 @@ def __init__(self, command: roboarm.command.CommandType | None = None, target_gr self.command = command self.target_grip = float(target_grip) - def new_session(self, context=None, *, now=None, gate=None) -> RemoteSession: + def new_session(self, context=None) -> RemoteSession: action = [{'robot_command': self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] return RemoteSession(_FakeInferenceSession(action)) @@ -536,7 +536,7 @@ 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, *, now=None, gate=None): + def new_session(self, context=None): return _StaticMetaSession(self._command) # Session.meta defaults to {} @property @@ -922,7 +922,7 @@ def test_timeout_during_inference_drops_the_chunk(world): ChunkedSchedule().wrap(policy), make_embodiment(simulated=True), task=Task(instruction='test', timeout=0.05), - trials=[{keys.INFERENCE_LATENCY: 0.2}], # the gate holds the answer well past the deadline + trials=[{keys.INFERENCE_LATENCY: 0.2}], # the charge holds the answer well past the deadline ) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() @@ -1088,7 +1088,7 @@ def __call__(self, obs): return [] class EmptyChunkPolicy(Policy): - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _EmptyChunkSession() harness = Harness(EmptyChunkPolicy(), make_embodiment()) @@ -1574,11 +1574,11 @@ def __init__(self, wall_sec: float = 0.0, span_sec: float = 0.2, steps: int = 10 self._span_sec = span_sec self._steps = steps - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _SlowSession(self._wall_sec, self._span_sec, self._steps) -class _ReplanEarly(SchedulingWrapper): +class _ReplanEarly(PolicyWrapper): """Infers on the first observation and again halfway through the chunk it returned. The re-query-before-exhaustion shape (RTC, temporal ensembling) that the substrate exists for: unlike @@ -1586,23 +1586,25 @@ class _ReplanEarly(SchedulingWrapper): """ class _Session(DelegatingSession): - def __init__(self, inner: Session, now): + def __init__(self, inner: Session, charge_sec: float): super().__init__(inner) - self._now = now + self._charge_sec = charge_sec self._replan_at: float | None = None def __call__(self, obs): - if self._replan_at is not None and self._now() < self._replan_at: + t0 = obs['obs_time_ns'] / 1e9 + if self._replan_at is not None and t0 < self._replan_at: return None result = self._inner(obs) assert result is not None, 'the inner policy of this test wrapper always returns a chunk' - now = self._now() - result = [{**action, 'timestamp': now + action['timestamp']} for action in result] - self._replan_at = now + (result[-1]['timestamp'] - now) / 2 + anchor = t0 + self._charge_sec + result = [{**action, 'timestamp': anchor + action['timestamp']} for action in result] + self._replan_at = t0 + (result[-1]['timestamp'] - t0) / 2 return result - def wrap_session(self, inner: Session, context, now): - return _ReplanEarly._Session(inner, now) + def wrap_session(self, inner: Session, context): + assert context is not None # the harness always passes the trial context + return _ReplanEarly._Session(inner, float(context[keys.INFERENCE_LATENCY])) class _TimedRecorder(pimm.SignalEmitter): @@ -1702,7 +1704,7 @@ def __call__(self, obs): return [{keys.ROBOT_COMMAND: command, 'timestamp': i * 0.01} for i in range(10)] class _GripThenArmPolicy(Policy): - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _GripThenArm() harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index c743e1583..5d6584492 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -164,12 +164,12 @@ class _ChunkPolicy(Policy): def __init__(self, actions: list[dict]): self._actions = actions - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _FixedSession(list(self._actions)) class _SinglePolicy(Policy): - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _FixedSession({'v': 42}) @@ -183,7 +183,7 @@ def encode(self, data): class _MetaPolicy(Policy): - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _FixedSession({}) @property diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index e6704a314..89e98f458 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -27,7 +27,7 @@ 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 - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self.session_count += 1 return _TrackingSession(self._actions, {'policy_key': 'policy_value'}) @@ -57,7 +57,7 @@ def __init__(self, rec, actions): self._actions = actions self.last_session = None - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self.last_session = _CapturingSession(self._rec, self._actions) return self.last_session @@ -146,7 +146,7 @@ def __call__(self, obs): return None class _NonePolicy(Policy): - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _NoneSession() session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index ee5f58f44..5ed109a3c 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -1,10 +1,7 @@ -"""Unit tests for PolicyWrapper composition, the InferenceGate, ChunkedSchedule, TemporalStack, and the +"""Unit tests for PolicyWrapper composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" -import concurrent.futures -import threading import time -from concurrent.futures import ThreadPoolExecutor from typing import Any import numpy as np @@ -20,7 +17,7 @@ JointDeltaAction, RelativePositionAction, ) -from positronic.policy.base import InferenceGate, LatencyMode, Policy, PolicyWrapper, Session +from positronic.policy.base import Policy, PolicyWrapper, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -35,19 +32,6 @@ from positronic.policy.wrappers import ChunkedSchedule, TemporalStack -class _FakeClock: - """Minimal clock stub for unit tests — caller sets ``t`` directly.""" - - def __init__(self, t: float = 0.0): - self.t = t - - def now(self) -> float: - return self.t - - def now_ns(self) -> int: - return int(self.t * 1e9) - - class _ConstSession(Session): def __init__(self, actions, wall_sec: float = 0.0): self._actions = actions @@ -61,17 +45,21 @@ def __call__(self, obs): class _ConstPolicy(Policy): - def __init__(self, actions): + def __init__(self, actions, wall_sec: float = 0.0): self._actions = actions + self._wall_sec = wall_sec self._session: _ConstSession | None = None - def new_session(self, context=None, *, now=None, gate=None): - self._session = _ConstSession(self._actions) + def new_session(self, context=None): + self._session = _ConstSession(self._actions, wall_sec=self._wall_sec) return self._session _ONE_ACTION = [{'v': 1, 'timestamp': 0.0}] +# The sim default: model calls charge the world nothing, so chunks anchor at their observation instant. +_NO_LATENCY = {keys.INFERENCE_LATENCY: 0.0} + def _obs(now_sec=0.0): return {keys.OBS_TIME_NS: int(now_sec * 1e9)} @@ -80,11 +68,9 @@ def _obs(now_sec=0.0): class TestChunkedSchedule: def test_first_call_runs_inference(self): # Relative timestamps: trajectory of duration 0.5s - clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - policy = ChunkedSchedule().wrap(inner) - session = policy.new_session(now=clock.now) - result = session(_obs()) + session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + result = session(_obs(1.0)) assert result is not None assert len(result) == 2 # Timestamps stamped to absolute by ChunkedSchedule. @@ -92,78 +78,88 @@ def test_first_call_runs_inference(self): assert result[1]['timestamp'] == 1.5 def test_returns_none_while_trajectory_active(self): - # Trajectory starts at clock=1.0, ends at 1.0+0.5=1.5. - clock = _FakeClock(t=1.0) + # Trajectory starts at 1.0, ends at 1.0+0.5=1.5. inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - policy = ChunkedSchedule().wrap(inner) - session = policy.new_session(now=clock.now) - session(_obs()) - clock.t = 1.2 - assert session(_obs()) is None - clock.t = 1.4 - assert session(_obs()) is None + session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + session(_obs(1.0)) + assert session(_obs(1.2)) is None + assert session(_obs(1.4)) is None def test_re_infers_after_trajectory_consumed(self): - clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(now=clock.now) - session(_obs()) # trajectory ends at clock=1.5 - clock.t = 1.3 - assert session(_obs()) is None - clock.t = 1.6 - result = session(_obs()) + session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + session(_obs(1.0)) # trajectory ends at 1.5 + assert session(_obs(1.3)) is None + result = session(_obs(1.6)) assert result is not None assert inner._session.call_count == 2 def test_single_action_refires_immediately_after(self): - """Single action at ts=0 → trajectory_end = now → next tick re-infers.""" - clock = _FakeClock(t=1.0) - policy = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(now=clock.now) - session(_obs()) - clock.t = 1.01 - result = session(_obs()) + """Single action at ts=0 → trajectory_end = the observation instant → next tick re-infers.""" + session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(_NO_LATENCY) + session(_obs(1.0)) + assert session(_obs(1.01)) is not None + + def test_constant_charge_anchors_at_obs_time_plus_charge(self): + """The anchor is ``t0 + C`` whatever the call's wall duration — the reproducible mode.""" + inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.02) + session = ChunkedSchedule().wrap(inner).new_session({keys.INFERENCE_LATENCY: 0.5}) + result = session(_obs(1.0)) assert result is not None + assert result[0]['timestamp'] == 1.5 + assert session(_obs(1.4)) is None # within the charged window the chunk is still due to play + assert session(_obs(1.6)) is not None + + def test_measured_charge_anchors_at_obs_time_plus_wall_duration(self): + inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.05) + session = ChunkedSchedule().wrap(inner).new_session({keys.INFERENCE_LATENCY: True}) + result = session(_obs(1.0)) + assert result is not None + assert 1.05 <= result[0]['timestamp'] < 1.5 + + def test_no_latency_key_charges_wall_duration(self): + """Hardware passes no latency key; each call is charged what it really took.""" + inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.05) + session = ChunkedSchedule().wrap(inner).new_session() + result = session(_obs(1.0)) + assert result is not None + assert 1.05 <= result[0]['timestamp'] < 1.5 class TestPipelineComposition: """Test | operator across PolicyWrapper and Codec types.""" def test_wrapper_pipe_wrapper(self): - clock = _FakeClock(t=1.0) pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) policy = pipeline.wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(now=clock.now) + session = policy.new_session(_NO_LATENCY) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None assert result[0]['v'] == 1 def test_codec_pipe_wrapper(self): - clock = _FakeClock(t=1.0) codec = ActionTimestamp(fps=10.0) pipeline = codec | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) policy = pipeline.wrap(_ConstPolicy([{'action': 'test', 'timestamp': 0.0}])) - session = policy.new_session(now=clock.now) + session = policy.new_session(_NO_LATENCY) result = session(_obs()) assert result is not None def test_full_pipeline(self): - clock = _FakeClock(t=1.0) codec = ActionTimestamp(fps=10.0) pipeline = ChunkedSchedule() | codec assert isinstance(pipeline, PolicyWrapper) # 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 (clock=1.0). + # → ChunkedSchedule shifts to 1.0, 1.1, 1.2, 1.3, 1.4 (obs at 1.0). policy = pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)])) - session = policy.new_session(now=clock.now) - result = session(_obs()) + session = policy.new_session(_NO_LATENCY) + result = session(_obs(1.0)) assert result is not None assert result[0]['timestamp'] == 1.0 # Second call within trajectory window returns None (ChunkedSchedule). - clock.t = 1.2 - assert session(_obs()) is None + assert session(_obs(1.2)) is None def test_codec_and_stays_codec_only(self): """& only works between codecs, not wrappers.""" @@ -247,19 +243,17 @@ class TestTemporalStack: OFFSETS = (-0.2, -0.1, 0.0) def test_pad_start_repeats_oldest(self): - clock = _FakeClock(t=0.0) inner = _CapturePolicy() - session = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS).wrap(inner).new_session(now=clock.now) + session = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS).wrap(inner).new_session() session(_stack_obs(0.0, 1.0)) stack = inner.session.seen[0]['v'] assert stack.shape == (3, 1) assert (stack == 1.0).all() def test_no_pad_start_grows_from_one(self): - clock = _FakeClock(t=0.0) inner = _CapturePolicy() wrapper = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS, pad_start=False) - session = wrapper.wrap(inner).new_session(now=clock.now) + session = wrapper.wrap(inner).new_session() session(_stack_obs(0.0, 1.0)) assert inner.session.seen[0]['v'].shape == (1, 1) @@ -276,10 +270,9 @@ def test_no_pad_start_full_window_matches_padded(self): offsets = self.OFFSETS stacks = {} for pad_start in (True, False): - clock = _FakeClock(t=0.0) inner = _CapturePolicy() wrapper = TemporalStack(keys=('v',), offsets_sec=offsets, pad_start=pad_start) - session = wrapper.wrap(inner).new_session(now=clock.now) + session = wrapper.wrap(inner).new_session() for i in range(4): session(_stack_obs(0.1 * i, float(i))) stacks[pad_start] = inner.session.seen[-1]['v'] @@ -465,23 +458,20 @@ def test_pipe_composes_no_further(self): _ = pipeline | spec.PolicySource(_ConstPolicy([])) def test_inline_full_pipe(self): - clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'action': f'a{i}'} for i in range(5)]) policy = spec.inline(ChunkedSchedule() | spec.remote | ActionTimestamp(fps=10.0) | spec.PolicySource(inner)) assert isinstance(policy, Policy) - session = policy.new_session(now=clock.now) - result = session(_obs()) + session = policy.new_session(_NO_LATENCY) + result = session(_obs(1.0)) assert result is not None assert result[0]['timestamp'] == 1.0 - clock.t = 1.2 - assert session(_obs()) is None + assert session(_obs(1.2)) is None def test_inline_tolerates_marker_less_pipe(self): - clock = _FakeClock(t=1.0) inner = _ConstPolicy(_ONE_ACTION) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) - session = policy.new_session(now=clock.now) - result = session(_obs()) + session = policy.new_session(_NO_LATENCY) + result = session(_obs(1.0)) assert result is not None and result[0]['timestamp'] == 1.0 def test_inline_bare_source_pipe_is_the_loaded_policy(self): @@ -581,104 +571,3 @@ def test_survives_a_wire_round_trip(self): rebuilt = spec.from_spec(RestrictImageSize(64, 48).to_spec()) assert isinstance(rebuilt, RestrictImageSize) assert rebuilt.encode({'cam': _image(480, 640)})['cam'].shape == (48, 64, 3) - - -class _BlockingSession(Session): - """Holds the call inside the gate until the test releases it, and reports when it was entered.""" - - def __init__(self): - self.entered = threading.Event() - self.release = threading.Event() - - def __call__(self, obs): - self.entered.set() - self.release.wait(timeout=5.0) - return _ONE_ACTION - - -def _run_gated(gate: InferenceGate, inner: Session, executor: ThreadPoolExecutor): - return executor.submit(gate.wrap(inner), _obs()) - - -class TestInferenceGate: - def test_declared_mode_parks_the_call_until_its_delay_has_passed(self): - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) - with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) - concurrent.futures.wait([future], timeout=0.05) - assert not future.done(), 'the gate let the call through before its declared delay' - clock.t = 1.5 - assert future.result(timeout=5.0) is not None - assert gate.t0 == 1.0 - - def test_measured_mode_parks_the_call_for_its_own_wall_duration(self): - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.MEASURED) - with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession(_ONE_ACTION, wall_sec=0.1), executor) - concurrent.futures.wait([future], timeout=0.3) - assert not future.done(), 'the gate charged nothing for a call that took 0.1s' - clock.t = 1.5 # past t0 + the wall duration, whatever it measured - assert future.result(timeout=5.0) is not None - - @pytest.mark.parametrize( - 'mode,delay_sec', [(LatencyMode.LIVE, 0.0), (LatencyMode.DECLARED, 0.0)], ids=['live', 'declared-zero'] - ) - def test_a_call_owed_nothing_is_never_parked(self, mode, delay_sec): - """Hardware pays what the model took, and sim's default charges nothing: neither owes the wrapper a - wait, so the call returns with the clock standing still.""" - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, mode, delay_sec) - with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) - assert future.result(timeout=1.0) is not None # the clock never moves and the call still returns - assert gate.entered is False - - def test_entered_is_visible_while_the_call_is_at_the_model(self): - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.LIVE) - inner = _BlockingSession() - with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, inner, executor) - assert inner.entered.wait(timeout=5.0) - assert gate.entered is True - inner.release.set() - future.result(timeout=5.0) - assert gate.entered is False - - def test_cancel_unparks_a_call_the_harness_no_longer_wants(self): - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) - with ThreadPoolExecutor(max_workers=1) as executor: - future = _run_gated(gate, _ConstSession(_ONE_ACTION), executor) - concurrent.futures.wait([future], timeout=0.05) - assert not future.done() - gate.cancel() - assert future.result(timeout=5.0) is not None # released without waiting out the 60s delay - - def test_gate_wraps_the_inner_session_of_a_scheduling_wrapper(self): - """The cost lands below the wrapper: by the time ``ChunkedSchedule`` anchors its chunk, the delay - has already been paid, so the anchor is the release instant.""" - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.DECLARED, 0.5) - policy = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(now=clock.now, gate=gate) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(session, _obs()) - concurrent.futures.wait([future], timeout=0.05) - assert not future.done() - clock.t = 1.5 - chunk = future.result(timeout=5.0) - assert chunk is not None and chunk[0]['timestamp'] == 1.5 - - def test_a_stack_without_a_scheduling_wrapper_gets_no_gate(self): - """``TemporalStack`` owns no plan, so it is not what the platform charges; its call runs straight - through and the gate is never entered.""" - clock = _FakeClock(t=1.0) - gate = InferenceGate(clock.now, LatencyMode.DECLARED, 60.0) - policy = TemporalStack(keys=('v',), offsets_sec=(0.0,)).wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(now=clock.now, gate=gate) - - assert session({**_obs(), 'v': np.array([1.0])}) is not None - assert gate.entered is False diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index a6e7eb006..c1b6019b5 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -1,18 +1,18 @@ """Composable policy wrappers — scheduling and temporal frame stacking. Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is -outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only -``ChunkedSchedule`` needs the live clock — it anchors a chunk to the instant the inference gate releases -its call, which the pre-inference observation stamp cannot give — so the harness passes ``now`` (a -``Callable[[], float]`` in seconds) to ``new_session``, and it reaches every session in the stack. +outermost), exactly like codecs. All timing comes from the observation: ``obs_time_ns`` is the world +instant the harness built the observation at, and ``ChunkedSchedule`` derives the chunk's anchor from +it plus the trial's inference charge (``keys.INFERENCE_LATENCY`` in the session context). """ +import time from collections import deque import numpy as np from positronic import keys -from positronic.policy.base import DelegatingSession, Now, PolicyWrapper, SchedulingWrapper, Session +from positronic.policy.base import DelegatingSession, PolicyWrapper, Session def _obs_time(obs) -> float: @@ -20,32 +20,29 @@ def _obs_time(obs) -> float: return obs[keys.OBS_TIME_NS] / 1e9 -class ChunkedSchedule(SchedulingWrapper): +class ChunkedSchedule(PolicyWrapper): """Wait for the current trajectory to finish before calling the inner policy again. Owns relative→absolute time conversion: inner layers (codecs, models) emit relative timestamps; - this wrapper anchors them to ``now()`` *after* the inner call returns, which is the instant the - inference gate releases it. Returns ``None`` ("keep executing the current trajectory") until the last - action's timestamp is reached, then calls the inner policy. + this wrapper anchors them to the instant the model's answer is due — the observation instant plus + the trial's inference charge. Returns ``None`` ("keep executing the current trajectory") until the + last action's timestamp is reached, then calls the inner policy. """ class _Session(DelegatingSession): """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" - def __init__(self, inner: Session, now: Now | None): + def __init__(self, inner: Session, charge_sec: float | None): super().__init__(inner) - self._now = now + # Seconds the trial charges each model call, or ``None`` to charge the call's wall duration. + self._charge_sec = charge_sec self._trajectory_end: float | None = None def __call__(self, obs): - if self._now is None: - raise ValueError( - 'ChunkedSchedule needs a clock to run inference: pass now (a callable returning seconds) to ' - 'new_session. The harness supplies it; a direct RemotePolicy.new_session() outside the harness ' - 'must too.' - ) - if self._trajectory_end is not None and self._now() < self._trajectory_end: + t0 = _obs_time(obs) + if self._trajectory_end is not None and t0 < self._trajectory_end: return None + wall_t0 = time.monotonic() result = self._inner(obs) if result is not None: # A single-action session may return a bare dict, and a no-codec path may omit @@ -53,9 +50,9 @@ def __call__(self, obs): # immediate action executes instead of raising. if isinstance(result, dict): result = [result] + anchor = t0 + (time.monotonic() - wall_t0 if self._charge_sec is None else self._charge_sec) # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). - now = self._now() - result = [{**r, keys.ACTION_TIMESTAMP: now + r.get(keys.ACTION_TIMESTAMP, 0.0)} for r in result] + 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 @@ -63,8 +60,10 @@ def cancel(self): self._trajectory_end = None super().cancel() - def wrap_session(self, inner: Session, context, now: Now | None): - return ChunkedSchedule._Session(inner, now) + def wrap_session(self, inner: Session, context): + # ``True`` (or no key — the hardware case) charges each call what it really took. + latency = (context or {}).get(keys.INFERENCE_LATENCY, True) + return ChunkedSchedule._Session(inner, None if latency is True else float(latency)) def to_spec(self): return {'name': 'chunked_schedule'} @@ -157,7 +156,7 @@ def __init__(self, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_st 'in-range targets and the stack would be empty' ) - def wrap_session(self, inner: Session, context, now: Now): + def wrap_session(self, inner: Session, context): return TemporalStack._Session(inner, self._keys, self._offsets_sec, self._pad_start) def to_spec(self): diff --git a/positronic/probe.py b/positronic/probe.py index 12cc70642..a01ec93a0 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -134,7 +134,7 @@ def main( image_keys = [k for k in obs if k.startswith(keys.IMAGE_PREFIX)] rec = Recorder(pos3.sync(output_dir)) - session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None, now=time.time) + session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None) meta = dict(session.meta) name = label or _recording_name(meta) try: diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 3633e180a..6a7515940 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -312,7 +312,7 @@ def __init__(self, command: roboarm_command.CommandType, chunk_len: int): self.chunk_len = chunk_len self.chunks = 0 - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): return _JointposChunkSession(self) diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index 45ebd9bf5..7c9a97e43 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -286,7 +286,7 @@ class DreamZeroPolicy(Policy): def __init__(self, sp: DreamZeroSubprocess): self._subprocess = sp - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): client = RoboarenaClient(port=self._subprocess.roboarena_port) client.connect() return _DreamZeroSession(client, str(uuid.uuid4())) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 0ae1f7458..eef076907 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -220,7 +220,7 @@ def __init__(self, groot: Gr00tSubprocess, checkpoint_path: str): self._groot = groot self._checkpoint_path = checkpoint_path - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self._groot.client.reset() return _Gr00tSession(self._groot.client) diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index 53c21a7ef..440288815 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -102,7 +102,7 @@ def config(self) -> PreTrainedConfig: """The checkpoint's own declaration of what this policy takes.""" return self._policy.config - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self._policy.reset() return _LerobotSession(self._policy, self._preprocessor, self._postprocessor, self._device, self._meta) diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 8c41eb11b..589ca2518 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -92,7 +92,7 @@ def __init__(self, policy: PreTrainedPolicy, device: str | None = None, extra_me self._policy = policy.to(self._device) self._meta = extra_meta or {} - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): self._policy.reset() return _LerobotSession(self._policy, self._device, self._meta) diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index 76516a148..5b6165ea6 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -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, *, now=None, gate=None) -> Session: + def new_session(self, context=None) -> Session: 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 57f20f047..ffaa5ed81 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -140,7 +140,7 @@ class OpenpiPolicy(Policy): def __init__(self, subproc: OpenpiSubprocess): self._subproc = subproc - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None): client = self._subproc.client client.reset() return _OpenpiSession(client) From feb924efdc9b3c3b98db69415f53058a397a6ca9 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 22:12:07 +0300 Subject: [PATCH 04/64] Make `TrajectoryPlayer` consume a deque instead of tracking an index --- positronic/policy/harness.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 867194c8c..bc3b37b5a 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,5 +1,6 @@ import concurrent.futures import time +from collections import deque from collections.abc import Generator, Iterable, Iterator from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass @@ -40,16 +41,14 @@ class TrajectoryPlayer: to emit.""" def __init__(self): - self._trajectory: Trajectory = [] - self._index: int = 0 + self._pending: deque[tuple[int, Any]] = deque() def set(self, trajectory: Trajectory): - self._trajectory = trajectory - self._index = 0 + self._pending = deque(trajectory) def next_due(self) -> int | None: """Timestamp of the earliest waypoint not yet played, or ``None`` once the schedule is exhausted.""" - return self._trajectory[self._index][0] if self._index < len(self._trajectory) else None + return self._pending[0][0] if self._pending else None def advance(self, current_time: int): """The single value due at ``current_time``, or ``None`` when no waypoint has come due since the @@ -59,9 +58,8 @@ def advance(self, current_time: int): the final one. Revisit if pacing turns out not to hold one waypoint due per round. """ value = None - while self._index < len(self._trajectory) and self._trajectory[self._index][0] <= current_time: - value = self._trajectory[self._index][1] - self._index += 1 + while self._pending and self._pending[0][0] <= current_time: + value = self._pending.popleft()[1] return value From 6aab45b520f71012bf3059fa2517f473a122402e Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 22:37:08 +0300 Subject: [PATCH 05/64] Hand the policy stack a `now` clock; only the harness reads `inference_latency` --- ARCHITECTURE.md | 12 ++-- .../offboard/tests/test_remote_policy.py | 4 +- positronic/offboard/tests/test_server.py | 9 ++- positronic/policy/base.py | 24 ++++--- positronic/policy/codec.py | 2 +- positronic/policy/harness.py | 16 ++++- positronic/policy/recording.py | 6 +- positronic/policy/remote.py | 6 +- .../policy/tests/test_golden_pipeline.py | 2 +- positronic/policy/tests/test_harness.py | 28 ++++---- positronic/policy/tests/test_policy_io.py | 6 +- positronic/policy/tests/test_recording.py | 6 +- positronic/policy/tests/test_wrappers.py | 72 ++++++++----------- positronic/policy/wrappers.py | 42 +++++------ positronic/probe.py | 2 +- .../env_server/tests/test_remote_env.py | 2 +- positronic/vendors/dreamzero/server.py | 2 +- positronic/vendors/gr00t/server.py | 2 +- positronic/vendors/lerobot/policy.py | 2 +- positronic/vendors/lerobot_0_3_3/policy.py | 2 +- positronic/vendors/molmoact2/policy.py | 2 +- positronic/vendors/openpi/server.py | 2 +- 22 files changed, 129 insertions(+), 122 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 16633701c..25514e6d3 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -132,12 +132,12 @@ episode lifecycle — nothing else. Scheduling, blending, history stacking and e the wrapper stack around the policy; a session returning `None` means "keep executing the current trajectory". -**Inference cost is a fact of the trial, charged twice from the same source.** The trial context -carries `inference_latency` — a fixed charge in seconds, or the call's own wall duration — and both -sides read it independently: the scheduling wrapper stamps its chunk for the observation instant -plus the charge, and the harness withholds the trajectory, and the world clock, until that instant -arrives. Neither side signals the other; the observation's own timestamp is the shared anchor, so -the charge is deterministic wherever a fixed one is asked for. +**Inference cost is a fact of the trial, owned by the harness.** The trial context carries +`inference_latency` — a fixed charge in seconds, or the call's own wall duration — and the harness +alone reads it: it withholds a returned trajectory, and the world clock, until the charge is paid, +and the clock it hands the policy stack (`now`) reads the instant the in-flight call's output takes +effect. A scheduling wrapper stamps its chunk at `now()` and never learns the mode, so the charge +is deterministic wherever a fixed one is asked for. **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 388c31fc1..bc2c40507 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -391,7 +391,7 @@ def test_empty_declaration_fails_before_motion(): def test_declared_stack_built_at_session_open(): """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 = policy.new_session({keys.INFERENCE_LATENCY: 0.0}) + session = policy.new_session(now=lambda: 1.0) actions = session({keys.OBS_TIME_NS: int(1e9)}) assert actions == [{'a': 1, 'timestamp': 1.0}] @@ -511,7 +511,7 @@ def test_remote_policy_lifecycle(inference_server, mock_policy): host, port = inference_server policy = RemotePolicy(f'{host}:{port}') - session = policy.new_session({keys.INFERENCE_LATENCY: 0.0}) + session = policy.new_session(now=lambda: 0.0) meta = session.meta assert meta['server.model_name'] == 'test_model' diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 8710bbf48..4196ee09c 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -12,7 +12,6 @@ from websockets.exceptions import InvalidStatus from websockets.sync.client import connect -from positronic import keys from positronic.offboard.client import InferenceClient, InferenceSession from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer from positronic.offboard.server_utils import warmup @@ -244,7 +243,7 @@ def __call__(self, obs): class _ScriptedPolicy(Policy): """Deterministic base policy: every session returns the same untimestamped chunk.""" - def new_session(self, context=None) -> Session: + def new_session(self, context=None, now=None) -> Session: return _ScriptedSession() @@ -254,12 +253,12 @@ def test_in_process_equals_remote_for_same_pipeline(start_server): def pipeline(): return ChunkedSchedule() | remote | ActionTimestamp(fps=10.0) | PolicySource(_ScriptedPolicy()) - context = {keys.INFERENCE_LATENCY: 0.0} + clock = [100.0] host, port, _server = start_server(pipeline()) - remote_session = RemotePolicy(f'{host}:{port}').new_session(context) + remote_session = RemotePolicy(f'{host}:{port}').new_session(now=lambda: clock[0]) - local_session = inline(pipeline()).new_session(context) + local_session = inline(pipeline()).new_session(now=lambda: clock[0]) remote_actions = remote_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) local_actions = local_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 3cd6b435f..07930815b 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,9 +1,12 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any +Now = Callable[[], float] + + # Structural keys of the wire spec: ``|`` serializes as ``{SEQ: [...]}``, ``&`` as ``{PAR: [...]}``. SEQ = 'seq' PAR = 'par' @@ -78,11 +81,16 @@ class Policy(ABC): """ @abstractmethod - def new_session(self, context: dict[str, Any] | None = None) -> Session: + def new_session(self, context: dict[str, Any] | None = None, now: Now | None = None) -> Session: """Create a new inference session for an episode. Args: context: Episode context (task description, eval metadata, etc.). + now: The policy stack's clock (seconds), supplied by the harness and passed down to every + wrapped session. During a session call it reads the trial instant the call's output takes + effect — the observation instant plus the inference charge paid so far — so a chunk + stamped at ``now()`` lands when its call is paid for. ``None`` where no runtime clock + exists (server-side, warmup). """ @property @@ -100,8 +108,8 @@ class DelegatingPolicy(Policy): def __init__(self, inner: Policy): self._inner = inner - def new_session(self, context=None): - return self._inner.new_session(context) + def new_session(self, context=None, now=None): + return self._inner.new_session(context, now) @property def meta(self): @@ -133,12 +141,12 @@ class PolicyWrapper: def wrap(self, policy: Policy) -> Policy: """Apply this wrapper to a policy. Default: wrap every session it creates via ``wrap_session``. - Composition happens at config time; the trial context reaches the wrapped + Composition happens at config time; the trial context and clock reach the wrapped sessions through ``new_session``. """ return _WrapperPolicy(policy, self) - def wrap_session(self, inner: Session, context: dict[str, Any] | None) -> Session: + def wrap_session(self, inner: Session, context: dict[str, Any] | None, now: Now | None) -> Session: """Wrap a single session. Subclasses override this for per-session wrapping.""" raise NotImplementedError('Override wrap_session or wrap') @@ -176,8 +184,8 @@ def __init__(self, inner: Policy, wrapper: PolicyWrapper): super().__init__(inner) self._wrapper = wrapper - def new_session(self, context=None): - return self._wrapper.wrap_session(self._inner.new_session(context), context) + def new_session(self, context=None, now=None): + return self._wrapper.wrap_session(self._inner.new_session(context, now), context, now) @property def meta(self): diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 38ae68021..16ae3a000 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -80,7 +80,7 @@ def training_encoder(self) -> EpisodeTransform: def meta(self) -> dict: return {} - def wrap_session(self, inner: Session, context): + def wrap_session(self, inner: Session, context, now): return _CodecSession(inner, self) @final diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index bc3b37b5a..ac8ce746e 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -184,8 +184,10 @@ class Harness(pimm.ControlSystem): error recovery, blending, absolute time stamping — lives in the policy/session layer: the wrapper owns the plan, the harness plays it, one command per channel per round. The session call runs on a worker thread so playing continues while the model does; the harness withholds the trajectory, and the world - clock, until the trial's inference charge (``inference_latency``) is paid. The RUN context is handed - whole to the task's scene reset, which reads the per-trial keys it needs (e.g. ``eval.seed``). + clock, until the trial's inference charge (``inference_latency``) is paid, and the ``now`` it hands + ``new_session`` reads time the same way — so wrappers stamp chunks for the paid instant without + knowing the mode. The RUN context is handed whole to the task's scene reset, which reads the + per-trial keys it needs (e.g. ``eval.seed``). A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's @@ -349,7 +351,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._policy_session = self.policy.new_session(self.context) + self._policy_session = self.policy.new_session(self.context, self._policy_now) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -423,6 +425,14 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it return inputs + def _policy_now(self) -> float: + """The trial instant the in-flight call's output takes effect: its observation instant plus the + charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; + the loop thread writes the call's start fields before submitting it. + """ + charge = time.monotonic() - self._wall_t0 if self._charge is None else self._charge + return self._t0_ns / 1e9 + charge + def _step(self, clock: pimm.Clock) -> None: """Keep one session call in flight and install the trajectory it returns. diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index d9ed32bf8..5a094f796 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -36,7 +36,7 @@ (each action stamped at the inference-request time plus its horizon offset), so a ``TimeSeriesView`` reads commanded values with real axes. That anchor is the pre-inference ``obs_time_ns``, so it precedes true execution — which the scheduling wrapper anchors at the instant -the inference gate releases its call — by the inference latency. Select ``action_time`` to see them. +its call's charge is paid — by the inference latency. Select ``action_time`` to see them. Entity paths are ``{tap_name}/{data_key}``. A tap's incoming observation keys and outgoing action keys share that namespace; in the rare case the same key appears on @@ -222,7 +222,7 @@ def _log_action_series(path: str, arr: np.ndarray, horizon: np.ndarray, base_ns: Each action is stamped at ``base_ns + horizon_i``, where ``base_ns`` is the inference-request time; successive chunks lay out along one clock so a ``TimeSeriesView`` has real axes. This precedes true execution by the inference latency: the scheduling wrapper - anchors commands at the instant the gate releases its call, which a recorder tap cannot observe. + anchors commands at the instant its call's charge is paid, which a recorder tap cannot observe. """ arr = np.asarray(arr, dtype=np.float64) if arr.ndim == 1: @@ -342,7 +342,7 @@ def __init__(self, rec: Recorder, name: str): self._rec = rec self._name = name - def wrap_session(self, inner: Session, context) -> Session: + def wrap_session(self, inner: Session, context, now) -> Session: stream = self._rec._open_stream() return _RecordingTapSession(inner, self._rec, self._name, stream) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 438e66af6..cb5174436 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -102,7 +102,7 @@ def server_meta(self) -> dict[str, Any]: ws_session.close() return self._server_meta - def new_session(self, context=None) -> RemoteSession: + def new_session(self, context=None, now=None) -> RemoteSession: compress = bool(self.server_meta().get(keys.COMPRESS_IMAGES)) ws_session = self._client.new_session() return RemoteSession(ws_session, compress_images=compress) @@ -166,8 +166,8 @@ def _policy(self) -> Policy: self._stacked = stack.wrap(self._endpoint) return self._stacked - def new_session(self, context=None) -> Session: - return self._policy().new_session(context) + def new_session(self, context=None, now=None) -> Session: + return self._policy().new_session(context, now) @property def meta(self) -> dict[str, Any]: diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index e14bf5b1b..f0241c74e 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -82,7 +82,7 @@ class ScriptedProportionalPolicy(Policy): clock, no images. Codec stamps/truncates; the harness anchors/schedules. """ - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _ScriptedSession() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index c3aeb6216..c5992c35c 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -82,7 +82,7 @@ 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): + def new_session(self, context=None, now=None): self.reset_calls += 1 self.last_reset_context = context return _SpySession(self) @@ -127,7 +127,7 @@ def __init__( def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None): + def new_session(self, context=None, now=None): self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -157,7 +157,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.counter = 0 - def new_session(self, context=None): + def new_session(self, context=None, now=None): self.reset_calls += 1 self.last_reset_context = context return _ChunkSession(self) @@ -192,7 +192,7 @@ def __init__(self, command: roboarm.command.CommandType | None = None, target_gr self.command = command self.target_grip = float(target_grip) - def new_session(self, context=None) -> RemoteSession: + def new_session(self, context=None, now=None) -> RemoteSession: action = [{'robot_command': self.command, 'target_grip': self.target_grip, 'timestamp': 0.0}] return RemoteSession(_FakeInferenceSession(action)) @@ -536,7 +536,7 @@ 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): + def new_session(self, context=None, now=None): return _StaticMetaSession(self._command) # Session.meta defaults to {} @property @@ -1088,7 +1088,7 @@ def __call__(self, obs): return [] class EmptyChunkPolicy(Policy): - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _EmptyChunkSession() harness = Harness(EmptyChunkPolicy(), make_embodiment()) @@ -1574,7 +1574,7 @@ def __init__(self, wall_sec: float = 0.0, span_sec: float = 0.2, steps: int = 10 self._span_sec = span_sec self._steps = steps - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _SlowSession(self._wall_sec, self._span_sec, self._steps) @@ -1586,9 +1586,9 @@ class _ReplanEarly(PolicyWrapper): """ class _Session(DelegatingSession): - def __init__(self, inner: Session, charge_sec: float): + def __init__(self, inner: Session, now): super().__init__(inner) - self._charge_sec = charge_sec + self._now = now self._replan_at: float | None = None def __call__(self, obs): @@ -1597,14 +1597,14 @@ def __call__(self, obs): return None result = self._inner(obs) assert result is not None, 'the inner policy of this test wrapper always returns a chunk' - anchor = t0 + self._charge_sec + anchor = self._now() result = [{**action, 'timestamp': anchor + action['timestamp']} for action in result] self._replan_at = t0 + (result[-1]['timestamp'] - t0) / 2 return result - def wrap_session(self, inner: Session, context): - assert context is not None # the harness always passes the trial context - return _ReplanEarly._Session(inner, float(context[keys.INFERENCE_LATENCY])) + def wrap_session(self, inner: Session, context, now): + assert now is not None # the harness always passes its clock + return _ReplanEarly._Session(inner, now) class _TimedRecorder(pimm.SignalEmitter): @@ -1704,7 +1704,7 @@ def __call__(self, obs): return [{keys.ROBOT_COMMAND: command, 'timestamp': i * 0.01} for i in range(10)] class _GripThenArmPolicy(Policy): - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _GripThenArm() harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) diff --git a/positronic/policy/tests/test_policy_io.py b/positronic/policy/tests/test_policy_io.py index 5d6584492..452637d3b 100644 --- a/positronic/policy/tests/test_policy_io.py +++ b/positronic/policy/tests/test_policy_io.py @@ -164,12 +164,12 @@ class _ChunkPolicy(Policy): def __init__(self, actions: list[dict]): self._actions = actions - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _FixedSession(list(self._actions)) class _SinglePolicy(Policy): - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _FixedSession({'v': 42}) @@ -183,7 +183,7 @@ def encode(self, data): class _MetaPolicy(Policy): - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _FixedSession({}) @property diff --git a/positronic/policy/tests/test_recording.py b/positronic/policy/tests/test_recording.py index 89e98f458..e88acba1a 100644 --- a/positronic/policy/tests/test_recording.py +++ b/positronic/policy/tests/test_recording.py @@ -27,7 +27,7 @@ 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 - def new_session(self, context=None): + def new_session(self, context=None, now=None): self.session_count += 1 return _TrackingSession(self._actions, {'policy_key': 'policy_value'}) @@ -57,7 +57,7 @@ def __init__(self, rec, actions): self._actions = actions self.last_session = None - def new_session(self, context=None): + def new_session(self, context=None, now=None): self.last_session = _CapturingSession(self._rec, self._actions) return self.last_session @@ -146,7 +146,7 @@ def __call__(self, obs): return None class _NonePolicy(Policy): - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _NoneSession() session = Recorder(tmp_path).tap('t').wrap(_NonePolicy()).new_session() diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 5ed109a3c..ccb4c9403 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -1,7 +1,6 @@ """Unit tests for PolicyWrapper composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" -import time from typing import Any import numpy as np @@ -33,33 +32,27 @@ class _ConstSession(Session): - def __init__(self, actions, wall_sec: float = 0.0): + def __init__(self, actions): self._actions = actions - self._wall_sec = wall_sec self.call_count = 0 def __call__(self, obs): self.call_count += 1 - time.sleep(self._wall_sec) return self._actions class _ConstPolicy(Policy): - def __init__(self, actions, wall_sec: float = 0.0): + def __init__(self, actions): self._actions = actions - self._wall_sec = wall_sec self._session: _ConstSession | None = None - def new_session(self, context=None): - self._session = _ConstSession(self._actions, wall_sec=self._wall_sec) + def new_session(self, context=None, now=None): + self._session = _ConstSession(self._actions) return self._session _ONE_ACTION = [{'v': 1, 'timestamp': 0.0}] -# The sim default: model calls charge the world nothing, so chunks anchor at their observation instant. -_NO_LATENCY = {keys.INFERENCE_LATENCY: 0.0} - def _obs(now_sec=0.0): return {keys.OBS_TIME_NS: int(now_sec * 1e9)} @@ -69,7 +62,7 @@ class TestChunkedSchedule: def test_first_call_runs_inference(self): # Relative timestamps: trajectory of duration 0.5s inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 1.0) result = session(_obs(1.0)) assert result is not None assert len(result) == 2 @@ -80,50 +73,47 @@ def test_first_call_runs_inference(self): def test_returns_none_while_trajectory_active(self): # Trajectory starts at 1.0, ends at 1.0+0.5=1.5. inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 1.0) session(_obs(1.0)) assert session(_obs(1.2)) is None assert session(_obs(1.4)) is None def test_re_infers_after_trajectory_consumed(self): + clock = [1.0] inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(_NO_LATENCY) + session = ChunkedSchedule().wrap(inner).new_session(now=lambda: clock[0]) session(_obs(1.0)) # trajectory ends at 1.5 assert session(_obs(1.3)) is None + clock[0] = 1.6 result = session(_obs(1.6)) assert result is not None assert inner._session.call_count == 2 def test_single_action_refires_immediately_after(self): - """Single action at ts=0 → trajectory_end = the observation instant → next tick re-infers.""" - session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(_NO_LATENCY) + """Single action at ts=0 → trajectory_end = the anchor instant → next tick re-infers.""" + session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(now=lambda: 1.0) session(_obs(1.0)) assert session(_obs(1.01)) is not None - def test_constant_charge_anchors_at_obs_time_plus_charge(self): - """The anchor is ``t0 + C`` whatever the call's wall duration — the reproducible mode.""" - inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.02) - session = ChunkedSchedule().wrap(inner).new_session({keys.INFERENCE_LATENCY: 0.5}) + def test_anchors_at_now_not_at_the_observation(self): + """``now`` reads the instant the call's charge is paid; the chunk is stamped there, not at the obs.""" + session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(now=lambda: 1.5) result = session(_obs(1.0)) assert result is not None assert result[0]['timestamp'] == 1.5 - assert session(_obs(1.4)) is None # within the charged window the chunk is still due to play - assert session(_obs(1.6)) is not None - def test_measured_charge_anchors_at_obs_time_plus_wall_duration(self): - inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.05) - session = ChunkedSchedule().wrap(inner).new_session({keys.INFERENCE_LATENCY: True}) - result = session(_obs(1.0)) - assert result is not None - assert 1.05 <= result[0]['timestamp'] < 1.5 + def test_expiry_is_judged_at_the_observation_instant(self): + """Whether the trajectory has run out is a question about the observation, not about ``now``.""" + inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 2.0) + session(_obs(1.0)) # anchored at now() = 2.0, so the trajectory ends at 2.5 + assert session(_obs(2.4)) is None + assert session(_obs(2.6)) is not None - def test_no_latency_key_charges_wall_duration(self): - """Hardware passes no latency key; each call is charged what it really took.""" - inner = _ConstPolicy(_ONE_ACTION, wall_sec=0.05) - session = ChunkedSchedule().wrap(inner).new_session() - result = session(_obs(1.0)) - assert result is not None - assert 1.05 <= result[0]['timestamp'] < 1.5 + def test_running_inference_without_a_clock_raises(self): + session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session() + with pytest.raises(ValueError, match='clock'): + session(_obs(1.0)) class TestPipelineComposition: @@ -133,7 +123,7 @@ def test_wrapper_pipe_wrapper(self): pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) policy = pipeline.wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(_NO_LATENCY) + session = policy.new_session(now=lambda: 1.0) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None assert result[0]['v'] == 1 @@ -143,7 +133,7 @@ def test_codec_pipe_wrapper(self): pipeline = codec | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) policy = pipeline.wrap(_ConstPolicy([{'action': 'test', 'timestamp': 0.0}])) - session = policy.new_session(_NO_LATENCY) + session = policy.new_session(now=lambda: 1.0) result = session(_obs()) assert result is not None @@ -154,7 +144,7 @@ def test_full_pipeline(self): # 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 (obs at 1.0). policy = pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)])) - session = policy.new_session(_NO_LATENCY) + session = policy.new_session(now=lambda: 1.0) result = session(_obs(1.0)) assert result is not None assert result[0]['timestamp'] == 1.0 @@ -231,7 +221,7 @@ class _CapturePolicy(Policy): def __init__(self): self.session = _CaptureSession() - def new_session(self, context=None, *, now=None, gate=None): + def new_session(self, context=None, now=None): return self.session @@ -461,7 +451,7 @@ 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)) assert isinstance(policy, Policy) - session = policy.new_session(_NO_LATENCY) + session = policy.new_session(now=lambda: 1.0) result = session(_obs(1.0)) assert result is not None assert result[0]['timestamp'] == 1.0 @@ -470,7 +460,7 @@ def test_inline_full_pipe(self): def test_inline_tolerates_marker_less_pipe(self): inner = _ConstPolicy(_ONE_ACTION) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) - session = policy.new_session(_NO_LATENCY) + session = policy.new_session(now=lambda: 1.0) result = session(_obs(1.0)) assert result is not None and result[0]['timestamp'] == 1.0 diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index c1b6019b5..d97b8ae43 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -1,18 +1,17 @@ """Composable policy wrappers — scheduling and temporal frame stacking. Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is -outermost), exactly like codecs. All timing comes from the observation: ``obs_time_ns`` is the world -instant the harness built the observation at, and ``ChunkedSchedule`` derives the chunk's anchor from -it plus the trial's inference charge (``keys.INFERENCE_LATENCY`` in the session context). +outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only +``ChunkedSchedule`` needs the clock — it anchors a chunk at the instant its call's charge is paid, +which the observation stamp cannot give — so the harness passes ``now`` (a ``Callable[[], float]`` +in seconds) to ``new_session`` and it reaches that one session. """ -import time from collections import deque import numpy as np -from positronic import keys -from positronic.policy.base import DelegatingSession, PolicyWrapper, Session +from positronic.policy.base import DelegatingSession, Now, PolicyWrapper, Session def _obs_time(obs) -> float: @@ -24,25 +23,28 @@ class ChunkedSchedule(PolicyWrapper): """Wait for the current trajectory to finish before calling the inner policy again. Owns relative→absolute time conversion: inner layers (codecs, models) emit relative timestamps; - this wrapper anchors them to the instant the model's answer is due — the observation instant plus - the trial's inference charge. Returns ``None`` ("keep executing the current trajectory") until the - last action's timestamp is reached, then calls the inner policy. + this wrapper anchors them to ``now()`` after the inner call — the instant the call's charge is + paid. Returns ``None`` ("keep executing the current trajectory") until the last action's timestamp + is reached at the observation instant, then calls the inner policy. """ class _Session(DelegatingSession): """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" - def __init__(self, inner: Session, charge_sec: float | None): + def __init__(self, inner: Session, now: Now | None): super().__init__(inner) - # Seconds the trial charges each model call, or ``None`` to charge the call's wall duration. - self._charge_sec = charge_sec + self._now = now self._trajectory_end: float | None = None def __call__(self, obs): - t0 = _obs_time(obs) - if self._trajectory_end is not None and t0 < self._trajectory_end: + if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: return None - wall_t0 = time.monotonic() + if self._now is None: + raise ValueError( + 'ChunkedSchedule needs a clock to run inference: pass now (a callable returning seconds) to ' + 'new_session. The harness supplies it; a direct RemotePolicy.new_session() outside the harness ' + 'must too.' + ) result = self._inner(obs) if result is not None: # A single-action session may return a bare dict, and a no-codec path may omit @@ -50,7 +52,7 @@ def __call__(self, obs): # immediate action executes instead of raising. if isinstance(result, dict): result = [result] - anchor = t0 + (time.monotonic() - wall_t0 if self._charge_sec is None else self._charge_sec) + anchor = self._now() # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). 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 @@ -60,10 +62,8 @@ def cancel(self): self._trajectory_end = None super().cancel() - def wrap_session(self, inner: Session, context): - # ``True`` (or no key — the hardware case) charges each call what it really took. - latency = (context or {}).get(keys.INFERENCE_LATENCY, True) - return ChunkedSchedule._Session(inner, None if latency is True else float(latency)) + def wrap_session(self, inner: Session, context, now: Now | None): + return ChunkedSchedule._Session(inner, now) def to_spec(self): return {'name': 'chunked_schedule'} @@ -156,7 +156,7 @@ def __init__(self, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_st 'in-range targets and the stack would be empty' ) - def wrap_session(self, inner: Session, context): + def wrap_session(self, inner: Session, context, now): return TemporalStack._Session(inner, self._keys, self._offsets_sec, self._pad_start) def to_spec(self): diff --git a/positronic/probe.py b/positronic/probe.py index a01ec93a0..76c6d9cb9 100644 --- a/positronic/probe.py +++ b/positronic/probe.py @@ -134,7 +134,7 @@ def main( image_keys = [k for k in obs if k.startswith(keys.IMAGE_PREFIX)] rec = Recorder(pos3.sync(output_dir)) - session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None) + session = rec.tap(_TAP).wrap(policy).new_session({keys.TASK: task} if task else None, time.time) meta = dict(session.meta) name = label or _recording_name(meta) try: diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 6a7515940..e59ca2198 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -312,7 +312,7 @@ def __init__(self, command: roboarm_command.CommandType, chunk_len: int): self.chunk_len = chunk_len self.chunks = 0 - def new_session(self, context=None): + def new_session(self, context=None, now=None): return _JointposChunkSession(self) diff --git a/positronic/vendors/dreamzero/server.py b/positronic/vendors/dreamzero/server.py index 7c9a97e43..01c8c05f4 100644 --- a/positronic/vendors/dreamzero/server.py +++ b/positronic/vendors/dreamzero/server.py @@ -286,7 +286,7 @@ class DreamZeroPolicy(Policy): def __init__(self, sp: DreamZeroSubprocess): self._subprocess = sp - def new_session(self, context=None): + def new_session(self, context=None, now=None): client = RoboarenaClient(port=self._subprocess.roboarena_port) client.connect() return _DreamZeroSession(client, str(uuid.uuid4())) diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index eef076907..6071fb46e 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -220,7 +220,7 @@ def __init__(self, groot: Gr00tSubprocess, checkpoint_path: str): self._groot = groot self._checkpoint_path = checkpoint_path - def new_session(self, context=None): + def new_session(self, context=None, now=None): self._groot.client.reset() return _Gr00tSession(self._groot.client) diff --git a/positronic/vendors/lerobot/policy.py b/positronic/vendors/lerobot/policy.py index 440288815..fa83c9922 100644 --- a/positronic/vendors/lerobot/policy.py +++ b/positronic/vendors/lerobot/policy.py @@ -102,7 +102,7 @@ def config(self) -> PreTrainedConfig: """The checkpoint's own declaration of what this policy takes.""" return self._policy.config - def new_session(self, context=None): + def new_session(self, context=None, now=None): self._policy.reset() return _LerobotSession(self._policy, self._preprocessor, self._postprocessor, self._device, self._meta) diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 589ca2518..5e608c1b1 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -92,7 +92,7 @@ def __init__(self, policy: PreTrainedPolicy, device: str | None = None, extra_me self._policy = policy.to(self._device) self._meta = extra_meta or {} - def new_session(self, context=None): + def new_session(self, context=None, now=None): self._policy.reset() return _LerobotSession(self._policy, self._device, self._meta) diff --git a/positronic/vendors/molmoact2/policy.py b/positronic/vendors/molmoact2/policy.py index 5b6165ea6..b61a2e0e2 100644 --- a/positronic/vendors/molmoact2/policy.py +++ b/positronic/vendors/molmoact2/policy.py @@ -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) -> Session: + def new_session(self, context=None, now=None) -> Session: 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 ffaa5ed81..519b0ad37 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -140,7 +140,7 @@ class OpenpiPolicy(Policy): def __init__(self, subproc: OpenpiSubprocess): self._subproc = subproc - def new_session(self, context=None): + def new_session(self, context=None, now=None): client = self._subproc.client client.reset() return _OpenpiSession(client) From 1a67390538132c159dee412d1e2817c4eb30c8ba Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 22:49:49 +0300 Subject: [PATCH 06/64] Trim the diff to main: restore its text where behavior is unchanged --- .../offboard/tests/test_remote_policy.py | 13 +- positronic/offboard/tests/test_server.py | 9 +- positronic/policy/base.py | 14 +-- positronic/policy/tests/test_wrappers.py | 113 ++++++++++-------- positronic/policy/wrappers.py | 21 ++-- 5 files changed, 94 insertions(+), 76 deletions(-) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index bc2c40507..65af1fbb5 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -302,7 +302,7 @@ def test_no_truncation_without_horizon(self): def test_remote_session_normalizes_single_dict(): - """Server returning a single action dict is wrapped into a 1-element list.""" + """Server returning a single action dict (legacy shape) is wrapped into a 1-element list.""" endpoint, _ = _mock_endpoint(infer_return={keys.ROBOT_COMMAND: 'X', 'timestamp': 0.0}) session = endpoint.new_session() @@ -390,9 +390,10 @@ def test_empty_declaration_fails_before_motion(): def test_declared_stack_built_at_session_open(): """The server-declared local stack runs in front of the connection.""" + clock = [1.0] policy, mock_ws = _mock_remote_policy(CHUNKED_STACK, infer_return=[{'a': 1, 'timestamp': 0.0}]) - session = policy.new_session(now=lambda: 1.0) - actions = session({keys.OBS_TIME_NS: int(1e9)}) + session = policy.new_session(now=lambda: clock[0]) + actions = session({keys.OBS_TIME_NS: 0}) assert actions == [{'a': 1, 'timestamp': 1.0}] @@ -517,7 +518,7 @@ def test_remote_policy_lifecycle(inference_server, mock_policy): assert meta['server.model_name'] == 'test_model' assert meta['type'] == 'remote' - obs = {'dataset': 'test', 'obs_time_ns': 0} + obs = {'dataset': 'test'} action = session(obs) # Single-dict server response is normalized to a 1-element list (Session contract) and # anchored to absolute time by the declared ChunkedSchedule. @@ -526,7 +527,7 @@ def test_remote_policy_lifecycle(inference_server, mock_policy): session.close() # New session - session2 = policy.new_session() + session2 = policy.new_session(now=lambda: 0.0) session2.close() @@ -534,7 +535,7 @@ def test_remote_session_meta(inference_server): """Session meta must include server metadata.""" host, port = inference_server policy = RemotePolicy(f'{host}:{port}') - session = policy.new_session() + session = policy.new_session(now=lambda: 0.0) meta = session.meta assert meta['type'] == 'remote' diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 4196ee09c..3915e3cec 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -260,8 +260,8 @@ def pipeline(): local_session = inline(pipeline()).new_session(now=lambda: clock[0]) - remote_actions = remote_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) - local_actions = local_session({keys.OBS_TIME_NS: int(100.0 * 1e9)}) + remote_actions = remote_session({keys.OBS_TIME_NS: 0}) + local_actions = local_session({keys.OBS_TIME_NS: 0}) assert remote_actions == local_actions # Three scripted actions plus the chunk-closing validity sentinel ActionTimestamp appends. assert local_actions == [ @@ -272,8 +272,9 @@ def pipeline(): ] # Both gate identically while the chunk plays out. - assert remote_session({keys.OBS_TIME_NS: int(100.15 * 1e9)}) is None - assert local_session({keys.OBS_TIME_NS: int(100.15 * 1e9)}) is None + clock[0] = 100.15 + assert remote_session({keys.OBS_TIME_NS: 0}) is None + assert local_session({keys.OBS_TIME_NS: 0}) is None remote_session.close() diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 07930815b..46dd2eb6d 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -27,8 +27,9 @@ class Session(ABC): **Return contract**: ``list[dict] | None``. ``None`` means "no new trajectory, keep executing the current one" (used by scheduling wrappers). - A list is a new trajectory, replacing whatever is playing. Single-action - returns must be wrapped into a 1-element list by the producer. + 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. """ @abstractmethod @@ -86,11 +87,8 @@ def new_session(self, context: dict[str, Any] | None = None, now: Now | None = N Args: context: Episode context (task description, eval metadata, etc.). - now: The policy stack's clock (seconds), supplied by the harness and passed down to every - wrapped session. During a session call it reads the trial instant the call's output takes - effect — the observation instant plus the inference charge paid so far — so a chunk - stamped at ``now()`` lands when its call is paid for. ``None`` where no runtime clock - exists (server-side, warmup). + now: The runtime clock (current time in seconds), supplied by the harness and passed down + to every wrapped session. ``None`` where no runtime clock exists (server-side, warmup). """ @property @@ -141,7 +139,7 @@ class PolicyWrapper: def wrap(self, policy: Policy) -> Policy: """Apply this wrapper to a policy. Default: wrap every session it creates via ``wrap_session``. - Composition happens at config time; the trial context and clock reach the wrapped + Composition happens at config time; the runtime clock reaches the wrapped sessions through ``new_session``. """ return _WrapperPolicy(policy, self) diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index ccb4c9403..dfd4cd364 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -1,5 +1,4 @@ -"""Unit tests for PolicyWrapper composition, ChunkedSchedule, TemporalStack, and the -policy-pipeline algebra.""" +"""Unit tests for PolicyWrapper composition, ChunkedSchedule, TemporalStack, and the policy-pipeline algebra.""" from typing import Any @@ -31,6 +30,19 @@ from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +class _FakeClock: + """Minimal clock stub for unit tests — caller sets ``t`` directly.""" + + def __init__(self, t: float = 0.0): + self.t = t + + def now(self) -> float: + return self.t + + def now_ns(self) -> int: + return int(self.t * 1e9) + + class _ConstSession(Session): def __init__(self, actions): self._actions = actions @@ -51,9 +63,6 @@ def new_session(self, context=None, now=None): return self._session -_ONE_ACTION = [{'v': 1, 'timestamp': 0.0}] - - def _obs(now_sec=0.0): return {keys.OBS_TIME_NS: int(now_sec * 1e9)} @@ -61,9 +70,11 @@ def _obs(now_sec=0.0): class TestChunkedSchedule: def test_first_call_runs_inference(self): # Relative timestamps: trajectory of duration 0.5s + clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 1.0) - result = session(_obs(1.0)) + policy = ChunkedSchedule().wrap(inner) + session = policy.new_session(now=clock.now) + result = session(_obs()) assert result is not None assert len(result) == 2 # Timestamps stamped to absolute by ChunkedSchedule. @@ -71,85 +82,85 @@ def test_first_call_runs_inference(self): assert result[1]['timestamp'] == 1.5 def test_returns_none_while_trajectory_active(self): - # Trajectory starts at 1.0, ends at 1.0+0.5=1.5. + # Trajectory starts at clock=1.0, ends at 1.0+0.5=1.5. + clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 1.0) - session(_obs(1.0)) - assert session(_obs(1.2)) is None - assert session(_obs(1.4)) is None + policy = ChunkedSchedule().wrap(inner) + session = policy.new_session(now=clock.now) + session(_obs()) + clock.t = 1.2 + assert session(_obs()) is None + clock.t = 1.4 + assert session(_obs()) is None def test_re_infers_after_trajectory_consumed(self): - clock = [1.0] + clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(now=lambda: clock[0]) - session(_obs(1.0)) # trajectory ends at 1.5 + session = ChunkedSchedule().wrap(inner).new_session(now=clock.now) + session(_obs(1.0)) # trajectory ends at clock=1.5 assert session(_obs(1.3)) is None - clock[0] = 1.6 + clock.t = 1.6 result = session(_obs(1.6)) assert result is not None assert inner._session.call_count == 2 def test_single_action_refires_immediately_after(self): - """Single action at ts=0 → trajectory_end = the anchor instant → next tick re-infers.""" - session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(now=lambda: 1.0) + """Single action at ts=0 → trajectory_end = now → next tick re-infers.""" + clock = _FakeClock(t=1.0) + policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + session = policy.new_session(now=clock.now) session(_obs(1.0)) - assert session(_obs(1.01)) is not None - - def test_anchors_at_now_not_at_the_observation(self): - """``now`` reads the instant the call's charge is paid; the chunk is stamped there, not at the obs.""" - session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session(now=lambda: 1.5) - result = session(_obs(1.0)) + clock.t = 1.01 + result = session(_obs(1.01)) assert result is not None - assert result[0]['timestamp'] == 1.5 def test_expiry_is_judged_at_the_observation_instant(self): """Whether the trajectory has run out is a question about the observation, not about ``now``.""" inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) - session = ChunkedSchedule().wrap(inner).new_session(now=lambda: 2.0) - session(_obs(1.0)) # anchored at now() = 2.0, so the trajectory ends at 2.5 + session = ChunkedSchedule().wrap(inner).new_session(now=_FakeClock(t=2.0).now) + session(_obs(1.0)) # anchored at now()=2.0, so the trajectory ends at 2.5 assert session(_obs(2.4)) is None assert session(_obs(2.6)) is not None - def test_running_inference_without_a_clock_raises(self): - session = ChunkedSchedule().wrap(_ConstPolicy(_ONE_ACTION)).new_session() - with pytest.raises(ValueError, match='clock'): - session(_obs(1.0)) - class TestPipelineComposition: """Test | operator across PolicyWrapper and Codec types.""" def test_wrapper_pipe_wrapper(self): + clock = _FakeClock(t=1.0) pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy(_ONE_ACTION)) - session = policy.new_session(now=lambda: 1.0) + policy = pipeline.wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + session = policy.new_session(now=clock.now) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None assert result[0]['v'] == 1 def test_codec_pipe_wrapper(self): + clock = _FakeClock(t=1.0) codec = ActionTimestamp(fps=10.0) pipeline = codec | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) policy = pipeline.wrap(_ConstPolicy([{'action': 'test', 'timestamp': 0.0}])) - session = policy.new_session(now=lambda: 1.0) + session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None def test_full_pipeline(self): + clock = _FakeClock(t=1.0) codec = ActionTimestamp(fps=10.0) pipeline = ChunkedSchedule() | codec assert isinstance(pipeline, PolicyWrapper) # 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 (obs at 1.0). + # → ChunkedSchedule shifts to 1.0, 1.1, 1.2, 1.3, 1.4 (clock=1.0). policy = pipeline.wrap(_ConstPolicy([{'action': f'a{i}'} for i in range(5)])) - session = policy.new_session(now=lambda: 1.0) - result = session(_obs(1.0)) + session = policy.new_session(now=clock.now) + result = session(_obs()) assert result is not None assert result[0]['timestamp'] == 1.0 # Second call within trajectory window returns None (ChunkedSchedule). - assert session(_obs(1.2)) is None + clock.t = 1.2 + assert session(_obs()) is None def test_codec_and_stays_codec_only(self): """& only works between codecs, not wrappers.""" @@ -233,17 +244,19 @@ class TestTemporalStack: OFFSETS = (-0.2, -0.1, 0.0) def test_pad_start_repeats_oldest(self): + clock = _FakeClock(t=0.0) inner = _CapturePolicy() - session = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS).wrap(inner).new_session() + session = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS).wrap(inner).new_session(now=clock.now) session(_stack_obs(0.0, 1.0)) stack = inner.session.seen[0]['v'] assert stack.shape == (3, 1) assert (stack == 1.0).all() def test_no_pad_start_grows_from_one(self): + clock = _FakeClock(t=0.0) inner = _CapturePolicy() wrapper = TemporalStack(keys=('v',), offsets_sec=self.OFFSETS, pad_start=False) - session = wrapper.wrap(inner).new_session() + session = wrapper.wrap(inner).new_session(now=clock.now) session(_stack_obs(0.0, 1.0)) assert inner.session.seen[0]['v'].shape == (1, 1) @@ -260,9 +273,10 @@ def test_no_pad_start_full_window_matches_padded(self): offsets = self.OFFSETS stacks = {} for pad_start in (True, False): + clock = _FakeClock(t=0.0) inner = _CapturePolicy() wrapper = TemporalStack(keys=('v',), offsets_sec=offsets, pad_start=pad_start) - session = wrapper.wrap(inner).new_session() + session = wrapper.wrap(inner).new_session(now=clock.now) for i in range(4): session(_stack_obs(0.1 * i, float(i))) stacks[pad_start] = inner.session.seen[-1]['v'] @@ -448,20 +462,23 @@ def test_pipe_composes_no_further(self): _ = pipeline | spec.PolicySource(_ConstPolicy([])) def test_inline_full_pipe(self): + clock = _FakeClock(t=1.0) inner = _ConstPolicy([{'action': f'a{i}'} for i in range(5)]) policy = spec.inline(ChunkedSchedule() | spec.remote | ActionTimestamp(fps=10.0) | spec.PolicySource(inner)) assert isinstance(policy, Policy) - session = policy.new_session(now=lambda: 1.0) - result = session(_obs(1.0)) + session = policy.new_session(now=clock.now) + result = session(_obs()) assert result is not None assert result[0]['timestamp'] == 1.0 - assert session(_obs(1.2)) is None + clock.t = 1.2 + assert session(_obs()) is None def test_inline_tolerates_marker_less_pipe(self): - inner = _ConstPolicy(_ONE_ACTION) + clock = _FakeClock(t=1.0) + inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}]) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) - session = policy.new_session(now=lambda: 1.0) - result = session(_obs(1.0)) + session = policy.new_session(now=clock.now) + result = session(_obs()) assert result is not None and result[0]['timestamp'] == 1.0 def test_inline_bare_source_pipe_is_the_loaded_policy(self): diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index d97b8ae43..b54810a74 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -2,8 +2,8 @@ Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only -``ChunkedSchedule`` needs the clock — it anchors a chunk at the instant its call's charge is paid, -which the observation stamp cannot give — so the harness passes ``now`` (a ``Callable[[], float]`` +``ChunkedSchedule`` needs the live clock — it anchors a chunk to inference *completion*, which the +pre-inference observation stamp cannot give — so the harness passes ``now`` (a ``Callable[[], float]`` in seconds) to ``new_session`` and it reaches that one session. """ @@ -23,9 +23,9 @@ class ChunkedSchedule(PolicyWrapper): """Wait for the current trajectory to finish before calling the inner policy again. Owns relative→absolute time conversion: inner layers (codecs, models) emit relative timestamps; - this wrapper anchors them to ``now()`` after the inner call — the instant the call's charge is - paid. Returns ``None`` ("keep executing the current trajectory") until the last action's timestamp - is reached at the observation instant, then calls the inner policy. + this wrapper anchors them to ``now()`` *after* inner inference returns, so execution aligns to + inference-finish (not inference-start). Returns ``None`` ("keep executing the current trajectory") + until the last action's timestamp is reached, then calls the inner policy. """ class _Session(DelegatingSession): @@ -37,14 +37,14 @@ def __init__(self, inner: Session, now: Now | None): self._trajectory_end: float | None = None def __call__(self, obs): - if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: - return None if self._now is None: raise ValueError( 'ChunkedSchedule needs a clock to run inference: pass now (a callable returning seconds) to ' 'new_session. The harness supplies it; a direct RemotePolicy.new_session() outside the harness ' 'must too.' ) + if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: + return None result = self._inner(obs) if result is not None: # A single-action session may return a bare dict, and a no-codec path may omit @@ -52,9 +52,10 @@ def __call__(self, obs): # immediate action executes instead of raising. if isinstance(result, dict): result = [result] - anchor = self._now() + # Anchor to post-inference time so execution starts when inference *finished*. # Copy dicts so we don't mutate caller-owned data (sessions may reuse templates). - result = [{**r, keys.ACTION_TIMESTAMP: anchor + r.get(keys.ACTION_TIMESTAMP, 0.0)} for r in result] + now = self._now() + result = [{**r, keys.ACTION_TIMESTAMP: now + 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 @@ -156,7 +157,7 @@ def __init__(self, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_st 'in-range targets and the stack would be empty' ) - def wrap_session(self, inner: Session, context, now): + def wrap_session(self, inner: Session, context, now: Now | None): return TemporalStack._Session(inner, self._keys, self._offsets_sec, self._pad_start) def to_spec(self): From 4c6f07bf5ffdd4f303ed66a2be410702a5ca5f86 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 22:59:05 +0300 Subject: [PATCH 07/64] Keep `Session` typed as on main; accept the read-only view at the submit site --- positronic/offboard/tests/test_remote_policy.py | 4 +--- positronic/policy/base.py | 4 ++-- positronic/policy/harness.py | 3 ++- positronic/policy/remote.py | 6 +++--- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 65af1fbb5..6a3b2c417 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -57,9 +57,7 @@ class TestPrepareObs: def test_images_pass_through_untouched_by_default(self): session = RemoteSession(_mock_ws_session()) obs = {'cam': _make_image(480, 640), 'state': np.array([1.0])} - prepared = session._prepare_obs(obs) - assert prepared.keys() == obs.keys() - assert all(prepared[key] is value for key, value in obs.items()) # the same arrays, not re-encoded + assert session._prepare_obs(obs) is obs def test_compression_reaches_nested_images(self): session = RemoteSession(_mock_ws_session(), compress_images=True) diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 46dd2eb6d..56fbc36d6 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Callable, Mapping +from collections.abc import Callable from typing import Any Now = Callable[[], float] @@ -33,7 +33,7 @@ class Session(ABC): """ @abstractmethod - def __call__(self, obs: Mapping[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: """Predict actions for the given observation.""" @property diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index ac8ce746e..0ec2b7b52 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -456,7 +456,8 @@ def _step(self, clock: pimm.Clock) -> None: self._deadline = clock.now() + self._task.timeout self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() - self._future = self._executor.submit(session, frozen_view(obs)) + # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. + self._future = self._executor.submit(session, frozen_view(obs)) # pyright: ignore[reportArgumentType] # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: time.sleep(0) diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index cb5174436..142539a0d 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -26,9 +26,9 @@ def __init__(self, ws_session: InferenceSession, compress_images: bool = False): self._session = ws_session self._compress_images = compress_images - def _prepare_obs(self, obs: cabc.Mapping[str, Any]) -> dict[str, Any]: + def _prepare_obs(self, obs: dict[str, Any]) -> dict[str, Any]: if not self._compress_images: - return dict(obs) + return obs return {key: self._prepare_value(key, value) for key, value in obs.items()} def _prepare_value(self, key: str, value: Any) -> Any: @@ -42,7 +42,7 @@ def _prepare_value(self, key: str, value: Any) -> Any: return type(value)(self._prepare_value(key, v) for v in value) return value - def __call__(self, obs: cabc.Mapping[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: """Forwards the observation to the remote server and returns the action trajectory. Single-action server responses are wrapped into a 1-element list to honor From 0fd8f233b54bc18203c2622c51d2a019b03528b0 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 23:22:11 +0300 Subject: [PATCH 08/64] Let a cancelled episode's failed inference go, and name the clock for what it computes --- positronic/policy/harness.py | 11 +++++-- positronic/policy/tests/test_harness.py | 43 +++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 3 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 0ec2b7b52..bf79397a3 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,4 +1,5 @@ import concurrent.futures +import logging import time from collections import deque from collections.abc import Generator, Iterable, Iterator @@ -304,7 +305,11 @@ def _cancel_session(self) -> None: player.set([]) if self._future is not None: future, self._future = self._future, None - future.result() # nothing may close or re-enter the session while the worker is still inside it + concurrent.futures.wait([future]) # nothing may close or re-enter the session while the worker is inside + # rules-allow: swallowed-error — the cancelled call's failure must not pre-empt the stop and the home + # this cancel runs before; a live episode's failure still surfaces from ``_take``. + if (exc := future.exception()) is not None: + logging.warning(f'Inference failed on the call this episode cancelled: {exc}') if self._policy_session is not None: self._policy_session.cancel() @@ -351,7 +356,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._policy_session = self.policy.new_session(self.context, self._policy_now) + self._policy_session = self.policy.new_session(self.context, self._effect_time) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -425,7 +430,7 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it return inputs - def _policy_now(self) -> float: + def _effect_time(self) -> float: """The trial instant the in-flight call's output takes effect: its observation instant plus the charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; the loop thread writes the call's start fields before submitting it. diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index c5992c35c..09f7cc4af 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1684,6 +1684,49 @@ def test_harness_keeps_playing_while_a_call_is_in_flight(world): assert len(during) >= 3, f'the harness stopped playing during inference: {[t for t, _ in played]}' +@pytest.mark.timeout(20.0) +def test_finish_during_a_failing_call_still_stops_and_homes(world): + """A FINISH arriving while the call it cancels has failed still commits the recording and homes: the + failure is discarded with the schedule rather than unwinding the episode's end.""" + + class _FailingSession(Session): + """Fails after the directive lands, so the harness meets the failure only when it reaps the worker.""" + + def __call__(self, obs): + time.sleep(0.05) + raise RuntimeError('inference boom') + + class _FailingPolicy(Policy): + def new_session(self, context=None, now=None): + return _FailingSession() + + harness = Harness(_FailingPolicy(), make_embodiment(simulated=True)) + cmd_recorder = RecordingEmitter() + ds_recorder = RecordingEmitter() + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands['target_grip']._bind(RecordingEmitter()) + harness.ds_command._bind(ds_recorder) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + # The charge holds the failed call in flight, so FINISH is what reaps it. + (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=0.3)), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver, _Pacer()]), steps=2000) + + stops = [data for _, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] + assert len(stops) == 1 + assert isinstance(cmd_recorder.emitted[-1][1], Reset) + + @pytest.mark.timeout(3.0) def test_installed_trajectory_clears_the_channels_it_omits(world): """A trajectory naming only one channel replaces the whole schedule: the omitted channel stops being From 38c2cc1f235e12c9f4b0c5156037a36b9278e78a Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 23:29:50 +0300 Subject: [PATCH 09/64] Keep `recording.py` as on main bar the clock the wrapper now reads --- positronic/policy/recording.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/positronic/policy/recording.py b/positronic/policy/recording.py index 5a094f796..99cfdc4ea 100644 --- a/positronic/policy/recording.py +++ b/positronic/policy/recording.py @@ -35,8 +35,8 @@ Every field is *also* logged as ``rr.Scalars`` on a dedicated ``action_time`` timeline (each action stamped at the inference-request time plus its horizon offset), so a ``TimeSeriesView`` reads commanded values with real axes. That anchor is the pre-inference -``obs_time_ns``, so it precedes true execution — which the scheduling wrapper anchors at the instant -its call's charge is paid — by the inference latency. Select ``action_time`` to see them. +``obs_time_ns``, so it precedes the harness's true execution time (stamped after +inference by ``ChunkedSchedule``) by the inference latency. Select ``action_time`` to see them. Entity paths are ``{tap_name}/{data_key}``. A tap's incoming observation keys and outgoing action keys share that namespace; in the rare case the same key appears on @@ -221,8 +221,9 @@ def _log_action_series(path: str, arr: np.ndarray, horizon: np.ndarray, base_ns: Each action is stamped at ``base_ns + horizon_i``, where ``base_ns`` is the inference-request time; successive chunks lay out along one clock so a ``TimeSeriesView`` - has real axes. This precedes true execution by the inference latency: the scheduling wrapper - anchors commands at the instant its call's charge is paid, which a recorder tap cannot observe. + has real axes. This precedes true execution by the inference latency: the harness's + ``ChunkedSchedule`` anchors commands at ``now()`` *after* inference, which a recorder + tap sitting inside it cannot observe. """ arr = np.asarray(arr, dtype=np.float64) if arr.ndim == 1: @@ -305,7 +306,6 @@ def __init__( self._stream: rr.RecordingStream | None = None self._rrd_path: Path | None = None self._live = 0 - # Shared across a pipeline's taps and unlocked: the harness keeps one session call in flight. self._depth = 0 self._timeline_values: dict[str, Any] = {} self._image_paths: list[str] = [] From 842f3eff8b496c62c05eed71c9b3234427a5b86e Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 23:35:52 +0300 Subject: [PATCH 10/64] Call the fixed charge constant, arm the effect clock at episode start, drop the player TODO --- ARCHITECTURE.md | 4 ++-- positronic/keys.py | 6 +++--- positronic/policy/harness.py | 12 ++++++------ positronic/policy/tests/test_harness.py | 4 ++-- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 25514e6d3..689ee5e7c 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -133,11 +133,11 @@ the wrapper stack around the policy; a session returning `None` means "keep exec trajectory". **Inference cost is a fact of the trial, owned by the harness.** The trial context carries -`inference_latency` — a fixed charge in seconds, or the call's own wall duration — and the harness +`inference_latency` — a constant charge in seconds, or the call's own wall duration — and the harness alone reads it: it withholds a returned trajectory, and the world clock, until the charge is paid, and the clock it hands the policy stack (`now`) reads the instant the in-flight call's output takes effect. A scheduling wrapper stamps its chunk at `now()` and never learns the mode, so the charge -is deterministic wherever a fixed one is asked for. +is deterministic wherever a constant one is asked for. **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every diff --git a/positronic/keys.py b/positronic/keys.py index 66305bd34..2ab4c700a 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -89,7 +89,7 @@ def is_robot_command(name: str) -> bool: EVAL_SUCCESS = 'eval.success' EVAL_TERMINATED = 'eval.terminated' -# The inference cost each model call charges the world clock: seconds (fixed, reproducible), or ``True`` -# for the call's own wall duration. A sim trial without the key charges nothing (the world holds still per -# call); hardware always pays wall. +# The inference cost each model call charges the world clock: a constant number of seconds (reproducible), +# or ``True`` for the call's own wall duration. A sim trial without the key charges nothing (the world holds +# still per call); hardware always pays wall. INFERENCE_LATENCY = 'inference_latency' diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index bf79397a3..80e981a25 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -52,12 +52,8 @@ def next_due(self) -> int | None: return self._pending[0][0] if self._pending else None def advance(self, current_time: int): - """The single value due at ``current_time``, or ``None`` when no waypoint has come due since the - last call. - - TODO: several waypoints due at once collapse to the last, which drops the motion of every delta but - the final one. Revisit if pacing turns out not to hold one waypoint due per round. - """ + """The single value due at ``current_time`` — the last, when several came due since the previous + call — or ``None`` when none did.""" value = None while self._pending and self._pending[0][0] <= current_time: value = self._pending.popleft()[1] @@ -356,6 +352,10 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} + # Arm the clock before handing it out: a session reading it before its first call must see this + # episode's start, not the release time of the last episode's final call. + self._t0_ns = clock.now_ns() + self._wall_t0 = time.monotonic() self._policy_session = self.policy.new_session(self.context, self._effect_time) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 09f7cc4af..1164350a0 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1653,8 +1653,8 @@ def test_default_latency_pauses_the_world_for_the_call(world): @pytest.mark.timeout(20.0) @pytest.mark.parametrize('wall_sec', [0.0, 0.05]) -def test_declared_latency_ignores_what_the_call_really_took(world, wall_sec): - """The reproducible mode: the wrapper is released a fixed delay after the call started, so the played +def test_constant_latency_ignores_what_the_call_really_took(world, wall_sec): + """The reproducible mode: the wrapper is released a constant delay after the call started, so the played trace is the same against a fast server and a slow one.""" played = _run_sim_episode(world, SlowPolicy(wall_sec=wall_sec), ChunkedSchedule(), latency=0.3) From 6047f6df32ce8523ce50bd3a3d1e90ac92201f8b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Tue, 11 Aug 2026 23:40:34 +0300 Subject: [PATCH 11/64] Charge a real rig the wall time its calls take, whatever the trial asks --- positronic/policy/harness.py | 9 +++++--- positronic/policy/tests/test_harness.py | 30 ++++++++++++++++++------- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 80e981a25..c0dac4c1c 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -337,9 +337,12 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self.context = dict(context) if self._embodiment.simulated: # A sim trial that doesn't ask for latency simulation runs free of it: the world holds still for - # every model call. Hardware (no key) pays what the call really takes. - self.context.setdefault(keys.INFERENCE_LATENCY, False) - latency = self.context.get(keys.INFERENCE_LATENCY, True) + # every model call. + latency = self.context.setdefault(keys.INFERENCE_LATENCY, False) + else: + # The charge is a device for simulating a trial, so a real rig ignores it and pays the wall time + # its calls really take. + latency = True self._charge = None if latency is True else float(latency) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 1164350a0..be64c992d 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1618,9 +1618,12 @@ def emit(self, data, ts: int = -1): self.emitted.append((self._clock.now(), data)) -def _run_sim_episode(world, policy, wrapper, *, latency, steps=4000, run_sec=1.5) -> list[tuple[float, Any]]: - """One sim trial under ``latency``; returns the grip commands with the world time each went out at.""" - harness = Harness(wrapper.wrap(policy), make_embodiment(simulated=True)) +def _run_episode( + world, policy, wrapper, *, latency, simulated=True, steps=4000, run_sec=1.5 +) -> list[tuple[float, Any]]: + """One trial whose context asks for ``latency``; returns the grip commands with the world time each went + out at. A sim trial runs against a pacer, the sole time-master a real rig doesn't need.""" + harness = Harness(wrapper.wrap(policy), make_embodiment(simulated=simulated)) grip_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) harness.commands['target_grip']._bind(grip_recorder) @@ -1637,7 +1640,8 @@ def _run_sim_episode(world, policy, wrapper, *, latency, steps=4000, run_sec=1.5 (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), (None, run_sec), ]) - drive_scheduler(world.start([harness, driver, _Pacer()]), steps=steps) + systems = [harness, driver, _Pacer()] if simulated else [harness, driver] + drive_scheduler(world.start(systems), steps=steps) return grip_recorder.emitted[1:] # drop the startup home @@ -1645,7 +1649,7 @@ def _run_sim_episode(world, policy, wrapper, *, latency, steps=4000, run_sec=1.5 def test_default_latency_pauses_the_world_for_the_call(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 call really took.""" - played = _run_sim_episode(world, SlowPolicy(wall_sec=0.05), ChunkedSchedule(), latency=False) + played = _run_episode(world, SlowPolicy(wall_sec=0.05), ChunkedSchedule(), latency=False) assert played, 'no command was played' assert played[0][0] < 0.01, f'the world advanced during the call: first command at {played[0][0]}s' @@ -1656,7 +1660,7 @@ def test_default_latency_pauses_the_world_for_the_call(world): def test_constant_latency_ignores_what_the_call_really_took(world, wall_sec): """The reproducible mode: the wrapper is released a constant delay after the call started, so the played trace is the same against a fast server and a slow one.""" - played = _run_sim_episode(world, SlowPolicy(wall_sec=wall_sec), ChunkedSchedule(), latency=0.3) + played = _run_episode(world, SlowPolicy(wall_sec=wall_sec), ChunkedSchedule(), latency=0.3) assert played, 'no command was played' assert played[0][0] == pytest.approx(0.3, abs=0.02), f'first command at {played[0][0]}s, expected the 0.3s delay' @@ -1666,17 +1670,27 @@ def test_constant_latency_ignores_what_the_call_really_took(world, wall_sec): def test_measured_latency_charges_the_calls_own_wall_duration(world): """``inference_latency=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.""" - played = _run_sim_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=True) + played = _run_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=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 call took' +@pytest.mark.timeout(20.0) +def test_a_real_rig_ignores_the_latency_a_trial_asks_for(world): + """The charge simulates a trial; a real rig pays what its calls take, so a context carrying a constant + (the eval CLI writes one into every trial) does not hold the chunk back for it.""" + played = _run_episode(world, SlowPolicy(), ChunkedSchedule(), latency=5.0, simulated=False) + + assert played, 'no command was played' + assert played[0][0] < 1.0, f"first command at {played[0][0]}s: the trial's 5s charge was honoured" + + @pytest.mark.timeout(20.0) def test_harness_keeps_playing_while_a_call_is_in_flight(world): """A wrapper that replans before its chunk is exhausted leaves waypoints due during inference, and the harness emits them on time instead of standing still until the model answers.""" - played = _run_sim_episode(world, SlowPolicy(span_sec=0.4, steps=20), _ReplanEarly(), latency=0.15) + played = _run_episode(world, SlowPolicy(span_sec=0.4, steps=20), _ReplanEarly(), latency=0.15) # The second call starts halfway through the first chunk (0.2s in) and is owed 0.15s; the waypoints due # in that window have to keep going out. From a04cbe39d4bbc99ff76b43e4c40c40261e09035c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 00:47:05 +0300 Subject: [PATCH 12/64] Let the policy stop a faulted arm, and keep a constant charge off the machine clock --- .basedpyright/baseline.json | 16 -- docs/connect-your-model.md | 5 +- positronic/cfg/wrappers.py | 4 +- positronic/keys.py | 4 + positronic/policy/harness.py | 107 ++++++---- positronic/policy/spec.py | 3 +- .../policy/tests/golden_pipeline.json.gz | Bin 7893 -> 7883 bytes .../policy/tests/test_golden_pipeline.py | 13 +- positronic/policy/tests/test_harness.py | 192 ++++++++++++++---- positronic/policy/tests/test_wrappers.py | 34 +++- positronic/policy/wrappers.py | 26 ++- positronic/vendors/gr00t/server.py | 4 +- positronic/vendors/lerobot/server.py | 4 +- positronic/vendors/lerobot_0_3_3/policy.py | 4 +- positronic/vendors/lerobot_0_3_3/server.py | 4 +- positronic/vendors/molmoact2/server.py | 4 +- positronic/vendors/openpi/server.py | 4 +- 17 files changed, 314 insertions(+), 114 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 100110818..ecfd23cd5 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5835,14 +5835,6 @@ "lineCount": 1 } }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 35, - "endColumn": 40, - "lineCount": 1 - } - }, { "code": "reportArgumentType", "range": { @@ -5850,14 +5842,6 @@ "endColumn": 45, "lineCount": 1 } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 39, - "endColumn": 44, - "lineCount": 1 - } } ], "./positronic/policy/tests/test_policy_io.py": [ diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 5954d942e..34af81bd1 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -101,6 +101,7 @@ The client sends the full raw robot state as a dict. Keys are flat strings (the | `robot_state.ee_pose` | float32 | (7,) | End-effector pose: `x, y, z, qw, qx, qy, qz` (quaternion is **wxyz**, scalar first) | | `robot_state.q` | float32 | (7,) | Joint positions (radians) | | `robot_state.dq` | float32 | (7,) | Joint velocities (radians/s) | +| `robot_state.fault` | bool | scalar | Whether the arm is faulted. A faulted arm sends no `robot_state.ee_pose`/`q`/`dq` at all — it has no sample to give — so a stack that reads them needs `StopOnFault` ahead of it | | `grip` | float32 | scalar | Gripper closure in `[0, 1]`: 0 = open, 1 = closed | | `image.` | uint8 | (H, W, 3) | Camera RGB. Every eval target — PhAIL and each sim — sends `image.exterior` and `image.wrist`, whatever the underlying benchmark calls those cameras, so one codec reads them all; a target with more views adds its own names beside them (the MuJoCo sim adds `image.agent_view`) | | `obs_time_ns` | int | scalar | Harness-clock timestamp of this observation (ns) | @@ -206,14 +207,14 @@ 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 `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 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. 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 scheduler in particular, and on a real rig it wants `StopOnFault` outside that scheduler — a faulted arm is not tracking the plan it was given, and the wrapper answers the empty trajectory so the rig stops rather than resuming a chunk stamped before the fault. 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. `new_session`'s `now` argument is the runtime clock that wrappers scheduling against live time read; a policy that does no scheduling of its own just accepts and ignores it (server-side it is `None`). 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: ```python -ChunkedSchedule() | RestrictImageSize() | remote | codec | source +StopOnFault() | ChunkedSchedule() | 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. diff --git a/positronic/cfg/wrappers.py b/positronic/cfg/wrappers.py index d8c2300e4..6f2c0f6aa 100644 --- a/positronic/cfg/wrappers.py +++ b/positronic/cfg/wrappers.py @@ -1,7 +1,7 @@ import configuronic as cfn from positronic import keys as obs_keys -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack chunked_schedule = cfn.Config(ChunkedSchedule) temporal_stack = cfn.Config(TemporalStack) @@ -47,4 +47,4 @@ def video_context_wrappers(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 stack | ChunkedSchedule() + return StopOnFault() | stack | ChunkedSchedule() diff --git a/positronic/keys.py b/positronic/keys.py index 2ab4c700a..93e795e7e 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -29,6 +29,10 @@ def is_robot_command(name: str) -> bool: JOINTS = 'robot_state.q' JOINT_VEL = 'robot_state.dq' EE_POSE = 'robot_state.ee_pose' +# Whether the arm is faulted, in every observation the harness builds. A faulted arm has no sample to give +# and is not tracking the plan it was handed, so the policy stack — not the harness — decides what happens +# next: the rest of ``robot_state`` is absent from that observation. +ROBOT_FAULT = 'robot_state.fault' GRIP = 'grip' TASK = 'task' # The prefix that identifies a camera on the wire: an embodiment declares its cameras by naming diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index c0dac4c1c..bbf3554f4 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -14,6 +14,7 @@ from positronic import keys, telemetry, telemetry_keys from positronic.dataset.ds_writer_agent import DsWriterCommand from positronic.dataset.serializers import expand_suffixed +from positronic.drivers import roboarm from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Task from positronic.policy.base import Policy, Session @@ -53,7 +54,12 @@ def next_due(self) -> int | None: def advance(self, current_time: int): """The single value due at ``current_time`` — the last, when several came due since the previous - call — or ``None`` when none did.""" + call — or ``None`` when none did. + + Collapsing to the last is exact for an absolute setpoint and lossy for a relative one: a run of + deltas due together arrives as its final step alone. Pacing keeps one waypoint due per round + wherever a round is shorter than the spacing between waypoints. + """ value = None while self._pending and self._pending[0][0] <= current_time: value = self._pending.popleft()[1] @@ -217,8 +223,10 @@ def __init__( self._policy_session: Session | None = None # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. self._running = False - # One session call at a time, on a worker so the harness keeps playing while the model runs. - self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') + # One session call at a time, on a worker so the harness keeps playing while the model runs. The + # worker belongs to the episode: ending one abandons the call in flight rather than waiting for it, + # so the next episode must not queue behind it. + self._executor: ThreadPoolExecutor | None = None self._future: Future[list[dict[str, Any]] | None] | None = None # The in-flight call's start: the world instant its observation was built, and the wall instant it # was submitted. @@ -293,21 +301,33 @@ def _pace(self, clock: pimm.Clock) -> pimm.Command: return pimm.Sleep(min(POLL_PERIOD_SEC, max(min(due) - clock.now_ns(), 1) / 1e9)) def _cancel_session(self) -> None: - """Drop everything the episode has going: the schedule being played, the call on the worker, and the - session's scheduling state so the next inference is not held back. Devices hold their last commanded - position — nothing is buffered downstream to clear. + """Drop everything the episode has going: the schedule being played, and the call on the worker. + + The call is let go of rather than waited for — a model that hangs must not hold up the recording's + stop or the home — so its worker is retired with it and whatever it eventually answers, or raises, + lands nowhere. Devices hold their last commanded position; nothing is buffered downstream to clear. """ for player in self._players.values(): player.set([]) + self._retire_worker() + + @staticmethod + def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: + """Report the failure of a call nobody is waiting for any more.""" + # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no + # caller left to raise to; the log is the only place its failure can go. + if not future.cancelled() and (exc := future.exception()) is not None: + logging.error(f'Inference failed after the episode that asked for it ended: {exc}') + + def _retire_worker(self) -> None: + """Let go of this episode's worker and the call it is running: the answer lands nowhere and the + failure only reaches the log.""" if self._future is not None: - future, self._future = self._future, None - concurrent.futures.wait([future]) # nothing may close or re-enter the session while the worker is inside - # rules-allow: swallowed-error — the cancelled call's failure must not pre-empt the stop and the home - # this cancel runs before; a live episode's failure still surfaces from ``_take``. - if (exc := future.exception()) is not None: - logging.warning(f'Inference failed on the call this episode cancelled: {exc}') - if self._policy_session is not None: - self._policy_session.cancel() + self._future.add_done_callback(self._report_abandoned) + self._future = None + if self._executor is not None: + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None def _finalize_recording( self, clock: pimm.Clock, payload: dict[str, Any] | None = None @@ -359,6 +379,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: # episode's start, not the release time of the last episode's final call. self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') self._policy_session = self.policy.new_session(self.context, self._effect_time) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None @@ -400,17 +421,27 @@ def _handle_directive(self, directive: Directive, clock: pimm.Clock) -> Generato case _: raise ValueError(f'Unknown directive type: {directive.type}') + @staticmethod + def _is_faulted(value: Any) -> bool: + """Whether a raw observation is an arm reporting a fault. Every other not-ready sample is simply absent.""" + return isinstance(value, roboarm.State) and value.status is roboarm.RobotStatus.ERROR + def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: """Read every observation channel and assemble the policy input dict. Raises ``NoValueException`` if any channel has no value yet. Returns ``None`` while a serializer reports a sample is not ready (``robot_state`` during a ``RESETTING`` arm) or a channel still holds a pre-reset value — either way the harness skips inference rather than feed a partial or stale obs. + + A faulted arm is the exception: it has no sample either, but the plan being played was made for an + arm that is now somewhere else, so the observation goes to the policy stack carrying + ``keys.ROBOT_FAULT`` and without the arm's own entries. """ # Against the live model, not the one known at episode start: a remote env publishes its ``robot_meta`` # a turn after the reset that produced it, so at episode start there is no model to check. assert_default_frame(self._statics()) inputs: dict[str, Any] = {} + faulted = False for name, obs in self._embodiment.observations.items(): message = self.observations[name].read() if message is None: @@ -421,12 +452,16 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: if obs.serializer is not None: value = obs.serializer(value) if value is None: - return None + if not self._is_faulted(message.data): + return None + faulted = True + continue for full_name, v in expand_suffixed(name, value): if v is not None: inputs[full_name] = v if self._awaiting_obs: return None + inputs[keys.ROBOT_FAULT] = faulted inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() inputs.update(self.context) @@ -447,8 +482,8 @@ def _step(self, clock: pimm.Clock) -> None: The call goes to the worker so the harness keeps playing while the model runs; a wrapper that answers without inference still resolves in the round it was asked. """ - session = self._policy_session - assert session is not None # only a live episode steps + session, executor = self._policy_session, self._executor + assert session is not None and executor is not None # only a live episode steps if self._future is not None and not self._take(self._future, clock): return obs = self._build_obs(clock) @@ -465,39 +500,40 @@ def _step(self, clock: pimm.Clock) -> None: self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. - self._future = self._executor.submit(session, frozen_view(obs)) # pyright: ignore[reportArgumentType] - # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. - while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: - time.sleep(0) + self._future = executor.submit(session, frozen_view(obs)) # pyright: ignore[reportArgumentType] + if self._charge is None: + # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. + while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: + time.sleep(0) self._take(self._future, clock) def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> bool: """Install the call's trajectory once the world has paid for it; True once the future is consumed. - A skip (``None``) costs nothing and is consumed on sight. A model call's trajectory is stamped for - ``t0`` plus the trial's charge, so it is withheld until the world clock reaches that instant — and - the world is withheld from running past it: blocking here blocks the loop thread, which is what - advances a virtual clock. + Under a constant charge the world holds still until the call answers — blocking here blocks the loop + thread, which is what advances a virtual clock. Until a call answers there is no telling a skip from + a model call, so letting the world run meanwhile would spend trial time on whichever the machine + turned out to be slow at. A skip then costs nothing; a trajectory is stamped for ``t0`` plus the + charge and withheld until the world reaches that instant, playing what is already scheduled on the + way. A charge measured in wall time can hold nothing still, so there the world runs no further ahead + of the call's start than wall time has. """ - if future.done() and future.result() is None: - self._future = None - return True if self._charge is not None: - # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP - # and slips the install a full round. - if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): - return False concurrent.futures.wait([future]) elif not future.done(): - # The world may run no further ahead of the call's start than wall time has. ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) if ahead <= 0.0: return False concurrent.futures.wait([future], timeout=ahead) if not future.done(): return False - self._future = None actions = future.result() # taken on the loop thread, so a failing call still seals the episode + if actions is not None and self._charge is not None: + # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP + # and slips the install a full round. + if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): + return False # the schedule already playing carries the world to the release instant + self._future = None if actions is not None: self._install(actions, clock) return True @@ -561,8 +597,7 @@ def _shutdown(self) -> None: The harness does not own the policy's lifetime: the caller may run several harnesses over one policy (a multi-eval sweep), so it closes the policy once, after the last run. """ - self._future = None - self._executor.shutdown(wait=True, cancel_futures=True) + self._retire_worker() if self._policy_session is not None: self._policy_session.close() self._policy_session = None diff --git a/positronic/policy/spec.py b/positronic/policy/spec.py index f98b65199..6e66a02f8 100644 --- a/positronic/policy/spec.py +++ b/positronic/policy/spec.py @@ -46,7 +46,7 @@ RestrictImageSize, ) from positronic.policy.observation import ObservationCodec -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack class RemoteMarker(PolicyWrapper): @@ -150,6 +150,7 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) # names the whole-table fix, and the test above catches a desync meanwhile WIRE_WRAPPERS: dict[str, type[PolicyWrapper]] = { 'chunked_schedule': ChunkedSchedule, + 'stop_on_fault': StopOnFault, 'temporal_stack': TemporalStack, 'action_timestamp': ActionTimestamp, 'action_horizon': ActionHorizon, diff --git a/positronic/policy/tests/golden_pipeline.json.gz b/positronic/policy/tests/golden_pipeline.json.gz index 0df3117bead5ec06c2b34a4a9dd0a606eefa8808..48769162b0ef26a025b114bc990c0458e96b2b1c 100644 GIT binary patch literal 7883 zcmeI1c|4SB`@r#(ipP3RD9ebFEO{o$ZaO7XIIkr+BFr&aB7}o%Gs>G(%+#?{s0puK zX-FE|$P`+XrTB7bM+>PYv1ja}_>bp3)b`toA1&XRjoM-%`>Vr6Ok{cD)f z^3-G}48vxATv;9`msm`!%+E}cUt`&fnPsHvEMlV2(D##AtKuL&2BE+d%flZJGo~uO zh22_?iAA0@P?#1=#Xk=a50oSho#uYUP-G+U6l9T|*)0rLFrNh<-E zoRLBJ#rKMm54cU_t5-!tR-5KR=XdlvhzY{9!rK^DWTX%*5q&vIA-t{elK6wQhChA@ zBpFPxY#Zv4f-v>)M~v1rRHz)L>N0@7g5PdrJ7@QY68CWX1IB}T ziU3SCypZwWUzZc=Ea?5-@ygkymV$!MN-c#1he|C)1UEmh6f@Ir+HZB;<*tfElgr)x z5(4PE2P90&o-{CSVE6(mW{9$V)u@fRrf8GV&5H{}w5E)Xm6CeEfn^5b90j zC4gT<>M#l;O17D}m$NhEy#|G>Qo1S>z~P z-&I=I&kx9e{?MB`)D}Pv422kV{aXNe@G>M?kJ<_-fY%|YzJDv75;-6R?-xpfQ1w~c zfW2S_6j?7X0PF*Epse~V0YDWjhHC5O1p#%i9O|#n5|pM~_#`DVAXE>Dl2IUV0BnLF zaz6+h1m8lIWRwta2&6(6$o)cq4mb$iB%_4WDX0ml@d2UFP%T+r1TX|yP(L|K1VDhx z5Kn`=DDV@=50`Jq5(SQe!tlcl@?s7YOO{kKMlXfrE$(XIfQ#l+v_~1&#zczoD)O0Z zirp$D1}S{ipdUR-v6J=vPPco6ZGOM)!X0X8y{aQEhTO~Ni_{qOFp=qrA;|j8DB_|U zT5svtiz6)tJxpaBuMyP!X7>DPUIbMQnYsCU)BNFHKJtf=PYMGHz0Q6#|l#4{+g|685DF!SESSC0Mj z)4yYyJIX`uC5;evL;JtZ|AlF;heZUfh?yLO_dog$U3;4uZ6>pS(%(Gj4S}oqJ#H8f zzxvmTzWoGpLe}@fHS!+eWSujE-&G#tu3bU9@qT**><1-E>6O@ou#)3X8ZGHNjf_3@@h=neQWqj}dT-pE&oJyE zk&nv-I%11|8q(5}OL~#IRQ`MBKDFx<#o(Uh85P)p*WE2IyAe&^d(+~W&$W7Rb@t&} zzUl#*vx~1u!FjA>KfPoH>fZ`)c3d?zN}#h= zyrynE0x>V_^|Oy-naPy8-s0_uxO3{<=9)_%KAIhkE_rC;V1*9Ucu1x&vpWpSIF?q< zu%5_X4~UC+8F_j)JxSlpH!zVC(IUV9MDojjGxK^!vx0J~tg8y*IR3Y=$=#A=C-B`4!RRypvMp$8;WRA=WR^V0^bQ?~SCu?C zLAhmD^jc2aE|zruCuyw!P3?dcXx3WKH0 zYEzxx-COO^)8=S8xknAjN zGDKBNZaNc?p!-NNuxYFXL+ z-cW_0nOqm~QK)Qwfk;IMcH-Ux&Vp*+a-5A)Wt|j47Ip8meaRXN-VtHDV)&$)vj@fg z6mQSsalkDx3b`4L*cWF-@Z>FuDCK`rC$tT&A9eD2evcsv$EtKBOIhUIA)A1CHc%mba z#g6Q5+I2d@S>a$Yns|cf9^_UtFs2(ET2j7BTJ#*Yfn{6gKBgaA=4=ELFYYSpm#r3A0S|NF@9MBG^Dd?s) zdKA!}Iu*e$wG}y?rZ$f!-ReqZ=akWyLX9Z__B-0rK9kP}&$7~X>4i-?Vb$tdGMhIeRlnvJlg&`|XR2d~~t()%<;&MTpxoA-#w7WE;sr zQkG$rMy#B07;)D(g=u8bJ&GQ&!GF?gJ}Yw0zq&h~N()pQr`;UA z$Am9DwLI0St;U?q@W079xIIegK8X~qFeZUzo>xK|G~H1%`?!yDJ$`w`9&w(?SZ2|% z%dE+dA~5Xy99?8(dS*G(0)|*x>J+hRnmps1uOJ}p9uPR!cf`p~=!x9^wK&Bbq$pY?UAz4>3TI2q}z@!EeCk-wnsRs4=Sd9IH1F9_2NH}~A~Opy!Woz2%Da*6&+>gM8i zyD~S|Sfy?*e&0|zzGLZ;aPOgO)&hdWQcF;zQREge(*`DA9xivDcUO!0&uVqk|$?@ zfF}3~I!u-q0<^#m2u02k0<^(iD3UBM4CsQxr6C{-*DBZBFANxfQ_veSN(3+h7a#_? zUj#S|^1($LP@=#QPyh}!^os(N+6e&6Aaa1NErcBK5Hw8w}i~`QBp*2p2DuzT1dK{OLy++vMH?td; z)6klzLls9l^L>fUQ-qiPGp8Bfn`R?JxEFfx7H+?I@1N*>!xugiW5J_bVkHs`udU-P zk#H=C^*rWp!t_1NK_0xY{NLpC%Vr00tl`$a2I+6W{GVWYxb>;wlZy4M;QMsns%QTW z>AUc5c!rbsvd?iLR)cbOi=pRNIQwS~uOKinGrg1ie~Y?T$(){Ir{gBU_~wW?_^kpipaq6fj$Zl- zuuCZS-@sM>sr}6$qRcR%vom<-<*8Zc%-)qb^!&hNbWa$Zo)f%cg}HdK&%I1gm#8rl z9^?}?Q;~`3oRx+((gLp+{^DKJ+ZnXf_2J5Ore0d3lXrCqEwfV6Q;qCX*03*QkRvMv zgLo=)srvDcLpumh<1o8e$Fprrxl(d^#OdJu>dYrHg7*!8iGVJZkQPgJ2mCkpe)hgZ^yWf*tUhJu28V|pP=+&O9 zZrPYd=v&XC;vG6;ubdfQ?l5w9M%$aMI>7Y-kr-M{x-h5W^=LamuIH-Z8o3j#Rk>3A zb5WqG90%Sf2+yX`_tW31k;aE`8;)8i_xP$*s17b(ySeL>uVYS^y@RVtw)Lu1ANO?L zR)#)B`b?omTUXW0FJu%S;)chKu)sFdk3DyE#$ZmS-cpd03yZOau>hAj!v6Y%-!F=}MKL)D2m zRcw3W+TNNv6W_AqykgyY_)k`{gH~=P`-z2}*f5!B&t^4fmJME^m@P2Ut<_008SBHi zJ!VW2(#~7Jn1*9Lst>QUso8stnYdZytu5X@=J;{P6IC&1m;@mo_H-N$0SiEo8uA5G~(y`dv<6~5{zPW7`ip0k8qW9Ki+mBFx{(UtU zSic|%NBVBOfgJgoql+hKcUt3zwZhX#Da5yG8;41DU)<>%c|@mT7jemxpEuv0jb)7W z(A$o$pZ;IksGWJqgWgFGZ#O1Z40zecYfq=HFDUj%-Z-{6)HOZv(p%4RR(xG(&P>8->Aw4I)kXR1-~c!?O30UgRV_uq(x2rL>&h=RMO=xIiydi(d_AW;nr9 z`n_p+2+98a?~^{vw>ND{W;fg|@Fbe~#OmfI@*WF@xh>`=zb#%lseHwnBcHr}7_F9+ z)P_Ykt-!3kag=#F6zP?Zt2C9(&)IlI9myk!j$@9vPINqD zMiUNt1(f{U`!e)@uU&Hs&AwCnjAUxSWRH_Ullvqw86I+Hv3vb;sS=$m3Y?wQ%kZb| zhcY^>iauA;w9N84)?O0Uqv^Noj|~Ndyc{LY9y#ju$Tm%Tyd)U5esP&;_5@bYo!7I2 z1A*`m{e`if-(Do^g>Agm?k0>hwiS03z-$b&t6A|G&B~fmog%0aimi86%(3+vM@B=bx@ROyCy1*sh}bE|*d_omsp;qPdx zh2Bon??5=UHE#p!sjcezG*@E|PXA%zFU$68q;G`?y@b>h4(zCo|`zQ%2{=^iAIE8GQjV=fBQi>h5}gT=V%!!})D0nJ8)TO?OBmZ^%Q6*;L^LV;Qe++b zmPljGzGP?YV;P43jHTPX_kX+hIq&oMJ@0wOTBI|%gsqD@h(l5*$43ha zmKNVltSr~(rWB-2FOWN1KYK-TNb#@uFNi3;dP4C}H+&<-Pvsh}t#ptr7didv$)bOv z;aw^ITB&w*4!umh0wKG0-C7~L4&5#xyH4GSq$xcgHCU=8TBoQO>$!t@ALkh}!bK;c`GeUPWJSaaC>W%7JGZM zv8}yupl6Pf5qX)1IHgS%Vtepjxc$;62Dcu&ybv3(`-tv(#*a56j+g-wt$Q4&^qmEx z9{fh|vZG@7K-wg+7{v@1a!y|l?h*2}#^_~vr5zWBdhp(xfw-_7$M*|+|CBpQRx!7@ zeA3x}5>Z^i1n1uj@K(nI)D~bGR@r@CI1Xvq1z3+>VsrBVVQ<{(U3e3!9sR}I2{izCU%-hlwIONY&2`fI*jXuFyr5i2Gxl=FtG^d>2ZFs|s zQ_zs&z?9t{3cd~=vl2^2r6J>~isLv(*FDSlUigvK*hLf(8NZ@97q=waBYvVlk#c=| z2@g9~ysAKPm4yLhDwJWsG6GDXKw%*R76Pz?R|;hqvD*L`XjE9ph}{nGfk-BjjA$O? zHt1fM8rssG5ensk8KG0%8Mi|Zz#Pz%?u<;(gD@X7=)t%HdIT1VCU@6DWTv1eV2S8L zkKdS~!mvCv1NJundKy-ZmceG|F}WQyDk|KD-3OpSpQ2XxhGRE(m$83CTVfemp=V)3 z=u|8t8&nFmfS$xMvO|A^G3kI089AVrVK5!bhh)7%OMloTR2~NEe0}J?8+rpKqr>j$ zz6W{}rl6zX>An|w3x?Kt?CH)4Rfk#XeDw^_leq8iJqgu;`RIUNj9gHCSg4Mr7vp}Y zF)UUm)r*lEdJmSRGwH?11GR!x>VV!gda3sQbCXbeShvn&Z+Bj(6KqoFtG7EG>I!4f zW%qICgJNOqx(Yt-2ch0DKHbMYNqPcK{`yFs>dx>>T(!^e(S<&oFb?n#`ck2f8wVr* z7TK--On=R&42-5AR4cgurPz?s1s+WwJd4%r!7s9 z|8U^5Qeu$Ai(RN&y+-mpm>oa&TZudKnP^z_Z{7SD-U#uMWX=@_$p0HMt&k8}C1K{{LutZG@vn*WrX=w($D-ql%PzOW%?sT7jQf9UJHo>vIW> z-v1w%qeiJigzqfER#~MzyS(Lj{xxVuq*+#DBdxik*FT zJ@CQqyaOV~DjGf^?5(dw@_aZ}G5ZPOYJDw=M?|1vPZh#@%io5ng}WN9r^-ZIv|mhO z`*Y`|m5#3b)p~7Lhk3GiS8$3J(8Al0G+Nx&^*p7Dyks1wxv#)cg*0#H{e3=jLgZ|a z>&@xU2!{af5nJPkCrc7Zc>yj5zf`tY*R*SY#wZ`w)Qh!cfgj3<7#yf-xJEsF6MnhU zxi)gVQ|B-)Ct15WYI^?xw$FXnT+~0E+QV7)JuG)zn^5IlEh26sGT&_7=4+rQWpjQ0 zteRiSK$5=sfXiYwPPj4ucx+3)#j&Ro?HLn?rVH#$`Befd_kxksjQV%pE>HE+-RGV1@}Xmhx7G+4k*~i%3;tmVX4qHmdPB93I%(l1J(}-0~sH$$SrcT>vkEgx5 z=Qtd6Zl5T8@QL!f8|3OQW0~PesCS8D?*diRYa`jz{0U~iaf{i|R>4(l9syz&|Uc6Pwv_ z(~DA{&dom+H9e*Ml(rV-PvN6VfS@hw#NB2W!)7H^*at&9gMyARidu~B_+>UT1*bc? zLEg1FBUNi#Sc4JW(+sQ0VcCC>8u^KQdd|_Bws;moi(2NN=R8`q)D9JScV%bH);ZA@ z$na?!y@^w|A7E+#apCUfORrU`o`Oee!A@f=LxM!OBaE98!7U#^0B{dS|p`kf#U?~Hg6D(;6 zgM!L{h9!encBeR6ZhdaA#z3-1r*=gswx+7pCC86*_%GSAE{=q~h|@M0$?75)X^G9< ze9)IH^?O=Q^M_oo6C0<=wB^-^%ROU<*~kR(T)10?9W}&|K9#}8?asHpWyH6nIVg~` z<;{#S^E0Q{jN8fkMZP|Vzp>3!fV{)%LP4!gG z5+#OPf_8yL77OI_zuwar?0ktIzL0Ze{MHe5J+_gCmZ% zvPEUmmipQX=~Hv}mwTjRhvbG6lT@4p_5B}+v$)3Th0;!0|54oSawWouG6@xz<5MrZ zoLG|V5kk8UQ0od6ZD&p6eoADOL}KlOfP!aD)L6~hG1?923X{ShMH>t!l)QV(J#9z# zi#Do`RMLJkF5GtH`wu?GDGml)IY+QAnCwFyGm+S3aSB5*9yVM>mJ6e$RFV^}*a=6Ma zbH&wO!N|6X)I~m7@t-xkd!GM0S6FSOvuu;d`^YDQyfl(dYfJ0mJZGLZ$eC-NHo7YoRb!sUk=_f7Gi>bAdp4eig$h5l1%;oUuZ|c| zI`ppb^-F%=i^rYU7(3viux?LD$lFUR zPT_np4a#i@$pTJ)0_+Xa4tP)F_CF{*)+4xxZM2(kHCRb4*3(m0E!NA^S}oSw(?c!R zC+2|ubAiig3VWIL(-ihGd!{LHF~_7OvW+YfuFeX~IGo!_RUF1Cxi(Pn-{6fa0`n*W za_5TT6z-{O!!mvduCgky)PTWHa))#YP~T#8Sbd7P`@3RSp$xFCXg;?h2IyAU9yHpm zh!M&JJB*HUE7}4DV4~=5w;~9XJrw@Xj&Fab+ZOCj;3Sw;$N<3tfGGH_&<%oR1CM>Nq*b807%j3#fwVqNr_TK|Z~+EjR8cEpYax3& zy4$0O84AKm(IlCXpg3P$CNnUrh+#XH2e=P@D{|Y8S4wV$LVUYj>f~i!7ohBog5CWxz?AVI{JK{=d6Czlu{5_o2VB-2~ zf5Xp21@tB6zjO8vvRxk8}j!KUFs$cUFNCr z!E;-s^9Q<=5+6%Q?n24;4qfHB?u!?+NT;lqoHaN2Aq?>tg(R5v91o~xp}nWBei|&7e{*;P0UlD>2#eX>otb9IZb=;zBbvOl?w?qlIsJd0l zHrDBPTb!_p>`!<%j}gm#Jfqb!U6xHEcI700#kfuy=V?B#q&UE0NDF?SnD>s_o(Z5v zdgPC5+aP&Wy*g$}k4N^>(rQ9)Z7X>?I}1+QhNq+X+?5z*IeJ3AufrQ%AtV#Ow+!y< z#vJDs^i&d{CztTA*qoYq%vC2<;Wu+_qKlSm<(1_yFB{}t&{H>ur&_f7FAl^7%hA&; zO4S7JJf?cOLj}*NU863)(Y~uzTtPLDqw&XAd^_UQjY)?4^OAewLhjHLc{p8bTrL5B z1I{Vf*4D8*Yz@|rhbYk#dGU&r$8~c}ktW0U6{w|`m;hs;7^M?yX{74}dD@|8Ql-gA zP*%zjaPq>9rreotkV~Q`%K3~Y&GQ5Mms(IC9@iy_%3F{rmg^NyP+4OtF8149&ileU zk*7e&z4n2Wk^OG7S^S&12=JV0V*(ov#ox#a&(E)Kl-;Zu9ir8NTv5&$G_1%mxvCPj zIq_OSLgF`xw%NL4H*`#-llzrACfLXRnBi=e1%52#$*uP4DUGdmW498<#Lv84@30Pj zdzjcg=htqU!C5(uGnP%zKShv|%1{zqOQX0qm^f9l&x<^Fy%^IMfg7N5J_+Qbd!P%lQ?{J%gv2; z>x~;L+N1_uZogrP&B;(V(O67KUOx@Wl#7B=(plw3P4~c%wuwRSJb%A$@|zo<*f28B zU0eNVpVq125{sBSj>A%Gdo|3gWHK+vcRk~yysdRbS%Z-6@h@*T$Kn5mlGk;QT1 z7w#(FQJI=MUiPmx_Bcez0%bBj72=rd@;t_?ffvoPsN(Ss^J1Ik<`QV9G}g>nnyP&A z`x<}wyVl3gDHbHDjbk9KR5d>n{X$rmed8e1qT-)XBRW!qSa7*h@#4AKm3kzK;Ej#9 z!IXJ9u`azxnl6~KUUKZ$#75Rp@@tQ*ANG)yZ*l}}{)qk(m2lXA@@THXth_cw^uw0* zg|fI~l1Lruk7*!L*y3``PATNVn-zD)D9{sCHaf;dCoL>Mem2YHhH0~Y$LAC8_sB)l z69v1Yrfv>3n)RlYIBN89BC6#xWPX2~;iiDN0i#vgYE* z=J9>q|52MY?tM;xc_Ld~InOAiZUBRM(I2{bVYK;GOb^`32(=LnsOqq%9-@eOh4fyG zT)V*ZJ3J9$qLllv>OT2kaNd15zBxmCp=q=3%#$4Hp^tNqKE%wHNuDBvoG#C^AS{Z` zcf46Ui-db|X{*}gUUDnj-HC*j&|3|{mNzzjRJx)?g{|joeS_sQ>dGc;k_B&!x{xm1 zSiAjlNwvzR6s?im+Rm8Yad(4F*~(vJgh;Ocs3{GJ?-a__j#9`;>g2Q+XAKzsurMul z@4@Dsl1q}mUe-{q#lsV)ZdP!_`t~dS&9&l=OK_eQ?9Eto`Mr|ZDtPsJ!l+zk*iSQU z#qwI%+=+|d9V0bB%s==RcjNw&&c@ZD`n_>X(|3w#b@V-$F=uj`OwkbBzxe}lE?Ak8 zh!InH)ftKiEt9<+lOg{;;KiL)yCjR%sg;gZ>gUy^T6^tPYDX;+Ge7vM6GIIkICm_K zx;dB)aS=s&YkiRXwGj^5ub#}&y?i9i%jcqd`2%z>pG@=e`?9@h*gIx7mRT04K9*gU ztUi`o_Cg$};33o9D(OwlYQI@lAn2 z|C_D_XQ9LNntB+#(yOp@-Mat&n5I3}JpD~s;Y^VBqN@~|me`OE{V%)UFu^mdUegGD zLuG8d(M+k$qj&#j&pzP8k_3VO&*pv)KB>9cZ>{VPEwK@lUp}s@5nH^eFzC4$LVnfS z7lIA5{)ea#6^@~a9ju`6^Y$MeKbg|H_T}LZEwL^tS_BwqW!J3zxC47|bc0_bMX-vl zK4@nB57&Mzz?;!N8~m{LPsa1XSASP_^*Wgx4liW4{wSLJ1V zA6#rzfP!j7R!Q4-2`tO@i^~-ZqeCIOld_BABL$2MP(~ON8gyf1gl>bu(3WnDTcA5& z0_apXMhKJzCV`%GW86BhJ^Fx;mv~fR)D{RcAPg23@FQ5kgWZBC=v0Q)=7)uis`wO|d(T}n2tk83?Ni->H zba#qxFN+7rUW{gf90DGJU~xAS#Zkn-K>!Fz%=-^1kDDK1}NYpO;n7ZgSt3CXdfUDe5FzBkc$8- zLZ+0D19BPQM4(HfI3Os17vWaQw+nI=IEsj(iHdkBuMBvgDNR%W_dXNV(U#2rinVtu zVyUo?G(Y&CZ0~jYy+6wz=v7L5B5`Mz;f-D+xo#5OuJ^?wtnMETY$_#&OT5^PvhOvz z!9(@M|7LYxAdqw0&wSs9s^~Y$rkTjU?&b&KlZTESU^((H)K3~x-t@~!L(d{er(gd^ zcRm(QD8zI5{AZC6fUkd$F8SX@!h=~WoUHK1XNDhh`o}K;>4Cyz;@g$f8yMRCQ(>WN zg#OU^D9yF#7go64{T|$z^R#aMi;cfl=70@QaTxOJ8pwg)13v}n&Gm&!pB=anxFgsG z+8=tMs14z3UHzJ;UZCQ972=6?^&6g>$160d5MkEUZ+YsE|9zO`P`Kn#pI-|Y2g^^H z6J1o&5|T>xw%2x@2`!@v6|_xI@4=V6rRUgI54$$tsCdj6oEo%@ke-4+(CD*zQgNAg l-vi6OqV<55Gc_0JNBd`lTiwFIu(}FiAQ RobotStatus: + return self._status @property def q(self) -> np.ndarray: @@ -193,7 +198,7 @@ def _run_pipeline(tmp_path: Path) -> dict: # world's sole time-master — the shape a sim eval runs in. simulated=True, ) - harness = Harness(ChunkedSchedule().wrap(policy), embodiment) + harness = Harness((StopOnFault() | ChunkedSchedule()).wrap(policy), embodiment) ds_agent = wire.wire_embodiment(world, harness, embodiment, ds_writer, TimeMode.MESSAGE) world.connect(harness.ds_command, ds_agent.command) directive_em = world.pair(harness.directive) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index be64c992d..80430f4da 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -197,12 +197,27 @@ def new_session(self, context=None, now=None) -> RemoteSession: return RemoteSession(_FakeInferenceSession(action)) -class FakeRobotState: +class FakeRobotState(roboarm.State): def __init__(self, translation: np.ndarray, joints: np.ndarray, status: RobotStatus) -> None: - self.ee_pose = Transform3D(translation=translation, rotation=Rotation.identity) - self.q = joints - self.dq = np.zeros_like(joints) - self.status = status + self._ee_pose = Transform3D(translation=translation, rotation=Rotation.identity) + self._q = joints + self._status = status + + @property + def q(self) -> np.ndarray: + return self._q + + @property + def dq(self) -> np.ndarray: + return np.zeros_like(self._q) + + @property + def ee_pose(self) -> Transform3D: + return self._ee_pose + + @property + def status(self) -> RobotStatus: + return self._status @pytest.fixture @@ -329,6 +344,7 @@ def test_harness_emits_cartesian_move(world): keys.JOINT_VEL, keys.EE_POSE, keys.GRIP, + keys.ROBOT_FAULT, keys.TASK, 'descriptor', } @@ -1686,6 +1702,73 @@ def test_a_real_rig_ignores_the_latency_a_trial_asks_for(world): assert played[0][0] < 1.0, f"first command at {played[0][0]}s: the trial's 5s charge was honoured" +class _WallCost(PolicyWrapper): + """Burns real time on every call it passes down, answering or skipping — a stand-in for a stack that + copies buffers before the model ever sees the observation.""" + + def __init__(self, sec: float): + self._sec = sec + + class _Session(DelegatingSession): + def __init__(self, inner: Session, sec: float): + super().__init__(inner) + self._sec = sec + + def __call__(self, obs): + time.sleep(self._sec) + return self._inner(obs) + + def wrap_session(self, inner: Session, context, now): + return _WallCost._Session(inner, self._sec) + + +class _ObservedTicks(PolicyWrapper): + """Records the observation instant of every call that reaches it.""" + + def __init__(self): + self.seen: list[float] = [] + + class _Session(DelegatingSession): + def __init__(self, inner: Session, seen: list[float]): + super().__init__(inner) + self._seen = seen + + def __call__(self, obs): + self._seen.append(obs['obs_time_ns'] / 1e9) + return self._inner(obs) + + def wrap_session(self, inner: Session, context, now): + return _ObservedTicks._Session(inner, self.seen) + + +@pytest.mark.timeout(60.0) +def test_a_constant_charge_keeps_the_trace_off_the_machine_clock(): + """The reproducible mode's promise: what the machine spends — inside the model call, or above it in a + wrapper that ends up skipping — must not reach the trial. Two worlds, one slow, one not.""" + + def played(wall_sec: float) -> list[tuple[float, Any]]: + with pimm.World(virtual_time=True) as w: + policy = SlowPolicy(wall_sec=wall_sec, span_sec=0.3, steps=15) + return _run_episode(w, policy, _WallCost(wall_sec) | ChunkedSchedule(), latency=0.2, run_sec=0.8) + + assert played(0.0) == played(0.003) + + +@pytest.mark.timeout(30.0) +def test_the_wrappers_see_every_tick_outside_the_charge(world): + """What a temporal stack records: the charge is the only thing that keeps an observation from the + wrappers above the scheduler — never the machine, and never the rounds in between.""" + ticks = _ObservedTicks() + _run_episode( + world, SlowPolicy(wall_sec=0.01, span_sec=0.3, steps=15), ticks | ChunkedSchedule(), latency=0.2, run_sec=1.0 + ) + + period, charge = 0.005, 0.2 # the pacer's control period, and the trial's charge + gaps = [round(b - a, 4) for a, b in zip(ticks.seen, ticks.seen[1:], strict=False)] + assert all(gap <= period or charge <= gap <= charge + 2 * period for gap in gaps), gaps + assert len([gap for gap in gaps if gap > period]) >= 2, f'the trial inferred too few times: {gaps}' + + @pytest.mark.timeout(20.0) def test_harness_keeps_playing_while_a_call_is_in_flight(world): """A wrapper that replans before its chunk is exhausted leaves waypoints due during inference, and the @@ -1698,46 +1781,81 @@ def test_harness_keeps_playing_while_a_call_is_in_flight(world): assert len(during) >= 3, f'the harness stopped playing during inference: {[t for t, _ in played]}' -@pytest.mark.timeout(20.0) -def test_finish_during_a_failing_call_still_stops_and_homes(world): - """A FINISH arriving while the call it cancels has failed still commits the recording and homes: the - failure is discarded with the schedule rather than unwinding the episode's end.""" +@pytest.mark.timeout(3.0) +def test_a_faulted_arm_reaches_the_policy_without_its_state(world): + """A fault is not swallowed as "no observation": it goes up to the policy stack, which owns what to do + about it. The arm's own entries are absent — a faulted arm has no sample to give.""" + policy = SpyPolicy() + harness = Harness(policy, make_embodiment()) + p = _pair_all(world, harness) - class _FailingSession(Session): - """Fails after the directive lands, so the harness meets the failure only when it reaps the worker.""" + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='test')), 0.0), + ( + partial( + emit_ready_payload, + p['frame_em'], + p['robot_em'], + p['grip_em'], + make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR), + ), + 0.01, + ), + (None, 0.02), + ]) - def __call__(self, obs): - time.sleep(0.05) - raise RuntimeError('inference boom') + drive_scheduler(world.start([harness, driver]), steps=40) - class _FailingPolicy(Policy): - def new_session(self, context=None, now=None): - return _FailingSession() + assert policy.last_obs is not None, 'the fault never reached the policy' + assert policy.last_obs[keys.ROBOT_FAULT] is True + assert keys.JOINTS not in policy.last_obs + assert keys.EE_POSE not in policy.last_obs - harness = Harness(_FailingPolicy(), make_embodiment(simulated=True)) - cmd_recorder = RecordingEmitter() - ds_recorder = RecordingEmitter() - harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) - harness.commands['target_grip']._bind(RecordingEmitter()) - harness.ds_command._bind(ds_recorder) - frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) - grip_em = world.pair(harness.observations[keys.GRIP]) - directive_em = world.pair(harness.directive) +@pytest.mark.timeout(20.0) +def test_finish_does_not_wait_for_the_call_in_flight(): + """FINISH ends the episode where it lands: the recording stops and the arm homes while the model is + still inside its call, and the failure that call ends in reaches nobody. - robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) - driver = ManualDriver([ - # The charge holds the failed call in flight, so FINISH is what reaps it. - (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=0.3)), 0.0), - (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), - (partial(directive_em.emit, Directive.FINISH()), 0.05), - (None, 0.05), - ]) - drive_scheduler(world.start([harness, driver, _Pacer()]), steps=2000) + Real time, real rig: a wall-charged call is the only one the harness leaves in flight across rounds. + """ + hang_sec = 1.0 - stops = [data for _, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] + class _HangingSession(Session): + def __call__(self, obs): + time.sleep(hang_sec) + raise RuntimeError('inference boom') + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + cmd_recorder = RecordingEmitter() + ds_recorder = _TimedRecorder(world.clock) + harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) + harness.commands['target_grip']._bind(RecordingEmitter()) + harness.ds_command._bind(ds_recorder) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations['robot_state']) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + started = world.clock.now() + drive_scheduler(world.start([harness, driver]), steps=40) + + stops = [(t, data) for t, data in ds_recorder.emitted if data.type == DsWriterCommandType.STOP_EPISODE] assert len(stops) == 1 + assert stops[0][0] - started < hang_sec, 'the stop waited for the call to answer' assert isinstance(cmd_recorder.emitted[-1][1], Reset) diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index dfd4cd364..1dc4ba873 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -27,7 +27,7 @@ RestrictImageSize, ) from positronic.policy.observation import ObservationCodec -from positronic.policy.wrappers import ChunkedSchedule, TemporalStack +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault, TemporalStack class _FakeClock: @@ -63,8 +63,35 @@ def new_session(self, context=None, now=None): return self._session -def _obs(now_sec=0.0): - return {keys.OBS_TIME_NS: int(now_sec * 1e9)} +def _obs(now_sec=0.0, fault=False): + return {keys.OBS_TIME_NS: int(now_sec * 1e9), keys.ROBOT_FAULT: fault} + + +class TestStopOnFault: + def test_a_faulted_arm_stops_what_is_executing(self): + inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session(_obs(0.0, fault=True)) == [] + assert inner.call_count == 0, 'the model was asked about an arm that is not tracking it' + + def test_a_sound_arm_reaches_the_model(self): + inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session(_obs(0.0)) is not None + assert inner.call_count == 1 + + def test_recovery_plans_afresh_instead_of_resuming(self): + """The fault resets the scheduler below it, so the first sound observation infers again rather than + waiting out the chunk stamped before the fault.""" + inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 1.0}]) + session = (StopOnFault() | ChunkedSchedule()).wrap(inner).new_session(now=_FakeClock(t=0.0).now) + + assert session(_obs(0.0)) is not None # a chunk that runs until 1.0 + assert session(_obs(0.2, fault=True)) == [] + assert session(_obs(0.3)) is not None +>>>>>>> 0595b577 (Let the policy stop a faulted arm, and keep a constant charge off the machine clock) class TestChunkedSchedule: @@ -378,6 +405,7 @@ def test_non_deliverable_wrapper_fails_loudly(self): def test_wire_names_match_table(self): instances = { 'chunked_schedule': ChunkedSchedule(), + 'stop_on_fault': StopOnFault(), 'temporal_stack': TemporalStack(('v',), (0.0,)), 'action_timestamp': ActionTimestamp(fps=10.0), 'action_horizon': ActionHorizon(1.0), diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index b54810a74..d2c45573c 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -1,4 +1,4 @@ -"""Composable policy wrappers — scheduling and temporal frame stacking. +"""Composable policy wrappers — scheduling, fault handling and temporal frame stacking. Wrappers are composable serving-time concerns layered around a policy with ``|`` (left is outermost), exactly like codecs. Most read time from the observation (``obs_time_ns``); only @@ -11,6 +11,7 @@ import numpy as np +from positronic import keys from positronic.policy.base import DelegatingSession, Now, PolicyWrapper, Session @@ -19,6 +20,29 @@ def _obs_time(obs) -> float: return obs[keys.OBS_TIME_NS] / 1e9 +class StopOnFault(PolicyWrapper): + """Stop the arm while it is faulted, and plan afresh once it recovers. + + A faulted arm is not tracking the plan it was given, so the plan is worthless: this answers the empty + trajectory — stop what is executing — and resets the sessions below, so the first sound observation after + recovery reaches the model instead of resuming a chunk stamped before the fault. It belongs outside the + scheduling wrapper, which would otherwise answer "keep playing" without ever seeing the fault. + """ + + class _Session(DelegatingSession): + def __call__(self, obs): + if not obs[keys.ROBOT_FAULT]: + return self._inner(obs) + self.cancel() + return [] + + def wrap_session(self, inner: Session, context, now: Now | None): + return StopOnFault._Session(inner) + + def to_spec(self): + return {'name': 'stop_on_fault'} + + class ChunkedSchedule(PolicyWrapper): """Wait for the current trajectory to finish before calling the inner policy again. diff --git a/positronic/vendors/gr00t/server.py b/positronic/vendors/gr00t/server.py index 6071fb46e..04b2805de 100644 --- a/positronic/vendors/gr00t/server.py +++ b/positronic/vendors/gr00t/server.py @@ -19,7 +19,7 @@ from positronic.policy import Policy, Session from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints from positronic.utils.logging import init_logging 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 ChunkedSchedule() | RestrictImageSize(*gr00t.IMAGE_SIZE) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | 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 09f4784ec..443534182 100644 --- a/positronic/vendors/lerobot/server.py +++ b/positronic/vendors/lerobot/server.py @@ -12,7 +12,7 @@ from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, Pipeline, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.utils.logging import init_logging 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 ChunkedSchedule() | RestrictImageSize(512, 512) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(512, 512) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/lerobot_0_3_3/policy.py b/positronic/vendors/lerobot_0_3_3/policy.py index 5e608c1b1..3128394fb 100644 --- a/positronic/vendors/lerobot_0_3_3/policy.py +++ b/positronic/vendors/lerobot_0_3_3/policy.py @@ -15,7 +15,7 @@ from positronic.policy import Codec, Policy, Session from positronic.policy.observation import TASK_FIELD from positronic.policy.spec import PolicySource, inline -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import resolve_checkpoint from positronic.vendors.lerobot_0_3_3.backbone import register_all @@ -127,4 +127,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(ChunkedSchedule() | codec | PolicySource(base)) + return inline(StopOnFault() | ChunkedSchedule() | codec | PolicySource(base)) diff --git a/positronic/vendors/lerobot_0_3_3/server.py b/positronic/vendors/lerobot_0_3_3/server.py index 1996a9bb4..a6f740442 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.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import list_checkpoints, resolve_checkpoint from positronic.utils.logging import init_logging 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 ChunkedSchedule() | RestrictImageSize(224, 224) | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) | remote | codec | source ee = pipeline diff --git a/positronic/vendors/molmoact2/server.py b/positronic/vendors/molmoact2/server.py index d43a3f4c9..96cfb487d 100644 --- a/positronic/vendors/molmoact2/server.py +++ b/positronic/vendors/molmoact2/server.py @@ -9,7 +9,7 @@ from positronic.policy import Codec, Policy from positronic.policy.codec import RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.logging import init_logging 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 ChunkedSchedule() | RestrictImageSize() | remote | codec | source + return StopOnFault() | ChunkedSchedule() | RestrictImageSize() | remote | codec | source droid = pipeline diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 519b0ad37..bdb39b8b1 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -16,7 +16,7 @@ from positronic.policy import Codec, Policy, Session from positronic.policy.codec import ChangeEEFrame, RestrictImageSize from positronic.policy.spec import ModelSource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.utils.checkpoints import get_latest_checkpoint, list_checkpoints from positronic.utils.logging import init_logging from positronic.vendors import openpi @@ -258,7 +258,7 @@ 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 = ChunkedSchedule() | RestrictImageSize(224, 224) + local = StopOnFault() | ChunkedSchedule() | RestrictImageSize(224, 224) 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 From 35f32c49c1a52c4d74ebce8d990737c11855b95f Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 01:10:06 +0300 Subject: [PATCH 13/64] Follow main's key constants and keep `act_absolute` behind `StopOnFault` --- positronic/offboard/tests/test_server.py | 1 + positronic/policy/harness.py | 2 +- positronic/policy/tests/test_harness.py | 4 ++-- positronic/policy/tests/test_wrappers.py | 1 - 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/positronic/offboard/tests/test_server.py b/positronic/offboard/tests/test_server.py index 3915e3cec..d5f924285 100644 --- a/positronic/offboard/tests/test_server.py +++ b/positronic/offboard/tests/test_server.py @@ -12,6 +12,7 @@ from websockets.exceptions import InvalidStatus from websockets.sync.client import connect +from positronic import keys from positronic.offboard.client import InferenceClient, InferenceSession from positronic.offboard.server import AUTH_HEADER, AUTH_TOKEN_ENV, PolicyServer, bearer from positronic.offboard.server_utils import warmup diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index bbf3554f4..cc67a04bc 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -548,7 +548,7 @@ def _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: for name, player in self._players.items(): # Wrappers do action-timing math in float seconds; the schedule and every pimm channel are in ns. # This is the single explicit seconds->ns seam. - player.set([(int(a['timestamp'] * 1e9), a[name]) for a in actions if name in a]) + player.set([(int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in actions 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.""" diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 80430f4da..f0c44d615 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1608,7 +1608,7 @@ def __init__(self, inner: Session, now): self._replan_at: float | None = None def __call__(self, obs): - t0 = obs['obs_time_ns'] / 1e9 + 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) @@ -1734,7 +1734,7 @@ def __init__(self, inner: Session, seen: list[float]): self._seen = seen def __call__(self, obs): - self._seen.append(obs['obs_time_ns'] / 1e9) + self._seen.append(obs[keys.OBS_TIME_NS] / 1e9) return self._inner(obs) def wrap_session(self, inner: Session, context, now): diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 1dc4ba873..2e9e597fd 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -91,7 +91,6 @@ def test_recovery_plans_afresh_instead_of_resuming(self): assert session(_obs(0.0)) is not None # a chunk that runs until 1.0 assert session(_obs(0.2, fault=True)) == [] assert session(_obs(0.3)) is not None ->>>>>>> 0595b577 (Let the policy stop a faulted arm, and keep a constant charge off the machine clock) class TestChunkedSchedule: From fb205b18650ad1211e37f3623cd6d4a60c56c644 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 01:27:46 +0300 Subject: [PATCH 14/64] Land a stop in the round it is answered, whatever the trial charges --- ARCHITECTURE.md | 5 +-- positronic/policy/harness.py | 11 +++--- .../policy/tests/golden_pipeline.json.gz | Bin 7883 -> 7927 bytes positronic/policy/tests/test_harness.py | 33 +++++++++++++++++- 4 files changed, 41 insertions(+), 8 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 689ee5e7c..a20792fc5 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -136,8 +136,9 @@ trajectory". `inference_latency` — a constant charge in seconds, or the call's own wall duration — and the harness alone reads it: it withholds a returned trajectory, and the world clock, until the charge is paid, and the clock it hands the policy stack (`now`) reads the instant the in-flight call's output takes -effect. A scheduling wrapper stamps its chunk at `now()` and never learns the mode, so the charge -is deterministic wherever a constant one is asked for. +effect. An answer with no waypoints to place — a skip, or the empty trajectory that stops what is +executing — lands at once. A scheduling wrapper stamps its chunk at `now()` and never learns the +mode, so the charge is deterministic wherever a constant one is asked for. **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index cc67a04bc..d32ed1173 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -513,10 +513,11 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) Under a constant charge the world holds still until the call answers — blocking here blocks the loop thread, which is what advances a virtual clock. Until a call answers there is no telling a skip from a model call, so letting the world run meanwhile would spend trial time on whichever the machine - turned out to be slow at. A skip then costs nothing; a trajectory is stamped for ``t0`` plus the - charge and withheld until the world reaches that instant, playing what is already scheduled on the - way. A charge measured in wall time can hold nothing still, so there the world runs no further ahead - of the call's start than wall time has. + turned out to be slow at. What the charge then buys is the instant a trajectory takes effect: it is + stamped for ``t0`` plus the charge and withheld until the world reaches it, playing what is already + scheduled on the way. An answer with no waypoints to place — a skip, or the empty trajectory that + stops what is executing — has no such instant and lands at once. A charge measured in wall time can + hold nothing still, so there the world runs no further ahead of the call's start than wall time has. """ if self._charge is not None: concurrent.futures.wait([future]) @@ -528,7 +529,7 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) if not future.done(): return False actions = future.result() # taken on the loop thread, so a failing call still seals the episode - if actions is not None and self._charge is not None: + if actions and self._charge is not None: # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP # and slips the install a full round. if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): diff --git a/positronic/policy/tests/golden_pipeline.json.gz b/positronic/policy/tests/golden_pipeline.json.gz index 48769162b0ef26a025b114bc990c0458e96b2b1c..ac1cfbb81578b5a0cf270d771a6fa3916b1e6b61 100644 GIT binary patch literal 7927 zcmeI1c|4Ts`^S}Q{`>v$dp*}{Uf1h>p8LM9`+DyCdA&c^fNybe zy-g2v-nx3p!}WrLyX|GC%MPwi?hXph-X888<0&+(NNtQ>p99J}EK2xO-8sH(+JMmN zt>N#LomQdC6P8^XqKBZN+bl$orwo4|NC%B4T3$6b^*sDcW-QGJz4YFwB4~MSjEcuW{$O?Ji9{Xl_-(LGZd0(piWZpC~(j4^UGQy5j zi3$6OTTtMLD`B~0A}=zGcY5{%SUgB-caeQmZ#|;*PsRyaT_e7X1 zPEGaqlZ7L!swyKQczM@u_^}(u-a}s&NeIB47qj~D;W#0f%!h4<^cJ(Ce?Dlx-thU( zXn_APjWs<^0A{i1*pJ_E?)KJx6n#3?QVHh<%Lf&#J8p!RHskDP#b5!8BmMY@O-DsR zpKnlS?^ms3?@ir3^yGq~rl6IzT;Tni4&GBb+I$gYTQu)S8gG}zq;-%elE(UW83F<~?J|S}>g_Vb1a{hINZj^~v%be)tz#+|SFK|$ z$6ur4AQxBj43_<&$nYLKsg`0qg40LGG4LJW>Xoboc!}mH(|8JZ6CJmV9|r5KW-;Jh z4=04f`m&Z!d{hwydTRlxQmZGJDjY;{4$dt=07RUa&cVqE+y&7hif`fE4itbOi0NB6 zcL2p8A;e|z0?w8x06*v~0&$JB9RLDdK%`$I0RSSPD~Na3NIL-HpiqSPbrKig04NTD zxZaM@CgR1PrBS<1ga3m_Un=IU87ZRZ@1_%I+ z+YuwQ!2lF!*^VPpTO4p51hm(R)cyq!4id2siPV;Ofz23)-ez8b1X2#v`@In1*?yvc z7#YI1W&hR_+oHHYd=LqQ?l%p-z;@hIHR_KZga14*q9fuKiRbk^B7yyF8x?C))e@yz z$_GC-Faj6hxq;vBS#SWGTGbM#TDk{rIWTgXu9dZFKXA{Aw4& zrT~8*`}*^Km+N>ZQ&*vpw5m*`+5u-j@in7MkyIfJ1i&gsMc2M(-IX@0QQOX=Og1wx*~t=ImPbJ{K-* z%(UwWZZ~LRMYue4ELXY2r%|Y`Xk^@bLHR-9S~K9!G!LGA<@G{tdeQBKSG&NSwuKW( zDP2j-rDvVAvv*y_0|5;Z=ZkvI2RwQd)zH|Z@<(8WTg_u$zm(2>vnL;`Fz)6Rnm#yu zG`8>T%)}y%*K~fAJUGl*c&sUxj> zEAiHzMTwpAVVqqLzw~N-J!;fZZ)towr26e^fdXSC6>lZFiuWWLTp?qY*KzOcCxSIX zu)|;Wjnc8_N=ua2DAa}LX8CS53D;hg$9FlABmC2ppR?|vFD(S0zj=fqG++#?Fq~;3 z9BAx*op8ZRvf}B`QYk-KKZSg(tFZ5xOE$iti}ie;yDuT%1WLcR<4Bf9Uhkg~&84ID zJ*+kNH4s$zJ3)=5JVL)PkHn5SjL#3TS8OrkRw7T~b`G?=u?>uRotK$FWd+-(LBPsAdNW$6Q@8*;Z-6-j>3Toab$m>l<0`hb+lU zNgJ(J&Yii7o7C&;5OLT=3ky8Vit6%0H@DHGDI7fZ&rLQI;ShOUk=|qT5;kk zYhR+XdOVX+PD`OlqjyjF?oHRA^wESC%q1^*Nup6Rr4wZJaLZm0D~(sH3LT^^Yef}yuZ@RJ`vzRAdi2_NSR zz85R{#{&XYQYRA(6e5(EI?c*?briijtQb7Ri)abI=60V}mc{h6yVCSdJSZwxUr&_( z_&GqAGDo+^;k>(lw$p=QgKFoJ&jwz+v>@D+>iv-3?My@dss32W<8k?w7^?>tN5M)y z0u{>jk0*^U-Ky5f|Jmt&kqPi#Y12BJaWYuG`1N8Cyg?N{;;-9h^|BS6P^FWbu#|(2O-1s7EMV3us7N-TtJo7F+t#zq6$j2rTInu1JLG;1N zhoTDuI|!D7qEp6q3QNGS(Tf`1tSyFCx|-x~ImUQZDzhhVIG8iLdBCMlobGaUnw9TM z3!N#w6my3K=xb&3{UsK?DL5vz&RKpfvB&)LwXeDWAKL^frjI%16%fps^R{$$+_#^@ zX*9Kj1>{Me76vHTCM#!5_%_4kbqi}ttlGsI7S#{lAux;sAGCGa9LbEb4ODa)zUV#- z^ZM{uJ9yj-eoV|u^igArN`1VoFvwv7txk_X69O3&MKt66APx*$nVbGFxkCTQ)J%h+ z7N-@#sZ_Mb`<1bj!Vf$e5&^o)eADpzFo}6SC)~+EUB>=aaLjaM6#N2AVwUeb?p1)U zSYUCHM2&kb%5uE8I@NN#w7T4K{AD%Oa{%&J+Ld2jpO79m0cw)oc{vH`^lZJk$!9# zZksYD!+F5UL1&wO`t9bFvG5`Ns>-#7&v)A9XBsH%=|5xlnxk5VdvlBh?_Ys>cYr@( zO%iAX%}uqC;k@7t$HOX4q9#n{tJcofk9w##S?Jx54BZ6pHD}O7h=tz8173AJo0x8{ znBN%4=IL-2a$j60@HHoT>+TIqkL}i;=3>M8{Ih#&qF@8OKNGle{a2WNZ>I@BGY`W8 z?6^uDc{zTvjjwV&6D}6I8$6V@GSgE;GR9L;aRJ22|NwY_?F77a{ z2>%y2enn~?o~L<}X-va;q1%@5@4$Okr9QzMHE*sM&*PpM^{VWwF#Z(X>Y)Z24gtM1 zUWsOSaD-@c0JegF2(1w9EdT&W1Q8OVy%n$%B!?h|XmbK~fz%Nk*FI@G?zq>_BL&GK z61QTwfe^?mq82BH2dD}eArd(;yg)4o7fEXyh7WiWB1Ix@d$a>B#&J{ScM@VdrvT6b za+&mQJ0}ZZA%Bn%08XJ?(c>IIZEk=9=p}*!qYdoXbsv^=M=G2Mzz72`Ltc@zc3?z- zevlCoaR&wr41sV}YH?wH0Y*ZkDv4Z}{kdSyaAZ}8ig{(oPZ&wyU5H;L$4-nC@IEB5 zGGr%42AB&euH@jx$N>u=Z!1H%F$%6CSbbvspY83|`M%PvJco$M`=_6v{#i_Vy*P}u^qa3gO0LO`u! zm8@zBQ|;aj2M@UFillJ86eUFB_jz8~kM*r;!Kn)Efx`zJPl;&Uz)O2xIf#8x)sn#R zr3_ehd@d3!@WW{~bwQSx;eGtAP?f(?5Nu{|soc$GX_}+a01u1I*W8#Cjl=i;Av0%z zd5Czjg0_@*HvYMm-Hfj6!d!P&wMU*6Kn>K{2R6|7P&B#2s=7#g#9TF33#8Vb9DIkZ&kq3wM1^1vVj-KL@RV=p>>QjAJ(AJ_SfCLEcTLJ?!q&h1}d zk|L&kttG}xDf!KXqUa<38TN#x2T%L#|LUZ5_Hc*XwktY6z;86ovem3bysrjY!Cf$* zM#=+k8>wL;9nZFpE@q{IJlT`yylM@@nu$wRN6dC+{?4}G)W1?n-4s`?7ILjlMhBo& zGfiaJR&FqQRP)JV#+6mq{Fz^7&+2Egjd3Tw`w+NHI?3G>FJRkfd3vo?KND@6t=TLR z;PfI*{HN^X*6!wC{jAAnWD8ALHa1_7I;l_L=Ku7>y#BSjMoh`C4W;>aHX`f8wrRW_ z>QF{rLzJNln!7p^JJthW{jS(%rh*kk*gRjI1*h9JPkAy;$5Jh_KZ8dbdO&+gR>3N|_?yxGxImQajd&bX4CCOVdL(A43is^1Xd(o7Z~6a3eYS79Xka7H;*p6bihWpgWy- zp*{U@dzdQo5#^Vvz^-+f$EvOVAD zmu_x{M`oqy%h4^l6ogDsUn+R!u!6nf=8%~tkNT*77tLNy1)Kt-&g8Ikz?t=XAFoO& zce(4(d71;QAO3Mi&n3Edm$Z!X=8h4~mUR-0d4wt;aq7aBCT?2YFnv_f#HqX9*~7y5 z?Mf=4INY?OYN2p4H5uQK!BnHK=dUe5FR%T>{L{m)pY``NJY5*bLAn>0`{Fk;*ru|- zy^wOymsHj#rWxee-a2Y(Wd?T7t7yW{u5X6y`CQWo-^mLF#F*0RL#I$r zt!?JkdeM*rddasVu^+1PX-%{>dR5eYBP>Uoto$5H2Fl)(XmJ&4Q z&qSXrRPOKsm-o&6vDq8;dAL~jUADf^)Oa;;S}Lu&)23#~C)IoNc(es*5Za*{XQV~)Tb$IgxjhWp^u3mZ@TS)W!*kta8Ef+rc-p8Vk z74|BxzUJ~+?;G?g4ZXtny3+r_gZ^cuf9CCttN*Vo{$E-AUr-i6=O?~VfT=I4`sT+* z7*#GTFck-eLO_)vTo?&p6hyU><0s4kU>xKd3GoCKC5k0+Qi(e;vcM;hSCv}a7T-dC!SGJ*bF{u2oL56um|Et=GcW%1r9+H$sxNiYQS;GHxlBBf7BB8 z8knlJn4|)E2aLTvfAhx`!B~exF zXSmD2h@nX14g4|BfT5^MJWnZ1|q< z)A>E$D>9#p{a0Q$`TmPyzs-VUF>HR!}kmRi(-EQ!|@Ed`ysp7 zZ@N+Jr>9&0u-Jd4A9qCk`=$G@{|6BD2G$X`XNu$pA#zbp1ndKah7mNBV)3z@canYPQOG(%+#?{s0puK zX-FE|$P`+XrTB7bM+>PYv1ja}_>bp3)b`toA1&XRjoM-%`>Vr6Ok{cD)f z^3-G}48vxATv;9`msm`!%+E}cUt`&fnPsHvEMlV2(D##AtKuL&2BE+d%flZJGo~uO zh22_?iAA0@P?#1=#Xk=a50oSho#uYUP-G+U6l9T|*)0rLFrNh<-E zoRLBJ#rKMm54cU_t5-!tR-5KR=XdlvhzY{9!rK^DWTX%*5q&vIA-t{elK6wQhChA@ zBpFPxY#Zv4f-v>)M~v1rRHz)L>N0@7g5PdrJ7@QY68CWX1IB}T ziU3SCypZwWUzZc=Ea?5-@ygkymV$!MN-c#1he|C)1UEmh6f@Ir+HZB;<*tfElgr)x z5(4PE2P90&o-{CSVE6(mW{9$V)u@fRrf8GV&5H{}w5E)Xm6CeEfn^5b90j zC4gT<>M#l;O17D}m$NhEy#|G>Qo1S>z~P z-&I=I&kx9e{?MB`)D}Pv422kV{aXNe@G>M?kJ<_-fY%|YzJDv75;-6R?-xpfQ1w~c zfW2S_6j?7X0PF*Epse~V0YDWjhHC5O1p#%i9O|#n5|pM~_#`DVAXE>Dl2IUV0BnLF zaz6+h1m8lIWRwta2&6(6$o)cq4mb$iB%_4WDX0ml@d2UFP%T+r1TX|yP(L|K1VDhx z5Kn`=DDV@=50`Jq5(SQe!tlcl@?s7YOO{kKMlXfrE$(XIfQ#l+v_~1&#zczoD)O0Z zirp$D1}S{ipdUR-v6J=vPPco6ZGOM)!X0X8y{aQEhTO~Ni_{qOFp=qrA;|j8DB_|U zT5svtiz6)tJxpaBuMyP!X7>DPUIbMQnYsCU)BNFHKJtf=PYMGHz0Q6#|l#4{+g|685DF!SESSC0Mj z)4yYyJIX`uC5;evL;JtZ|AlF;heZUfh?yLO_dog$U3;4uZ6>pS(%(Gj4S}oqJ#H8f zzxvmTzWoGpLe}@fHS!+eWSujE-&G#tu3bU9@qT**><1-E>6O@ou#)3X8ZGHNjf_3@@h=neQWqj}dT-pE&oJyE zk&nv-I%11|8q(5}OL~#IRQ`MBKDFx<#o(Uh85P)p*WE2IyAe&^d(+~W&$W7Rb@t&} zzUl#*vx~1u!FjA>KfPoH>fZ`)c3d?zN}#h= zyrynE0x>V_^|Oy-naPy8-s0_uxO3{<=9)_%KAIhkE_rC;V1*9Ucu1x&vpWpSIF?q< zu%5_X4~UC+8F_j)JxSlpH!zVC(IUV9MDojjGxK^!vx0J~tg8y*IR3Y=$=#A=C-B`4!RRypvMp$8;WRA=WR^V0^bQ?~SCu?C zLAhmD^jc2aE|zruCuyw!P3?dcXx3WKH0 zYEzxx-COO^)8=S8xknAjN zGDKBNZaNc?p!-NNuxYFXL+ z-cW_0nOqm~QK)Qwfk;IMcH-Ux&Vp*+a-5A)Wt|j47Ip8meaRXN-VtHDV)&$)vj@fg z6mQSsalkDx3b`4L*cWF-@Z>FuDCK`rC$tT&A9eD2evcsv$EtKBOIhUIA)A1CHc%mba z#g6Q5+I2d@S>a$Yns|cf9^_UtFs2(ET2j7BTJ#*Yfn{6gKBgaA=4=ELFYYSpm#r3A0S|NF@9MBG^Dd?s) zdKA!}Iu*e$wG}y?rZ$f!-ReqZ=akWyLX9Z__B-0rK9kP}&$7~X>4i-?Vb$tdGMhIeRlnvJlg&`|XR2d~~t()%<;&MTpxoA-#w7WE;sr zQkG$rMy#B07;)D(g=u8bJ&GQ&!GF?gJ}Yw0zq&h~N()pQr`;UA z$Am9DwLI0St;U?q@W079xIIegK8X~qFeZUzo>xK|G~H1%`?!yDJ$`w`9&w(?SZ2|% z%dE+dA~5Xy99?8(dS*G(0)|*x>J+hRnmps1uOJ}p9uPR!cf`p~=!x9^wK&Bbq$pY?UAz4>3TI2q}z@!EeCk-wnsRs4=Sd9IH1F9_2NH}~A~Opy!Woz2%Da*6&+>gM8i zyD~S|Sfy?*e&0|zzGLZ;aPOgO)&hdWQcF;zQREge(*`DA9xivDcUO!0&uVqk|$?@ zfF}3~I!u-q0<^#m2u02k0<^(iD3UBM4CsQxr6C{-*DBZBFANxfQ_veSN(3+h7a#_? zUj#S|^1($LP@=#QPyh}!^os(N+6e&6Aaa1NErcBK5Hw8w}i~`QBp*2p2DuzT1dK{OLy++vMH?td; z)6klzLls9l^L>fUQ-qiPGp8Bfn`R?JxEFfx7H+?I@1N*>!xugiW5J_bVkHs`udU-P zk#H=C^*rWp!t_1NK_0xY{NLpC%Vr00tl`$a2I+6W{GVWYxb>;wlZy4M;QMsns%QTW z>AUc5c!rbsvd?iLR)cbOi=pRNIQwS~uOKinGrg1ie~Y?T$(){Ir{gBU_~wW?_^kpipaq6fj$Zl- zuuCZS-@sM>sr}6$qRcR%vom<-<*8Zc%-)qb^!&hNbWa$Zo)f%cg}HdK&%I1gm#8rl z9^?}?Q;~`3oRx+((gLp+{^DKJ+ZnXf_2J5Ore0d3lXrCqEwfV6Q;qCX*03*QkRvMv zgLo=)srvDcLpumh<1o8e$Fprrxl(d^#OdJu>dYrHg7*!8iGVJZkQPgJ2mCkpe)hgZ^yWf*tUhJu28V|pP=+&O9 zZrPYd=v&XC;vG6;ubdfQ?l5w9M%$aMI>7Y-kr-M{x-h5W^=LamuIH-Z8o3j#Rk>3A zb5WqG90%Sf2+yX`_tW31k;aE`8;)8i_xP$*s17b(ySeL>uVYS^y@RVtw)Lu1ANO?L zR)#)B`b?omTUXW0FJu%S;)chKu)sFdk3DyE#$ZmS-cpd03yZOau>hAj!v6Y%-!F=}MKL)D2m zRcw3W+TNNv6W_AqykgyY_)k`{gH~=P`-z2}*f5!B&t^4fmJME^m@P2Ut<_008SBHi zJ!VW2(#~7Jn1*9Lst>QUso8stnYdZytu5X@=J;{P6IC&1m;@mo_H-N$0SiEo8uA5G~(y`dv<6~5{zPW7`ip0k8qW9Ki+mBFx{(UtU zSic|%NBVBOfgJgoql+hKcUt3zwZhX#Da5yG8;41DU)<>%c|@mT7jemxpEuv0jb)7W z(A$o$pZ;IksGWJqgWgFGZ#O1Z40zecYfq=HFDUj%-Z-{6)HOZv(p%4RR(xG(&P>8->Aw4I)kXR1-~c!?O30UgRV_uq(x2rL>&h=RMO=xIiydi(d_AW;nr9 z`n_p+2+98a?~^{vw>ND{W;fg|@Fbe~#OmfI@*WF@xh>`=zb#%lseHwnBcHr}7_F9+ z)P_Ykt-!3kag=#F6zP?Zt2C9(&)IlI9myk!j$@9vPINqD zMiUNt1(f{U`!e)@uU&Hs&AwCnjAUxSWRH_Ullvqw86I+Hv3vb;sS=$m3Y?wQ%kZb| zhcY^>iauA;w9N84)?O0Uqv^Noj|~Ndyc{LY9y#ju$Tm%Tyd)U5esP&;_5@bYo!7I2 z1A*`m{e`if-(Do^g>Agm?k0>hwiS03z-$b&t6A|G&B~fmog%0aimi86%(3+vM@B=bx@ROyCy1*sh}bE|*d_omsp;qPdx zh2Bon??5=UHE#p!sjcezG*@E|PXA%zFU$68q;G`?y@b>h4(zCo|`zQ% fault_at + 3 * period] + assert not late, f'the faulted arm was still being driven at {late}' + + @pytest.mark.timeout(20.0) def test_finish_does_not_wait_for_the_call_in_flight(): """FINISH ends the episode where it lands: the recording stops and the arm homes while the model is From 58d684b7c4756371fe9136eb2e591be3a15cbd06 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 01:28:40 +0300 Subject: [PATCH 15/64] Name the grip channel with `keys.TARGET_GRIP` in the new stop test --- positronic/policy/tests/test_harness.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 285d7385d..471271db0 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1821,7 +1821,7 @@ def test_a_stop_lands_without_waiting_out_the_charge(world): harness = Harness(stack.wrap(SlowPolicy(span_sec=1.0, steps=50)), make_embodiment(simulated=True)) grip_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) - harness.commands['target_grip']._bind(grip_recorder) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) From 577a774c7e074a9fe433a5f7781fe8dd12d85aac Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 01:39:16 +0300 Subject: [PATCH 16/64] Read the fault flag as absent-means-sound, and name the wire keys in the new fixtures --- positronic/policy/tests/test_harness.py | 33 ++++++++++++++---------- positronic/policy/tests/test_wrappers.py | 8 ++++++ positronic/policy/wrappers.py | 5 +++- 3 files changed, 32 insertions(+), 14 deletions(-) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 471271db0..3581ac7c3 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1090,7 +1090,7 @@ def test_finish_stops_playing_the_live_chunk(world): stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'FINISH did not emit STOP_EPISODE' - grips_after = [data for lbl, data in events[stops[0] :] if lbl == 'target_grip'] + grips_after = [data for lbl, data in events[stops[0] :] if lbl == keys.TARGET_GRIP] assert grips_after == [0.0], f'the cancelled chunk kept playing past FINISH: {grips_after}' @@ -1293,7 +1293,7 @@ def test_shutdown_stops_playing_the_live_chunk(world): wrapped = ActionTimestamp(fps=5.0).wrap(ChunkPolicy()) # 1.8 s chunk — won't drain before shutdown harness = Harness(wrapped, make_embodiment()) harness.commands[keys.ROBOT_COMMAND]._bind(_LabeledRecorder(keys.ROBOT_COMMAND, events)) - harness.commands['target_grip']._bind(_LabeledRecorder('target_grip', events)) + harness.commands[keys.TARGET_GRIP]._bind(_LabeledRecorder(keys.TARGET_GRIP, events)) harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) @@ -1314,7 +1314,7 @@ def test_shutdown_stops_playing_the_live_chunk(world): stops = [i for i, (_, data) in enumerate(events) if getattr(data, 'type', None) is DsWriterCommandType.STOP_EPISODE] assert stops, 'shutdown did not emit STOP_EPISODE' - assert not [lbl for lbl, _ in events[stops[0] :] if lbl == 'target_grip'] + assert not [lbl for lbl, _ in events[stops[0] :] if lbl == keys.TARGET_GRIP] @pytest.mark.timeout(5.0) @@ -1579,7 +1579,11 @@ def __call__(self, obs): dt = self._span_sec / self._steps pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) return [ - {keys.ROBOT_COMMAND: CartesianPosition(pose=pose), 'target_grip': float(i), 'timestamp': i * dt} + { + keys.ROBOT_COMMAND: CartesianPosition(pose=pose), + keys.TARGET_GRIP: float(i), + keys.ACTION_TIMESTAMP: i * dt, + } for i in range(self._steps) ] @@ -1614,8 +1618,8 @@ def __call__(self, obs): result = self._inner(obs) assert result is not None, 'the inner policy of this test wrapper always returns a chunk' anchor = self._now() - result = [{**action, 'timestamp': anchor + action['timestamp']} for action in result] - self._replan_at = t0 + (result[-1]['timestamp'] - t0) / 2 + 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 wrap_session(self, inner: Session, context, now): @@ -1642,7 +1646,7 @@ def _run_episode( harness = Harness(wrapper.wrap(policy), make_embodiment(simulated=simulated)) grip_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) - harness.commands['target_grip']._bind(grip_recorder) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) @@ -1866,7 +1870,7 @@ def new_session(self, context=None, now=None): cmd_recorder = RecordingEmitter() ds_recorder = _TimedRecorder(world.clock) harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) - harness.commands['target_grip']._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) harness.ds_command._bind(ds_recorder) frame_em = world.pair(harness.observations[CAM]) @@ -1906,8 +1910,11 @@ def __call__(self, obs): 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 [{keys.ROBOT_COMMAND: command, 'target_grip': 0.5, 'timestamp': i * 0.01} for i in range(10)] - return [{keys.ROBOT_COMMAND: command, 'timestamp': i * 0.01} for i in range(10)] + return [ + {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)] class _GripThenArmPolicy(Policy): def new_session(self, context=None, now=None): @@ -1916,7 +1923,7 @@ def new_session(self, context=None, now=None): harness = Harness(ChunkedSchedule().wrap(_GripThenArmPolicy()), make_embodiment()) grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) - harness.commands['target_grip']._bind(grip_recorder) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) @@ -1944,7 +1951,7 @@ def test_home_and_manual_commands_are_emitted_as_plain_values(world): cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) - harness.commands['target_grip']._bind(grip_recorder) + harness.commands[keys.TARGET_GRIP]._bind(grip_recorder) harness.ds_command._bind(RecordingEmitter()) manual_em = world.pair(harness.manual_command) @@ -1964,7 +1971,7 @@ def test_abort_discards_a_call_that_is_still_in_flight(world): harness = Harness(ChunkedSchedule().wrap(SlowPolicy()), make_embodiment(simulated=True)) cmd_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) - harness.commands['target_grip']._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 2e9e597fd..30aceb527 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -82,6 +82,14 @@ def test_a_sound_arm_reaches_the_model(self): assert session(_obs(0.0)) is not None assert inner.call_count == 1 + def test_an_observation_with_no_flag_is_not_a_fault(self): + """Only the harness stamps the flag; a probe replaying a recording has no arm to fault.""" + inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + session = StopOnFault().wrap_session(inner, None, None) + + assert session({keys.OBS_TIME_NS: 0}) is not None + assert inner.call_count == 1 + def test_recovery_plans_afresh_instead_of_resuming(self): """The fault resets the scheduler below it, so the first sound observation infers again rather than waiting out the chunk stamped before the fault.""" diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index d2c45573c..27faedb97 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -27,11 +27,14 @@ class StopOnFault(PolicyWrapper): trajectory — stop what is executing — and resets the sessions below, so the first sound observation after recovery reaches the model instead of resuming a chunk stamped before the fault. It belongs outside the scheduling wrapper, which would otherwise answer "keep playing" without ever seeing the fault. + + The fault is ``keys.ROBOT_FAULT``, which the harness stamps on every observation it builds. An + observation from anywhere else — a probe replaying a recording — carries no arm to fault. """ class _Session(DelegatingSession): def __call__(self, obs): - if not obs[keys.ROBOT_FAULT]: + if not obs.get(keys.ROBOT_FAULT, False): return self._inner(obs) self.cancel() return [] From 8cdf5d996b36cad0e7538c3c2852db694bdc9268 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 17:46:49 +0300 Subject: [PATCH 17/64] Bound the env server's close acknowledgement, so a wedged peer ends the run A simulator stuck in its own teardown holds the socket open without ever answering, and the unbounded `recv` waited on it for as long as the run had. Found driving a real LIBERO eval: a policy error propagated correctly, then the eval sat in `EnvConnection.close` instead of exiting. --- positronic/simulator/env_server/client.py | 18 ++++++-- .../env_server/tests/test_remote_env.py | 43 ++++++++++++++++++- 2 files changed, 55 insertions(+), 6 deletions(-) diff --git a/positronic/simulator/env_server/client.py b/positronic/simulator/env_server/client.py index 4f6ca6f96..ee8dd2438 100644 --- a/positronic/simulator/env_server/client.py +++ b/positronic/simulator/env_server/client.py @@ -4,6 +4,7 @@ control system; tests use it directly to compare a socket rollout against an in-process one. """ +import logging import time from typing import Any @@ -12,6 +13,13 @@ from .protocol import decode, encode +logger = logging.getLogger(__name__) + +# How long ``close`` waits to be acknowledged. Teardown often runs while the peer is on its way out, and a +# simulator wedged in its own destructor holds the socket open without ever answering — an unbounded wait there +# hangs the run in place of ending it. +_CLOSE_ACK_TIMEOUT = 5.0 + class EnvConnection: """One websocket to an ``EnvServer``, opened with retry. ``reset``/``step`` block on the round-trip. @@ -52,12 +60,14 @@ def _request(self, msg: dict[str, Any]) -> dict[str, Any]: return result def close(self) -> None: - # Best-effort: ask the server to release, but a peer that is already gone (a crashed or killed server) is - # success too — the socket is closed regardless. try: self._ws.send(encode({'cmd': 'close'})) - self._ws.recv() + self._ws.recv(timeout=_CLOSE_ACK_TIMEOUT) except ConnectionClosed: - pass + pass # a peer already gone has released whatever the acknowledgement would have reported + except TimeoutError: + # Abandoning it is still better than hanging the run here, but a server that took the request and + # never answered is wedged rather than finished, and its resources are nobody's to reclaim now. + logger.error('Env server did not acknowledge close within %.1fs; abandoning it', _CLOSE_ACK_TIMEOUT) finally: self._ws.close() diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index e59ca2198..61729b9da 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -1,9 +1,12 @@ -from contextlib import nullcontext +import threading +import time +from contextlib import contextmanager, nullcontext from dataclasses import replace import numpy as np import pos3 import pytest +from websockets.sync.server import serve as websocket_serve import pimm from positronic import geom, keys @@ -20,7 +23,8 @@ from positronic.policy.tests.test_harness import StubPolicy from positronic.policy.wrappers import ChunkedSchedule from positronic.simulator.env_server.adapter import EnvAdapter, _in_env_control_frame, _wire_command -from positronic.simulator.env_server.client import EnvConnection +from positronic.simulator.env_server.client import _CLOSE_ACK_TIMEOUT, EnvConnection +from positronic.simulator.env_server.launcher import free_port from positronic.simulator.env_server.proxy import RemoteEnvControlSystem from positronic.simulator.env_server.server import EnvProtocol from positronic.simulator.env_server.tests.conftest import serve_env @@ -95,6 +99,41 @@ def test_transport_is_transparent(env_server): assert direct_step['control_dt'] == socket_step['control_dt'] +@contextmanager +def _mute_server(): + """A peer that accepts the connection and then answers nothing, holding the socket open. + + What a simulator looks like once it is wedged in its own teardown: the process is past serving but the + socket outlives it, so nothing ever closes the connection from that end. + """ + host, port = 'localhost', free_port() + release = threading.Event() + + def handler(connection): + for _ in connection: + release.wait(timeout=60.0) + + server = websocket_serve(handler, host, port) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield host, port + finally: + release.set() + server.shutdown() + thread.join(timeout=5.0) + + +@pytest.mark.timeout(60.0) +def test_close_gives_up_on_a_peer_that_never_answers(): + """Teardown ends the run: an unanswered goodbye is as good as a closed socket, and is not waited out.""" + with _mute_server() as (host, port): + conn = EnvConnection(host, port) + started = time.monotonic() + conn.close() + assert time.monotonic() - started < _CLOSE_ACK_TIMEOUT + 10.0 + + _HOLD = {'command': {'type': 'hold'}, 'grip': 0.0} From a33f100454d46640754add1593642ec7f4e1b81f Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 20:49:37 +0300 Subject: [PATCH 18/64] Mark the fault's raw-sample recovery as a HACK against #619 --- positronic/policy/harness.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index d32ed1173..1d0610ae8 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -452,6 +452,10 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: if obs.serializer is not None: value = obs.serializer(value) if value is None: + # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so the + # fault is recovered from the raw sample and stapled on below as `keys.ROBOT_FAULT` — a name + # already claiming to be part of `robot_state`. Emit it from the serializer and this branch, + # the raw-type check and the flag all go, and the fault reaches the recording as well. if not self._is_faulted(message.data): return None faulted = True From db2ef6eaa0929085e3dcb01db2cea57d41d5f061 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Wed, 12 Aug 2026 20:58:08 +0300 Subject: [PATCH 19/64] Put `StopOnFault` in the ready-model example, which reads the arm by name --- docs/connect-your-model.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 34af81bd1..de7a59a58 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -170,7 +170,7 @@ from positronic.drivers.roboarm import command from positronic.offboard import PolicyServer from positronic.policy import Policy, Session from positronic.policy.spec import PolicySource, remote -from positronic.policy.wrappers import ChunkedSchedule +from positronic.policy.wrappers import ChunkedSchedule, StopOnFault class MySession(Session): @@ -201,13 +201,13 @@ class MyPolicy(Policy): return {'type': 'my_model'} -pipeline = ChunkedSchedule() | remote | PolicySource(MyPolicy(load_my_model())) +pipeline = StopOnFault() | ChunkedSchedule() | 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 `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 `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 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 on a real rig it wants `StopOnFault` outside that scheduler — a faulted arm is not tracking the plan it was given, and the wrapper answers the empty trajectory so the rig stops rather than resuming a chunk stamped before the fault. 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 scheduler in particular, and `StopOnFault` outside that scheduler — a faulted arm is not tracking the plan it was given, so the wrapper answers the empty trajectory and the rig stops rather than resuming a chunk stamped before the fault. It is what makes the session above safe to write: a faulted observation arrives without `robot_state.ee_pose`/`q`/`dq`, so a session reading them by name needs the wrapper ahead of it. 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. `new_session`'s `now` argument is the runtime clock that wrappers scheduling against live time read; a policy that does no scheduling of its own just accepts and ignores it (server-side it is `None`). From 7e374374a0764a64905a19177c4f301be5cf100b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 13 Aug 2026 11:57:17 +0300 Subject: [PATCH 20/64] Give pimm's no-default sentinel a type instead of erasing it to `Any` `NODEFAULT: Any` let every unparameterized receiver infer `Any`, which is a blanket silencer rather than a contract. A dedicated `_NoDefault` in the parameter's union says the same thing without erasing anything: the checker reports exactly what it did before (1044 baseline entries either way). --- pimm/core.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/pimm/core.py b/pimm/core.py index 5a3be162b..6e4e3d6ab 100644 --- a/pimm/core.py +++ b/pimm/core.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass -from typing import Any, Generic, TypeVar, cast, final +from typing import Generic, TypeVar, cast, final T = TypeVar('T') U = TypeVar('U') @@ -11,7 +11,11 @@ class NoValueException(Exception): pass -NODEFAULT: Any = object() +class _NoDefault: + """No default at all, which `T | None` cannot express: `None` is itself a value a receiver may default to.""" + + +NODEFAULT = _NoDefault() @dataclass @@ -164,7 +168,7 @@ def emit(self, data: T, ts: int = -1): class ControlSystemReceiver(SignalReceiver[T]): """Receiver adaptor bound to a single upstream signal on behalf of a system.""" - def __init__(self, owner: ControlSystem, default: T | None = NODEFAULT, maxsize: int | None = None): + def __init__(self, owner: ControlSystem, default: T | None | _NoDefault = NODEFAULT, maxsize: int | None = None): self._owner = owner self._default = default self._internal: SignalReceiver[T] | None = None @@ -187,8 +191,8 @@ def read(self) -> Message[T] | None: value = self._internal.read() if value is not None: return value - if self._default is not NODEFAULT: - # Always not-updated; the check above excludes the sentinel, which `T | None` cannot express. + if not isinstance(self._default, _NoDefault): + # Always not-updated; `None` reaches here as the value it was given, not as an absent default. return Message(cast(T, self._default), -1, False) return None @@ -227,7 +231,9 @@ class ReceiverDict(dict[str, ControlSystemReceiver[U]]): Pass fake=True for all fake receivers, or fake={'key1', 'key2'} for specific keys. """ - def __init__(self, owner: ControlSystem, *, default: U | None = NODEFAULT, fake: bool | Iterable[str] = False): + def __init__( + self, owner: ControlSystem, *, default: U | None | _NoDefault = NODEFAULT, fake: bool | Iterable[str] = False + ): super().__init__() self._owner = owner self._default = default From fc1d010b4ce2d3c6b2a45a2f51fd7778ccce5788 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 13 Aug 2026 11:57:30 +0300 Subject: [PATCH 21/64] Name the action-timestamp and grip channels by their `keys` constants in tests --- positronic/policy/tests/test_wrappers.py | 34 ++++++++++++------------ positronic/tests/test_data_collection.py | 6 ++--- 2 files changed, 20 insertions(+), 20 deletions(-) diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 30aceb527..8e041c611 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -69,14 +69,14 @@ def _obs(now_sec=0.0, fault=False): class TestStopOnFault: def test_a_faulted_arm_stops_what_is_executing(self): - inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) session = StopOnFault().wrap_session(inner, None, None) assert session(_obs(0.0, fault=True)) == [] assert inner.call_count == 0, 'the model was asked about an arm that is not tracking it' def test_a_sound_arm_reaches_the_model(self): - inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) session = StopOnFault().wrap_session(inner, None, None) assert session(_obs(0.0)) is not None @@ -84,7 +84,7 @@ def test_a_sound_arm_reaches_the_model(self): def test_an_observation_with_no_flag_is_not_a_fault(self): """Only the harness stamps the flag; a probe replaying a recording has no arm to fault.""" - inner = _ConstSession([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstSession([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) session = StopOnFault().wrap_session(inner, None, None) assert session({keys.OBS_TIME_NS: 0}) is not None @@ -93,7 +93,7 @@ def test_an_observation_with_no_flag_is_not_a_fault(self): def test_recovery_plans_afresh_instead_of_resuming(self): """The fault resets the scheduler below it, so the first sound observation infers again rather than waiting out the chunk stamped before the fault.""" - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 1.0}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 1.0}]) session = (StopOnFault() | ChunkedSchedule()).wrap(inner).new_session(now=_FakeClock(t=0.0).now) assert session(_obs(0.0)) is not None # a chunk that runs until 1.0 @@ -105,20 +105,20 @@ class TestChunkedSchedule: def test_first_call_runs_inference(self): # Relative timestamps: trajectory of duration 0.5s clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) policy = ChunkedSchedule().wrap(inner) session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None assert len(result) == 2 # Timestamps stamped to absolute by ChunkedSchedule. - assert result[0]['timestamp'] == 1.0 - assert result[1]['timestamp'] == 1.5 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 + assert result[1][keys.ACTION_TIMESTAMP] == 1.5 def test_returns_none_while_trajectory_active(self): # Trajectory starts at clock=1.0, ends at 1.0+0.5=1.5. clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) policy = ChunkedSchedule().wrap(inner) session = policy.new_session(now=clock.now) session(_obs()) @@ -129,7 +129,7 @@ def test_returns_none_while_trajectory_active(self): def test_re_infers_after_trajectory_consumed(self): clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) session = ChunkedSchedule().wrap(inner).new_session(now=clock.now) session(_obs(1.0)) # trajectory ends at clock=1.5 assert session(_obs(1.3)) is None @@ -141,7 +141,7 @@ def test_re_infers_after_trajectory_consumed(self): def test_single_action_refires_immediately_after(self): """Single action at ts=0 → trajectory_end = now → next tick re-infers.""" clock = _FakeClock(t=1.0) - policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = ChunkedSchedule().wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) session(_obs(1.0)) clock.t = 1.01 @@ -150,7 +150,7 @@ def test_single_action_refires_immediately_after(self): def test_expiry_is_judged_at_the_observation_instant(self): """Whether the trajectory has run out is a question about the observation, not about ``now``.""" - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}, {'v': 2, 'timestamp': 0.5}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}, {'v': 2, keys.ACTION_TIMESTAMP: 0.5}]) session = ChunkedSchedule().wrap(inner).new_session(now=_FakeClock(t=2.0).now) session(_obs(1.0)) # anchored at now()=2.0, so the trajectory ends at 2.5 assert session(_obs(2.4)) is None @@ -164,7 +164,7 @@ def test_wrapper_pipe_wrapper(self): clock = _FakeClock(t=1.0) pipeline = TemporalStack(keys=('v',), offsets_sec=(0.0,)) | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy([{'v': 1, 'timestamp': 0.0}])) + policy = pipeline.wrap(_ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) result = session({keys.OBS_TIME_NS: int(1e9), 'v': np.array([5.0])}) assert result is not None @@ -175,7 +175,7 @@ def test_codec_pipe_wrapper(self): codec = ActionTimestamp(fps=10.0) pipeline = codec | ChunkedSchedule() assert isinstance(pipeline, PolicyWrapper) - policy = pipeline.wrap(_ConstPolicy([{'action': 'test', 'timestamp': 0.0}])) + policy = pipeline.wrap(_ConstPolicy([{'action': 'test', keys.ACTION_TIMESTAMP: 0.0}])) session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None @@ -191,7 +191,7 @@ def test_full_pipeline(self): session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None - assert result[0]['timestamp'] == 1.0 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 # Second call within trajectory window returns None (ChunkedSchedule). clock.t = 1.2 assert session(_obs()) is None @@ -504,17 +504,17 @@ def test_inline_full_pipe(self): session = policy.new_session(now=clock.now) result = session(_obs()) assert result is not None - assert result[0]['timestamp'] == 1.0 + assert result[0][keys.ACTION_TIMESTAMP] == 1.0 clock.t = 1.2 assert session(_obs()) is None def test_inline_tolerates_marker_less_pipe(self): clock = _FakeClock(t=1.0) - inner = _ConstPolicy([{'v': 1, 'timestamp': 0.0}]) + inner = _ConstPolicy([{'v': 1, keys.ACTION_TIMESTAMP: 0.0}]) policy = spec.inline(ChunkedSchedule() | spec.PolicySource(inner)) session = policy.new_session(now=clock.now) result = session(_obs()) - assert result is not None and result[0]['timestamp'] == 1.0 + assert result is not None and result[0][keys.ACTION_TIMESTAMP] == 1.0 def test_inline_bare_source_pipe_is_the_loaded_policy(self): inner = _ConstPolicy([]) diff --git a/positronic/tests/test_data_collection.py b/positronic/tests/test_data_collection.py index 56a1c1ded..303d1f1ce 100644 --- a/positronic/tests/test_data_collection.py +++ b/positronic/tests/test_data_collection.py @@ -157,7 +157,7 @@ def test_data_collection_with_mujoco_robot_gripper(tmp_path): writer_cm = LocalDatasetWriter(tmp_path) agent = DsWriterAgent(writer_cm.__enter__()) - agent.add_signal('target_grip') + agent.add_signal(keys.TARGET_GRIP) agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) agent.add_signal('controller_positions', controller_positions_serializer) agent.add_signal('robot_state', Serializers.robot_state) @@ -168,7 +168,7 @@ def test_data_collection_with_mujoco_robot_gripper(tmp_path): world.connect(dc.robot_commands, sim.commands) world.connect(dc.robot_commands, agent.inputs[keys.ROBOT_COMMAND]) world.connect(dc.target_grip, sim.target_grip) - world.connect(dc.target_grip, agent.inputs['target_grip']) + world.connect(dc.target_grip, agent.inputs[keys.TARGET_GRIP]) world.connect(sim.grip, agent.inputs[keys.GRIP]) world.connect(dc.ds_agent_commands, agent.command) @@ -208,7 +208,7 @@ def stop_episode(): ep = ds[0] assert isinstance(ep, Episode) - expected = {'target_grip', 'controller_positions.right', keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP} + expected = {keys.TARGET_GRIP, 'controller_positions.right', keys.JOINTS, keys.JOINT_VEL, keys.EE_POSE, keys.GRIP} assert expected.issubset(set(ep.keys())) # Robot/gripper signals should have at least one sample From f5a15c25bc9b5086a2b5a2a755711678ac85b9f4 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 13 Aug 2026 16:31:54 +0300 Subject: [PATCH 22/64] State that recording is best effort, and why a sample can go missing --- positronic/dataset/README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/positronic/dataset/README.md b/positronic/dataset/README.md index e31623873..376670a69 100644 --- a/positronic/dataset/README.md +++ b/positronic/dataset/README.md @@ -354,6 +354,7 @@ Each line of `edits.jsonl` is one JSON record carrying its op. `{"op": "set_stat Key ideas - Inputs are registered explicitly through `DsWriterAgent.add_signal(name, serializer=None)`. - The agent polls inputs at a configurable rate and appends only on updates. +- Recording is best effort, and this is a deliberate trade rather than an oversight. Each input arrives over a one-slot `pimm` signal where a new value overwrites one still unread, so a recorder that stalls for longer than the gap between two samples loses the older one — commands exactly as much as camera frames or arm state. An episode is what the recorder managed to observe, not a guaranteed-complete log of what happened; treat a missing sample as possible in any analysis that counts them. - A separate `command` channel controls episode lifecycle. - `time_mode` selects how timestamps are recorded: `CLOCK` (default) stamps samples when the agent ingests them (useful during live data collection so every signal reflects when the recorder could act on it), while `MESSAGE` preserves the timestamp attached by the emitting control system (ideal for analysing inference latency by keeping original emission times). From 602fd915c884c3244bdc621b86cac2ca0dc53696 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 13 Aug 2026 16:41:53 +0300 Subject: [PATCH 23/64] Reap the abandoned call before a new episode opens a session An in-process policy is one model across episodes, so `new_session` resets the object a call abandoned by the last episode may still be inside. The wait sits after the task reset, which usually covers it, rather than at the end of the episode a hung model must not hold up. --- positronic/policy/harness.py | 20 +++++++++-- positronic/policy/tests/test_harness.py | 47 +++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 1d0610ae8..362a96700 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -227,6 +227,7 @@ def __init__( # worker belongs to the episode: ending one abandons the call in flight rather than waiting for it, # so the next episode must not queue behind it. self._executor: ThreadPoolExecutor | None = None + self._retiring: ThreadPoolExecutor | None = None self._future: Future[list[dict[str, Any]] | None] | None = None # The in-flight call's start: the world instant its observation was built, and the wall instant it # was submitted. @@ -321,13 +322,27 @@ def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: def _retire_worker(self) -> None: """Let go of this episode's worker and the call it is running: the answer lands nowhere and the - failure only reaches the log.""" + failure only reaches the log. The worker is kept for ``_reap_worker`` to join at the next episode's + start, since ending an episode must not wait for a model that hangs.""" if self._future is not None: self._future.add_done_callback(self._report_abandoned) self._future = None if self._executor is not None: self._executor.shutdown(wait=False, cancel_futures=True) - self._executor = None + self._retiring, self._executor = self._executor, None + + def _reap_worker(self) -> None: + """Wait out the previous episode's abandoned call before this one opens a session. + + An in-process policy is a single model across episodes, so ``new_session`` resets the very object an + abandoned call may still be inside — a running thread survives ``shutdown(cancel_futures=True)``, + which cancels only what is still queued. Waiting here rather than at the end of the episode it + belongs to keeps a hung model from holding up that episode's recording and home, and the task reset + this follows usually covers the wait. + """ + if self._retiring is not None: + self._retiring.shutdown(wait=True) + self._retiring = None def _finalize_recording( self, clock: pimm.Clock, payload: dict[str, Any] | None = None @@ -375,6 +390,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} + self._reap_worker() # Arm the clock before handing it out: a session reading it before its first call must see this # episode's start, not the release time of the last episode's final call. self._t0_ns = clock.now_ns() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 3581ac7c3..4d5b8eb47 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1014,6 +1014,53 @@ def test_run_while_running_is_ignored(world): assert policy.reset_calls == 1 +class _AbandonedCallPolicy(Policy): + """Records the order of session openings and call completions, with a call that outlives its episode.""" + + def __init__(self, wall_sec: float): + self._wall_sec = wall_sec + self.events: list[str] = [] + + def new_session(self, context=None, now=None): + self.events.append('open') + return _AbandonedCallPolicy._Session(self) + + class _Session(Session): + def __init__(self, policy: '_AbandonedCallPolicy'): + self._policy = policy + + def __call__(self, obs): + time.sleep(self._policy._wall_sec) + self._policy.events.append('answered') + return None + + +@pytest.mark.timeout(10.0) +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 call + still inside the previous one — ``new_session`` resets the object that call is using.""" + policy = _AbandonedCallPolicy(wall_sec=0.4) + 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]) + emit_obs = partial(emit_ready_payload, p['frame_em'], p['robot_em'], p['grip_em'], robot_state) + + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='ep1')), 0.0), + (emit_obs, 0.01), # the observation that puts a call on the worker + (partial(p['directive_em'].emit, Directive.FINISH()), 0.02), # while that call is still running + (partial(p['directive_em'].emit, Directive.RUN(task='ep2')), 0.02), + (None, 0.02), + ]) + + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=40) + + assert policy.events[:3] == ['open', 'answered', 'open'], ( + f'the second session opened before the abandoned call answered: {policy.events}' + ) + + @pytest.mark.timeout(3.0) def test_run_calls_policy_reset_with_context(world): policy = StubPolicy() From db1fdf752c57d5c5670eab053fc951e594a171d0 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Thu, 13 Aug 2026 18:23:07 +0300 Subject: [PATCH 24/64] Hand the worker its own copy of the observation's arrays A camera renders into the array behind the adapter it re-emits, so a producer advancing while a wall-charged call runs rewrites what the worker is still reading. The loop thread yields there, which is what makes it reachable. --- .../dataset/tests/test_ds_writer_agent.py | 4 +- positronic/policy/harness.py | 13 ++++- positronic/policy/tests/test_harness.py | 58 ++++++++++++++++++- 3 files changed, 69 insertions(+), 6 deletions(-) diff --git a/positronic/dataset/tests/test_ds_writer_agent.py b/positronic/dataset/tests/test_ds_writer_agent.py index ecaa952fd..b214ba347 100644 --- a/positronic/dataset/tests/test_ds_writer_agent.py +++ b/positronic/dataset/tests/test_ds_writer_agent.py @@ -476,12 +476,12 @@ def test_pickles_with_every_constructor_argument_filled(): virtual_time=True, telemetry_span=partial(telemetry.span, telemetry_keys.SPAN_RECORD_IO), ) - agent.add_signal('robot_command', Serializers.robot_command) + agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) agent.add_signal('robot_state', Serializers.robot_state) loaded = pickle.loads(pickle.dumps(agent)) - assert set(loaded.inputs) == {'robot_command', 'robot_state'} + assert set(loaded.inputs) == {keys.ROBOT_COMMAND, 'robot_state'} with loaded._telemetry_span(): pass diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 362a96700..e723ce74e 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -8,6 +8,7 @@ from enum import Enum from typing import Any, TypeAlias +import numpy as np from opentelemetry.trace import Span import pimm @@ -66,6 +67,16 @@ def advance(self, current_time: int): return value +def _owned(obs: dict[str, Any]) -> dict[str, Any]: + """The observation with its arrays copied, so nothing rewrites what the worker is still reading. + + A producer may reuse one buffer for every sample it emits — a camera renders into the array behind the + adapter it re-emits each frame — and the loop thread yields while a call charged in wall time runs, so + that producer advances alongside the worker. Copying at dispatch pays once per call rather than per round. + """ + return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} + + def _assert_anchored(actions: 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 actions), default=0.0) @@ -520,7 +531,7 @@ def _step(self, clock: pimm.Clock) -> None: self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. - self._future = executor.submit(session, frozen_view(obs)) # pyright: ignore[reportArgumentType] + self._future = executor.submit(session, frozen_view(_owned(obs))) # pyright: ignore[reportArgumentType] if self._charge is None: # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 4d5b8eb47..a9df99158 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1014,6 +1014,58 @@ def test_run_while_running_is_ignored(world): assert policy.reset_calls == 1 +class _FrameWatchingPolicy(Policy): + """Reads its camera frame at both ends of a slow call, so a rewrite underneath it shows up as a difference.""" + + def __init__(self, wall_sec: float): + self._wall_sec = wall_sec + self.seen: list[tuple[np.ndarray, np.ndarray]] = [] + + def new_session(self, context=None, now=None): + return _FrameWatchingPolicy._Session(self) + + class _Session(Session): + def __init__(self, policy: '_FrameWatchingPolicy'): + self._policy = policy + + def __call__(self, obs): + entry = np.array(obs[CAM]) + time.sleep(self._policy._wall_sec) + self._policy.seen.append((entry, np.array(obs[CAM]))) + return None + + +@pytest.mark.timeout(10.0) +def test_a_producer_reusing_its_buffer_cannot_rewrite_a_pending_observation(world): + """A camera renders into the array behind the adapter it re-emits, and a wall-charged call runs while the + loop yields — so the observation handed to the worker has to be its own copy.""" + policy = _FrameWatchingPolicy(wall_sec=0.3) + 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]) + frame = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) + + def emit_frame(fill: int): + frame.array[:] = np.full((2, 2, 3), fill, dtype=np.uint8) + p['frame_em'].emit(frame) # the same adapter every time, as a camera does + p['robot_em'].emit(robot_state) + p['grip_em'].emit(0.25) + + driver = ManualDriver([ + (partial(p['directive_em'].emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: True})), 0.0), + (partial(emit_frame, 1), 0.01), + (partial(emit_frame, 9), 0.05), # rewrites the buffer while the first call is still running + (None, 0.4), + ]) + + scheduler = world.start([harness, driver]) + drive_scheduler(scheduler, steps=60) + + assert policy.seen, 'the policy was never called' + entry, exit_ = policy.seen[0] + np.testing.assert_array_equal(entry, exit_, 'the observation was rewritten while the call was in flight') + + class _AbandonedCallPolicy(Policy): """Records the order of session openings and call completions, with a call that outlives its episode.""" @@ -1703,7 +1755,7 @@ def _run_episode( robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([ - (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=latency)), 0.0), + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: latency})), 0.0), (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), (None, run_sec), ]) @@ -1882,7 +1934,7 @@ def test_a_stop_lands_without_waiting_out_the_charge(world): pose, joints = [0.1, 0.2, 0.3], [0.4, 0.5, 0.6] driver = ManualDriver([ - (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=charge)), 0.0), + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: charge})), 0.0), (partial(emit_ready_payload, frame_em, robot_em, grip_em, make_robot_state(pose, joints)), fault_at), (partial(robot_em.emit, make_robot_state(pose, joints, status=RobotStatus.ERROR)), charge + 0.2), ]) @@ -2028,7 +2080,7 @@ def test_abort_discards_a_call_that_is_still_in_flight(world): robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([ - (partial(directive_em.emit, Directive.RUN(task='t', inference_latency=1.0)), 0.0), + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: 1.0})), 0.0), (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), (None, 0.05), # well inside the 1.0s the gate owes the call (partial(directive_em.emit, Directive.ABORT()), 0.0), From 3b24948da74de9732b1c78efe2ed2ecce0efbf2a Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Fri, 14 Aug 2026 20:25:15 +0300 Subject: [PATCH 25/64] Drop the stale entry count from the wire-table waiver --- positronic/policy/spec.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/positronic/policy/spec.py b/positronic/policy/spec.py index 6e66a02f8..fbd9aa99e 100644 --- a/positronic/policy/spec.py +++ b/positronic/policy/spec.py @@ -146,8 +146,8 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) # and once as a key below — so a rename in one place silently desyncs the wire. Let each wrapper own its wire # name as a class attribute and build this mapping from it. ``test_wire_names_match_table`` is what catches a # desync until then. -# rules-allow: hardcoded-keys — fixing one of the thirteen leaves the rest on the old pattern; the TODO above -# names the whole-table fix, and the test above catches a desync meanwhile +# rules-allow: hardcoded-keys — fixing one entry leaves the rest on the old pattern; the TODO above names the +# whole-table fix, and the test above catches a desync meanwhile WIRE_WRAPPERS: dict[str, type[PolicyWrapper]] = { 'chunked_schedule': ChunkedSchedule, 'stop_on_fault': StopOnFault, From f72287389d70f36c4d0599d923a67302a7b23c86 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 15 Aug 2026 21:19:39 +0300 Subject: [PATCH 26/64] Type a session's observation as the read-only mapping it receives --- positronic/offboard/tests/test_remote_policy.py | 4 +++- positronic/policy/base.py | 4 ++-- positronic/policy/harness.py | 2 +- positronic/policy/remote.py | 6 +++--- 4 files changed, 9 insertions(+), 7 deletions(-) diff --git a/positronic/offboard/tests/test_remote_policy.py b/positronic/offboard/tests/test_remote_policy.py index 6a3b2c417..596847b91 100644 --- a/positronic/offboard/tests/test_remote_policy.py +++ b/positronic/offboard/tests/test_remote_policy.py @@ -57,7 +57,9 @@ class TestPrepareObs: def test_images_pass_through_untouched_by_default(self): session = RemoteSession(_mock_ws_session()) obs = {'cam': _make_image(480, 640), 'state': np.array([1.0])} - assert session._prepare_obs(obs) is obs + prepared = session._prepare_obs(obs) + assert prepared.keys() == obs.keys() + assert all(prepared[key] is value for key, value in obs.items()) def test_compression_reaches_nested_images(self): session = RemoteSession(_mock_ws_session(), compress_images=True) diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 56fbc36d6..46dd2eb6d 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Any Now = Callable[[], float] @@ -33,7 +33,7 @@ class Session(ABC): """ @abstractmethod - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: Mapping[str, Any]) -> list[dict[str, Any]] | None: """Predict actions for the given observation.""" @property diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index e723ce74e..0adb5626b 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -531,7 +531,7 @@ def _step(self, clock: pimm.Clock) -> None: self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. - self._future = executor.submit(session, frozen_view(_owned(obs))) # pyright: ignore[reportArgumentType] + self._future = executor.submit(session, frozen_view(_owned(obs))) if self._charge is None: # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 142539a0d..cb5174436 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -26,9 +26,9 @@ def __init__(self, ws_session: InferenceSession, compress_images: bool = False): self._session = ws_session self._compress_images = compress_images - def _prepare_obs(self, obs: dict[str, Any]) -> dict[str, Any]: + def _prepare_obs(self, obs: cabc.Mapping[str, Any]) -> dict[str, Any]: if not self._compress_images: - return obs + return dict(obs) return {key: self._prepare_value(key, value) for key, value in obs.items()} def _prepare_value(self, key: str, value: Any) -> Any: @@ -42,7 +42,7 @@ def _prepare_value(self, key: str, value: Any) -> Any: return type(value)(self._prepare_value(key, v) for v in value) return value - def __call__(self, obs: dict[str, Any]) -> list[dict[str, Any]] | None: + def __call__(self, obs: cabc.Mapping[str, Any]) -> list[dict[str, Any]] | None: """Forwards the observation to the remote server and returns the action trajectory. Single-action server responses are wrapped into a 1-element list to honor From 3ccdcf9b585bec8681c3e999d527b2f1a0bd6bf5 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 15 Aug 2026 21:26:49 +0300 Subject: [PATCH 27/64] Name the robot-state channel by a `keys` constant --- positronic/cfg/embodiment.py | 19 ++++++------ .../dataset/tests/test_ds_writer_agent.py | 10 +++---- positronic/keys.py | 13 ++++---- .../policy/tests/test_golden_pipeline.py | 2 +- positronic/policy/tests/test_harness.py | 30 +++++++++---------- positronic/simulator/env_server/proxy.py | 2 +- .../simulator/env_server/tests/mujoco_env.py | 2 +- .../env_server/tests/test_remote_env.py | 2 +- positronic/simulator/libero/adapter.py | 2 +- positronic/simulator/robolab/adapter.py | 2 +- positronic/tests/test_data_collection.py | 4 +-- positronic/wire.py | 4 +-- 12 files changed, 48 insertions(+), 44 deletions(-) diff --git a/positronic/cfg/embodiment.py b/positronic/cfg/embodiment.py index 29e5116af..42876b3e7 100644 --- a/positronic/cfg/embodiment.py +++ b/positronic/cfg/embodiment.py @@ -25,7 +25,7 @@ def droid(robot_arm, gripper, cameras): """Real single-arm Franka (DROID) + Robotiq gripper + ZED cameras.""" observations = { - 'robot_state': Observation(robot_arm.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot_arm.state, Serializers.robot_state), keys.GRIP: Observation(gripper.grip, None), **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } @@ -48,7 +48,7 @@ def droid(robot_arm, gripper, cameras): def yam(robot_arm, cameras): """Real single-arm i2rt YAM: the arm driver carries the gripper (they share one CAN chain).""" observations = { - 'robot_state': Observation(robot_arm.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot_arm.state, Serializers.robot_state), keys.GRIP: Observation(robot_arm.grip, None), **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } @@ -96,21 +96,22 @@ def yam_bimanual(left_channel: str, right_channel: str, mounts: dict[str, list[f for side, channel in (('left', left_channel), ('right', right_channel)) } observations = { - **{f'robot_state.{s}': Observation(arm.state, Serializers.robot_state) for s, arm in arms.items()}, - **{f'grip.{s}': Observation(arm.grip, None) for s, arm in arms.items()}, + **{f'{keys.ROBOT_STATE}.{s}': Observation(arm.state, Serializers.robot_state) for s, arm in arms.items()}, + **{f'{keys.GRIP}.{s}': Observation(arm.grip, None) for s, arm in arms.items()}, **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { **{ - f'robot_command.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) + f'{keys.ROBOT_COMMAND}.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) for s, arm in arms.items() }, - **{f'target_grip.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, + **{f'{keys.TARGET_GRIP}.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, } - joint_signals = {side: f'robot_state.{side}.q' for side in arms} + joint_signals = {side: f'{keys.ROBOT_STATE}.{side}.q' for side in arms} static_meta = { keys.JOINT_SIGNALS: list(joint_signals.values()), - keys.POSE_SIGNALS: [f'robot_state.{s}.ee_pose' for s in arms] + [f'robot_command.{s}.pose' for s in arms], + keys.POSE_SIGNALS: [f'{keys.ROBOT_STATE}.{s}.ee_pose' for s in arms] + + [f'{keys.ROBOT_COMMAND}.{s}.pose' for s in arms], keys.MOUNTS: {sig: mounts[side] for side, sig in joint_signals.items()}, } return Embodiment( @@ -133,7 +134,7 @@ def mujoco_franka(sim, camera_dict): Mujoco does not render the second image when using only 2 cameras. """ observations = { - 'robot_state': Observation(sim.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(sim.state, Serializers.robot_state), keys.GRIP: Observation(sim.grip, None), **{name: Observation(sim.cameras[orig], Serializers.camera_images) for name, orig in camera_dict.items()}, } diff --git a/positronic/dataset/tests/test_ds_writer_agent.py b/positronic/dataset/tests/test_ds_writer_agent.py index b214ba347..fc2a45c87 100644 --- a/positronic/dataset/tests/test_ds_writer_agent.py +++ b/positronic/dataset/tests/test_ds_writer_agent.py @@ -376,7 +376,7 @@ def status(self): def test_robot_state_serializer_drops_reset_and_emits_components(world): ds = FakeDatasetWriter() - agent, cmd_em, emitters = build_agent_with_pipes({'robot_state': Serializers.robot_state}, ds, world) + agent, cmd_em, emitters = build_agent_with_pipes({keys.ROBOT_STATE: Serializers.robot_state}, ds, world) q = np.arange(7, dtype=np.float32) dq = np.arange(7, dtype=np.float32) + 10 @@ -385,8 +385,8 @@ def test_robot_state_serializer_drops_reset_and_emits_components(world): script = [ (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.START_EPISODE)), 0.001), - (lambda: emitters['robot_state'].emit(_FakeState(q, dq, pose, RobotStatus.RESETTING)), 0.001), - (lambda: emitters['robot_state'].emit(_FakeState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(_FakeState(q, dq, pose, RobotStatus.RESETTING)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(_FakeState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.STOP_EPISODE)), 0.001), ] @@ -477,11 +477,11 @@ def test_pickles_with_every_constructor_argument_filled(): telemetry_span=partial(telemetry.span, telemetry_keys.SPAN_RECORD_IO), ) agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) - agent.add_signal('robot_state', Serializers.robot_state) + agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) loaded = pickle.loads(pickle.dumps(agent)) - assert set(loaded.inputs) == {keys.ROBOT_COMMAND, 'robot_state'} + assert set(loaded.inputs) == {keys.ROBOT_COMMAND, keys.ROBOT_STATE} with loaded._telemetry_span(): pass diff --git a/positronic/keys.py b/positronic/keys.py index 93e795e7e..a9149e59b 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -26,13 +26,16 @@ def is_robot_command(name: str) -> bool: return name == ROBOT_COMMAND or name.startswith(f'{ROBOT_COMMAND}.') -JOINTS = 'robot_state.q' -JOINT_VEL = 'robot_state.dq' -EE_POSE = 'robot_state.ee_pose' +# The arm's state channel, and the signals a recorded state unfolds into. As on the command side, the +# suffixes are ``Serializers.robot_state``'s, so the names derive from the channel rather than restating it. +ROBOT_STATE = 'robot_state' +JOINTS = f'{ROBOT_STATE}.q' +JOINT_VEL = f'{ROBOT_STATE}.dq' +EE_POSE = f'{ROBOT_STATE}.ee_pose' # Whether the arm is faulted, in every observation the harness builds. A faulted arm has no sample to give # and is not tracking the plan it was handed, so the policy stack — not the harness — decides what happens -# next: the rest of ``robot_state`` is absent from that observation. -ROBOT_FAULT = 'robot_state.fault' +# next: the rest of ``ROBOT_STATE`` is absent from that observation. +ROBOT_FAULT = f'{ROBOT_STATE}.fault' GRIP = 'grip' TASK = 'task' # The prefix that identifies a camera on the wire: an embodiment declares its cameras by naming diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index d073e2eb0..a4f6bccee 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -185,7 +185,7 @@ def _run_pipeline(tmp_path: Path) -> dict: embodiment = Embodiment( descriptor='', observations={ - 'robot_state': Observation(robot.state, Serializers.robot_state), + keys.ROBOT_STATE: Observation(robot.state, Serializers.robot_state), keys.GRIP: Observation(gripper.grip, None), }, commands={ diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index a9df99158..040f7a8ef 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -50,7 +50,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simu home values, descriptor — is read by the Harness. """ observations = { - 'robot_state': Observation(pimm.NoOpEmitter(), Serializers.robot_state), + keys.ROBOT_STATE: Observation(pimm.NoOpEmitter(), Serializers.robot_state), keys.GRIP: Observation(pimm.NoOpEmitter(), None), } for cam in cameras: @@ -257,7 +257,7 @@ def _pair_all(world, harness): harness.ds_command._bind(ds_recorder) return { 'frame_em': world.pair(harness.observations[CAM]), - 'robot_em': world.pair(harness.observations['robot_state']), + 'robot_em': world.pair(harness.observations[keys.ROBOT_STATE]), 'grip_em': world.pair(harness.observations[keys.GRIP]), 'directive_em': world.pair(harness.directive), 'command_rx': world.pair(harness.commands[keys.ROBOT_COMMAND]), @@ -311,7 +311,7 @@ def test_harness_emits_cartesian_move(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -370,7 +370,7 @@ def test_harness_passes_descriptor_to_policy(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -400,7 +400,7 @@ def test_robot_model_stays_out_of_the_observation(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations['grip']) directive_em = world.pair(harness.directive) @@ -471,7 +471,7 @@ def test_harness_waits_for_complete_inputs(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -948,7 +948,7 @@ def test_timeout_during_inference_drops_the_chunk(world): harness.ds_command._bind(ds_recorder) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) @@ -1172,7 +1172,7 @@ def test_finish_stops_playing_the_live_chunk(world): harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1214,7 +1214,7 @@ def new_session(self, context=None, now=None): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1396,7 +1396,7 @@ def test_shutdown_stops_playing_the_live_chunk(world): harness.ds_command._bind(_LabeledRecorder('ds_command', events)) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1749,7 +1749,7 @@ def _run_episode( harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1928,7 +1928,7 @@ def test_a_stop_lands_without_waiting_out_the_charge(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -1973,7 +1973,7 @@ def new_session(self, context=None, now=None): harness.ds_command._bind(ds_recorder) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) @@ -2026,7 +2026,7 @@ def new_session(self, context=None, now=None): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) @@ -2074,7 +2074,7 @@ def test_abort_discards_a_call_that_is_still_in_flight(world): harness.ds_command._bind(RecordingEmitter()) frame_em = world.pair(harness.observations[CAM]) - robot_em = world.pair(harness.observations['robot_state']) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) grip_em = world.pair(harness.observations[keys.GRIP]) directive_em = world.pair(harness.directive) diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index 55ed9d624..4896cf99c 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -153,7 +153,7 @@ def remote_franka_embodiment( server cannot import positronic to emit it via ``robot_meta``). """ observations = { - 'robot_state': Observation(proxy.observations['robot_state'], Serializers.robot_state), + keys.ROBOT_STATE: Observation(proxy.observations[keys.ROBOT_STATE], Serializers.robot_state), keys.GRIP: Observation(proxy.observations[keys.GRIP], None), **{logical: Observation(proxy.observations[logical], Serializers.camera_images) for logical in camera_dict}, } diff --git a/positronic/simulator/env_server/tests/mujoco_env.py b/positronic/simulator/env_server/tests/mujoco_env.py index 416ce45c8..1a6e29e03 100644 --- a/positronic/simulator/env_server/tests/mujoco_env.py +++ b/positronic/simulator/env_server/tests/mujoco_env.py @@ -171,7 +171,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = geom.Transform3D(raw_obs['ee_pos'], geom.Rotation.from_quat(raw_obs['ee_quat'])) state.encode(raw_obs['q'], raw_obs['dq'], ee_pose) state.array[14 + 7] = float(raw_obs['status']) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} for logical, model_name in self._camera_dict.items(): frame = raw_obs['cameras'][model_name] adapter = pimm.shared_memory.NumpySMAdapter(shape=frame.shape, dtype=frame.dtype) diff --git a/positronic/simulator/env_server/tests/test_remote_env.py b/positronic/simulator/env_server/tests/test_remote_env.py index 61729b9da..cb37caff2 100644 --- a/positronic/simulator/env_server/tests/test_remote_env.py +++ b/positronic/simulator/env_server/tests/test_remote_env.py @@ -183,7 +183,7 @@ def test_robolab_reports_and_drives_the_same_frame(self): 'joint_vel': np.zeros(7), 'grip': 0.0, } - reported = adapter.observations(raw)['robot_state'].ee_pose + reported = adapter.observations(raw)[keys.ROBOT_STATE].ee_pose commanded = _in_env_control_frame(roboarm_command.CartesianPosition(reported), adapter.env_control_frame).pose np.testing.assert_allclose(commanded.as_vector(self.rotmat), eef.as_vector(self.rotmat), atol=1e-6) diff --git a/positronic/simulator/libero/adapter.py b/positronic/simulator/libero/adapter.py index acb88ceac..02d2681bb 100644 --- a/positronic/simulator/libero/adapter.py +++ b/positronic/simulator/libero/adapter.py @@ -42,7 +42,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = geom.Transform3D(raw_obs['eef_pos'], geom.Rotation.from_quat_xyzw(raw_obs['eef_quat'])) state = MujocoFrankaState() state.encode(raw_obs['joint_pos'], raw_obs['joint_vel'], ee_pose) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} for logical, env_key in self._camera_dict.items(): # robosuite renders bottom-up; flip to standard top-down orientation (LIBERO's own video path # flips the same way). diff --git a/positronic/simulator/robolab/adapter.py b/positronic/simulator/robolab/adapter.py index 34a4413ab..e4afbda3e 100644 --- a/positronic/simulator/robolab/adapter.py +++ b/positronic/simulator/robolab/adapter.py @@ -34,7 +34,7 @@ def observations(self, raw_obs: dict[str, Any]) -> dict[str, Any]: ee_pose = eef_pose * self.env_control_frame.inv state = MujocoFrankaState() state.encode(raw_obs['joint_pos'], raw_obs['joint_vel'], ee_pose) - obs: dict[str, Any] = {'robot_state': state, keys.GRIP: float(raw_obs['grip'])} + obs: dict[str, Any] = {keys.ROBOT_STATE: state, keys.GRIP: float(raw_obs['grip'])} # TODO: honour a camera_dict naming any other RoboLab camera. env.py renders only the WRIST_LEFT # preset (over_shoulder_left + wrist) and hard-codes emitting those two, so a request for e.g. # over_shoulder_right_camera raises below. The full fix threads the requested set end-to-end: carry diff --git a/positronic/tests/test_data_collection.py b/positronic/tests/test_data_collection.py index 303d1f1ce..b84babab2 100644 --- a/positronic/tests/test_data_collection.py +++ b/positronic/tests/test_data_collection.py @@ -160,11 +160,11 @@ def test_data_collection_with_mujoco_robot_gripper(tmp_path): agent.add_signal(keys.TARGET_GRIP) agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) agent.add_signal('controller_positions', controller_positions_serializer) - agent.add_signal('robot_state', Serializers.robot_state) + agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) agent.add_signal(keys.GRIP) world.connect(sim.state, dc.robot_state) - world.connect(sim.state, agent.inputs['robot_state']) + world.connect(sim.state, agent.inputs[keys.ROBOT_STATE]) world.connect(dc.robot_commands, sim.commands) world.connect(dc.robot_commands, agent.inputs[keys.ROBOT_COMMAND]) world.connect(dc.target_grip, sim.target_grip) diff --git a/positronic/wire.py b/positronic/wire.py index 6fa0ef1c2..0be4dbe79 100644 --- a/positronic/wire.py +++ b/positronic/wire.py @@ -45,7 +45,7 @@ def wire( # noqa: C901 ds_agent.add_signal(signal_name, Serializers.camera_images) if robot_arm is not None: ds_agent.add_signal(keys.ROBOT_COMMAND, Serializers.robot_command) - ds_agent.add_signal('robot_state', Serializers.robot_state) + ds_agent.add_signal(keys.ROBOT_STATE, Serializers.robot_state) if gripper is not None: ds_agent.add_signal(keys.TARGET_GRIP) ds_agent.add_signal(keys.GRIP) @@ -54,7 +54,7 @@ def wire( # noqa: C901 world.connect(emitter, ds_agent.inputs[signal_name]) if robot_arm is not None: world.connect(harness.robot_commands, ds_agent.inputs[keys.ROBOT_COMMAND]) - world.connect(robot_arm.state, ds_agent.inputs['robot_state']) + world.connect(robot_arm.state, ds_agent.inputs[keys.ROBOT_STATE]) if gripper is not None: world.connect(harness.target_grip, ds_agent.inputs[keys.TARGET_GRIP]) world.connect(gripper.grip, ds_agent.inputs[keys.GRIP]) From 2fc58e67a0b8bd94cd1b423f514e48f9ca4bec00 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 15 Aug 2026 21:33:49 +0300 Subject: [PATCH 28/64] Let a fault on one arm outrank another arm's not-ready sample --- positronic/drivers/roboarm/so101/driver.py | 30 +++++---- positronic/policy/harness.py | 76 +++++++++++++--------- positronic/policy/tests/test_harness.py | 42 ++++++++++++ 3 files changed, 102 insertions(+), 46 deletions(-) diff --git a/positronic/drivers/roboarm/so101/driver.py b/positronic/drivers/roboarm/so101/driver.py index 35b6fdd5f..82a558670 100644 --- a/positronic/drivers/roboarm/so101/driver.py +++ b/positronic/drivers/roboarm/so101/driver.py @@ -62,6 +62,9 @@ def __init__(self, motor_bus: MotorBus, home_joints: list[float] | None = None): self.commands = pimm.ControlSystemReceiver[roboarm_command.CommandType](self) self.target_grip = pimm.ControlSystemReceiver[float](self) self._last_grip: float = 0.0 + # The arm half of the motor setpoint, in normalized units. ``None`` until the first arm command: + # a gripper target arriving before one has no arm position to pair with. + self._last_qpos: np.ndarray | None = None self.grip: pimm.SignalEmitter[float] = pimm.ControlSystemEmitter(self) self.state: pimm.SignalEmitter[SO101State] = pimm.ControlSystemEmitter(self) @@ -87,28 +90,27 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p while not should_stop.value: cmd_msg = self.commands.read() grip_msg = self.target_grip.read() - if grip_msg is not None and grip_msg.updated: - self._last_grip = grip_msg.data - if cmd_msg is not None and cmd_msg.updated: - match cmd_msg.data: + grip = grip_msg.data if grip_msg is not None and grip_msg.updated else None + command = cmd_msg.data if cmd_msg is not None and cmd_msg.updated else None + if grip is not None: + self._last_grip = grip + if command is not None: + match command: case roboarm_command.Reset(): raise NotImplementedError('Reset not implemented') case roboarm_command.CartesianPosition(pose): - qpos = self._solve_ik(state, pose) - q_with_gripper = np.concatenate([qpos, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) + self._last_qpos = self._solve_ik(state, pose) case roboarm_command.CartesianDelta() as delta_cmd: ee_pose, _ = self._forward_kinematics(self.motor_bus.position) - target = delta_cmd.apply(ee_pose) - qpos = self._solve_ik(state, target) - q_with_gripper = np.concatenate([qpos, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) + self._last_qpos = self._solve_ik(state, delta_cmd.apply(ee_pose)) case roboarm_command.JointPosition(qpos): - q_norm = self.rad_to_norm(qpos) - q_with_gripper = np.concatenate([q_norm, [self._last_grip]]) - self.motor_bus.set_target_position(q_with_gripper) + self._last_qpos = self.rad_to_norm(qpos) case other: raise ValueError(f'Unknown command: {other}') + # The arm and the gripper are one setpoint on a shared bus, but they arrive as two channels that + # need not carry a value in the same round, so either one changing rewrites the whole vector. + if (command is not None or grip is not None) and self._last_qpos is not None: + self.motor_bus.set_target_position(np.concatenate([self._last_qpos, [self._last_grip]])) q = self.motor_bus.position dq = self.motor_bus.velocity[:-1] diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 0adb5626b..32f7180d6 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -17,7 +17,7 @@ from positronic.dataset.serializers import expand_suffixed from positronic.drivers import roboarm from positronic.drivers.roboarm.ik import assert_default_frame -from positronic.eval import Embodiment, Task +from positronic.eval import Embodiment, Observation, Task from positronic.policy.base import Policy, Session from positronic.utils import flatten_dict, frozen_view @@ -67,16 +67,6 @@ def advance(self, current_time: int): return value -def _owned(obs: dict[str, Any]) -> dict[str, Any]: - """The observation with its arrays copied, so nothing rewrites what the worker is still reading. - - A producer may reuse one buffer for every sample it emits — a camera renders into the array behind the - adapter it re-emits each frame — and the loop thread yields while a call charged in wall time runs, so - that producer advances alongside the worker. Copying at dispatch pays once per call rather than per round. - """ - return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} - - def _assert_anchored(actions: 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 actions), default=0.0) @@ -453,6 +443,28 @@ def _is_faulted(value: Any) -> bool: """Whether a raw observation is an arm reporting a fault. Every other not-ready sample is simply absent.""" return isinstance(value, roboarm.State) and value.status is roboarm.RobotStatus.ERROR + def _read_channel(self, name: str, obs: Observation) -> tuple[dict[str, Any] | None, bool]: + """This channel's entries under their full names, and whether the arm behind it is faulted. + + The entries are ``None`` when the channel has no sample to give — a resetting or faulted arm alike. + Raises ``NoValueException`` before the channel has produced anything at all. + """ + message = self.observations[name].read() + if message is None: + raise pimm.NoValueException + if message.updated: + self._awaiting_obs.discard(name) + value = message.data + if obs.serializer is not None: + value = obs.serializer(value) + if value is None: + # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so the + # fault is recovered from the raw sample and stapled on by the caller as `keys.ROBOT_FAULT` — a + # name already claiming to be part of `robot_state`. Emit it from the serializer and this + # branch, the raw-type check and the flag all go, and the fault reaches the recording as well. + return None, self._is_faulted(message.data) + return {full: v for full, v in expand_suffixed(name, value) if v is not None}, False + def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: """Read every observation channel and assemble the policy input dict. @@ -469,27 +481,16 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: assert_default_frame(self._statics()) inputs: dict[str, Any] = {} faulted = False + not_ready = False for name, obs in self._embodiment.observations.items(): - message = self.observations[name].read() - if message is None: - raise pimm.NoValueException - if message.updated: - self._awaiting_obs.discard(name) - value = message.data - if obs.serializer is not None: - value = obs.serializer(value) - if value is None: - # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so the - # fault is recovered from the raw sample and stapled on below as `keys.ROBOT_FAULT` — a name - # already claiming to be part of `robot_state`. Emit it from the serializer and this branch, - # the raw-type check and the flag all go, and the fault reaches the recording as well. - if not self._is_faulted(message.data): - return None - faulted = True - continue - for full_name, v in expand_suffixed(name, value): - if v is not None: - inputs[full_name] = v + entries, channel_faulted = self._read_channel(name, obs) + faulted = faulted or channel_faulted + not_ready = not_ready or entries is None + inputs.update(entries or {}) + # Every channel is read before this decision, so a bimanual rig cannot hide one arm's fault behind + # another arm's not-ready sample: whichever channel comes first, the fault still reaches the stack. + if not_ready and not faulted: + return None if self._awaiting_obs: return None inputs[keys.ROBOT_FAULT] = faulted @@ -507,6 +508,17 @@ def _effect_time(self) -> float: charge = time.monotonic() - self._wall_t0 if self._charge is None else self._charge return self._t0_ns / 1e9 + charge + @staticmethod + def _owned(obs: dict[str, Any]) -> dict[str, Any]: + """The observation with its arrays copied, so nothing rewrites what the worker is still reading. + + A producer may reuse one buffer for every sample it emits — a camera renders into the array behind + the adapter it re-emits each frame — and the loop thread yields while a call charged in wall time + runs, so that producer advances alongside the worker. Copying at dispatch pays once per call rather + than per round. + """ + return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} + def _step(self, clock: pimm.Clock) -> None: """Keep one session call in flight and install the trajectory it returns. @@ -531,7 +543,7 @@ def _step(self, clock: pimm.Clock) -> None: self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. - self._future = executor.submit(session, frozen_view(_owned(obs))) + self._future = executor.submit(session, frozen_view(self._owned(obs))) if self._charge is None: # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 040f7a8ef..344839349 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1915,6 +1915,48 @@ def test_a_faulted_arm_reaches_the_policy_without_its_state(world): assert keys.EE_POSE not in policy.last_obs +@pytest.mark.timeout(3.0) +def test_one_arm_resetting_does_not_hide_another_arms_fault(world): + """A resetting arm and a faulted one both serialize to nothing. The fault wins whichever channel the + embodiment happens to list first, so a bimanual rig cannot carry on driving a faulted arm.""" + left, right = f'{keys.ROBOT_STATE}.left', f'{keys.ROBOT_STATE}.right' + embodiment = Embodiment( + '', + { + left: Observation(pimm.NoOpEmitter(), Serializers.robot_state), + right: Observation(pimm.NoOpEmitter(), Serializers.robot_state), + keys.GRIP: Observation(pimm.NoOpEmitter(), None), + }, + {keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command)}, + {}, + pimm.NoOpEmitter(), + ) + policy = SpyPolicy() + harness = Harness(policy, embodiment) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + left_em = world.pair(harness.observations[left]) + right_em = world.pair(harness.observations[right]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = world.pair(harness.directive) + + def emit_states(): + left_em.emit(make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.RESETTING)) + right_em.emit(make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR)) + grip_em.emit(0.0) + + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='test')), 0.0), + (emit_states, 0.01), + (None, 0.02), + ]) + + drive_scheduler(world.start([harness, driver]), steps=40) + + assert policy.last_obs is not None, 'the resetting arm swallowed the other arm’s fault' + assert policy.last_obs[keys.ROBOT_FAULT] is True + + @pytest.mark.timeout(20.0) def test_a_stop_lands_without_waiting_out_the_charge(world): """A charge places a trajectory's waypoints, and a stop has none: an arm that faults mid-chunk stops in From 8964abdb297ca95d78206158ef79641dbd63a0cb Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 15 Aug 2026 21:38:59 +0300 Subject: [PATCH 29/64] Sum the delta commands a late round overtook, instead of dropping them --- positronic/drivers/roboarm/command.py | 30 ++++++++++++++++++++ positronic/policy/harness.py | 37 +++++++++++++++---------- positronic/policy/tests/test_harness.py | 25 ++++++++++++++++- 3 files changed, 77 insertions(+), 15 deletions(-) diff --git a/positronic/drivers/roboarm/command.py b/positronic/drivers/roboarm/command.py index f7c62e016..b6372627b 100644 --- a/positronic/drivers/roboarm/command.py +++ b/positronic/drivers/roboarm/command.py @@ -1,5 +1,6 @@ """Collection of commands that can be sent to the robot.""" +from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any @@ -119,3 +120,32 @@ def from_wire(wire: dict[str, Any]) -> CommandType: ) case _: raise ValueError(f'Unknown command type: {wire["type"]}') + + +def _combine(acc: CommandType, cmd: CommandType) -> CommandType: + match (acc, cmd): + case (CartesianDelta(a, frame_a), CartesianDelta(b, frame_b)): + if not np.allclose(frame_a.as_matrix, frame_b.as_matrix): + raise ValueError('Cannot accumulate cartesian deltas expressed in different frames') + return CartesianDelta(_compose_delta(a, b), frame_a) + case (JointDelta(a), JointDelta(b)): + return JointDelta(a + b) + case (CartesianDelta() | JointDelta(), _) | (_, CartesianDelta() | JointDelta()): + raise ValueError(f'Cannot reduce {type(acc).__name__} then {type(cmd).__name__} in one tick') + case _: + return cmd + + +def reduce(due: Sequence[CommandType]) -> CommandType: + """Collapse the commands due in one round into the single command to execute. + + Folds the batch in timestamp order. A run of same-space deltas accumulates (their motion is summed, so a + round spanning several waypoints catches up rather than dropping them); a run of absolute commands keeps + the last. Mixing an absolute with a delta, or two delta spaces, has no faithful single-command form — a + delta binds to the pose measured when it is consumed, which an absolute target or a foreign space cannot + supply — and raises. + """ + result = due[0] + for cmd in due[1:]: + result = _combine(result, cmd) + return result diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 32f7180d6..27a1a613c 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -2,7 +2,7 @@ import logging import time from collections import deque -from collections.abc import Generator, Iterable, Iterator +from collections.abc import Callable, Generator, Iterable, Iterator, Sequence from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum @@ -39,12 +39,23 @@ Trajectory: TypeAlias = list[tuple[int, Any]] +def _last(due: Sequence[Any]) -> Any: + """The trailing value wins — the right collapse for absolute setpoints and gripper targets.""" + return due[-1] + + class TrajectoryPlayer: """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value - to emit.""" + to emit. + + ``reduce`` collapses the waypoints that came due together in one round. Keeping the last is right for a + value that states where to be; a channel carrying deltas passes ``roboarm.command.reduce`` so their + motion is summed rather than dropped. + """ - def __init__(self): + def __init__(self, reduce: Callable[[Sequence[Any]], Any] = _last): self._pending: deque[tuple[int, Any]] = deque() + self._reduce = reduce def set(self, trajectory: Trajectory): self._pending = deque(trajectory) @@ -54,17 +65,12 @@ def next_due(self) -> int | None: return self._pending[0][0] if self._pending else None def advance(self, current_time: int): - """The single value due at ``current_time`` — the last, when several came due since the previous - call — or ``None`` when none did. - - Collapsing to the last is exact for an absolute setpoint and lossy for a relative one: a run of - deltas due together arrives as its final step alone. Pacing keeps one waypoint due per round - wherever a round is shorter than the spacing between waypoints. - """ - value = None + """The single value due at ``current_time``, collapsed by ``reduce`` when several came due since the + previous call, or ``None`` when none did.""" + due = [] while self._pending and self._pending[0][0] <= current_time: - value = self._pending.popleft()[1] - return value + due.append(self._pending.popleft()[1]) + return self._reduce(due) if due else None def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: @@ -255,7 +261,10 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] - self._players = {name: TrajectoryPlayer() for name in embodiment.commands} + self._players = { + name: TrajectoryPlayer(roboarm.command.reduce if keys.is_robot_command(name) else _last) + for name in embodiment.commands + } self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 344839349..9f5e9dd90 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -13,7 +13,7 @@ from positronic.dataset.serializers import Serializers from positronic.drivers import roboarm from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, Reset, from_wire, to_wire +from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, JointDelta, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model from positronic.eval import Command, Embodiment, Observation, Task from positronic.geom import Rotation, Transform3D @@ -1362,6 +1362,29 @@ def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): assert player.advance(40) is None +def test_trajectory_player_sums_the_deltas_a_late_round_overtook(): + """A delta states how far to move, so the ones a round overtook are motion still owed. Keeping only the + last would silently shorten the trajectory.""" + player = TrajectoryPlayer(roboarm.command.reduce) + player.set([(10, JointDelta(np.array([0.1, 0.0]))), (20, JointDelta(np.array([0.2, 0.5])))]) + + caught_up = player.advance(25) + + assert isinstance(caught_up, JointDelta) + np.testing.assert_allclose(caught_up.velocities, [0.3, 0.5]) + + +def test_trajectory_player_refuses_to_collapse_a_delta_onto_an_absolute(): + """No single command carries both: a delta binds to the pose measured when it is consumed, which an + absolute target cannot supply.""" + player = TrajectoryPlayer(roboarm.command.reduce) + pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) + player.set([(10, CartesianPosition(pose=pose)), (20, JointDelta(np.array([0.1, 0.0])))]) + + with pytest.raises(ValueError, match='Cannot reduce'): + player.advance(25) + + def test_cartesian_delta_applies_in_world_frame(): current = Transform3D(np.array([0.5, 0.1, 0.3]), Rotation.from_rotvec(np.array([0.2, 0.1, 0.4]))) delta = Transform3D(np.array([0.02, -0.01, 0.05]), Rotation.from_rotvec(np.array([0.1, 0.0, 0.0]))) From 5f4e79b5d70361f76b5dc969e08bb3117126697d Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy Date: Sat, 15 Aug 2026 23:20:31 +0200 Subject: [PATCH 30/64] Let the human review declines one at a time, on their go --- .claude/skills/address-review/SKILL.md | 43 ++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index bd094b2a7..cff99fe9a 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -256,10 +256,53 @@ Summarize: - which threads you resolved (fixes only) vs left open (declines / defers / discussion), - any follow-ups the user should track. +Then **offer** the walk below, in one line, and go to Step 7 without waiting: + +> 3 threads are open on my call: , <title>, <title>. Want to go through them? + A bot will re-review on push and may add comments. **Don't hand the watch back to the user** — go to Step 7, which watches for that re-review in the background and loops you through another pass automatically until the reviewer converges. +## Walking your declines — only on an explicit go + +Every thread you declined or deferred is a decision the human has not seen. They are entitled to +review each one, during the cycle or after it converges. **Offer it; never start it uninvited.** A +proposal is one line (above). Presenting a thread before they say go is the failure this section +exists to prevent — as is presenting all of them at once. + +On their go, **one thread per message, four sentences, no headings, no code blocks**: + +1. what the thing is, +2. what the reviewer wanted, +3. why you declined it, +4. the decision you need — bolded, on its own line, with the options named. + +Then **stop**. Do not append the next thread, a summary, or what you plan to do after. + +Everything beyond those four sentences is available on request and withheld until then — the diff, the +evidence, the options table, the code. `git show`, measurements, a longer argument: all fine to *have*, +none of it goes in unless asked. + +- **"too long" / "too complicated" means cut it in half** — not rewrite it at the same length. If they + say it twice, you are still writing for yourself. +- **A follow-up question gets that answer alone.** "why declined?" is not an invitation to re-present + the thread; "which three?" wants three names and nothing else. +- Assume they have not read the code and will not. No "see the call site". + +Their verdicts are terse. What each one means: + +| They say | You do | +|---|---| +| "next" / "resolve, next" | reply, resolve, present the next thread — same shape, no re-asking | +| "leave it" | reply that it is accepted as a known limit, resolve | +| "fix it" | fix, run the gates, commit, push, reply with the SHA, resolve | +| "file it" / "open a bug" | file the issue, reply linking it, resolve | +| "why declined?" | the reasoning only | + +This is the one path on which a declined thread gets resolved: the human closed it, thread by thread. +Absent that, Step 5 stands and it stays open. + ## Step 7: Watch for convergence in the background (so the user doesn't have to) Two signals decide whether a push actually lands the PR, and **both are asynchronous**: CI From 023ce60e804a60503ae0361062d3199ea5cf34d3 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sat, 15 Aug 2026 23:43:35 +0200 Subject: [PATCH 31/64] Say that a confused reader needs unpacking, not cutting --- .claude/skills/address-review/SKILL.md | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index cff99fe9a..65e264895 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -271,21 +271,30 @@ review each one, during the cycle or after it converges. **Offer it; never start proposal is one line (above). Presenting a thread before they say go is the failure this section exists to prevent — as is presenting all of them at once. -On their go, **one thread per message, four sentences, no headings, no code blocks**: +On their go, **one thread per message, no section headings**, covering four things: 1. what the thing is, 2. what the reviewer wanted, 3. why you declined it, 4. the decision you need — bolded, on its own line, with the options named. +**Simple beats short.** The failure mode is density, not length: four sentences carrying three ideas +each are worse than twelve carrying one each. Unpack — one idea per sentence, plain words, and +describe the thing before arguing about it. Length is the budget for being understood; spend it on +unpacking and on nothing else. + +**A short code snippet is welcome** where code says it faster than prose — the branch that decides, +the line that was removed. A few lines, and say what it shows. Long listings, file tours and full +diffs are not. + Then **stop**. Do not append the next thread, a summary, or what you plan to do after. -Everything beyond those four sentences is available on request and withheld until then — the diff, the -evidence, the options table, the code. `git show`, measurements, a longer argument: all fine to *have*, -none of it goes in unless asked. +Withheld until asked: the measurements, the alternative designs, the options table, the wider diff. +Fine to *have* — `git show`, instrumentation, a longer argument — none of it goes in unasked. -- **"too long" / "too complicated" means cut it in half** — not rewrite it at the same length. If they - say it twice, you are still writing for yourself. +- **"too long" means cut it in half** — not rewrite it at the same length. +- **"too complicated" / "I don't understand" means unpack it** — more sentences, fewer ideas per + sentence, and a snippet if one exists. Cutting further is the wrong reflex and makes it worse. - **A follow-up question gets that answer alone.** "why declined?" is not an invitation to re-present the thread; "which three?" wants three names and nothing else. - Assume they have not read the code and will not. No "see the call site". From d7429ada86d5c82c1ad04877bb9bfedea3588319 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 00:15:43 +0200 Subject: [PATCH 32/64] Return a named outcome from `_take` instead of a bare bool --- positronic/policy/harness.py | 23 +++++++++++++++-------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 27a1a613c..3e32980ba 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -83,6 +83,14 @@ def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: ) +class _Answer(Enum): + """What became of the call ``_take`` was handed: its trajectory is installed and the future spent, or the + world has not yet paid the charge and the same future comes back next round.""" + + CONSUMED = 'consumed' + PENDING = 'pending' + + class DirectiveType(Enum): RUN = 'run' FINISH = 'finish' @@ -536,7 +544,7 @@ def _step(self, clock: pimm.Clock) -> None: """ session, executor = self._policy_session, self._executor assert session is not None and executor is not None # only a live episode steps - if self._future is not None and not self._take(self._future, clock): + if self._future is not None and self._take(self._future, clock) is _Answer.PENDING: return obs = self._build_obs(clock) if obs is None: @@ -551,7 +559,6 @@ def _step(self, clock: pimm.Clock) -> None: self._deadline = clock.now() + self._task.timeout self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() - # Sessions declare ``dict`` but must not mutate the obs, so they get a read-only view. self._future = executor.submit(session, frozen_view(self._owned(obs))) if self._charge is None: # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. @@ -559,8 +566,8 @@ def _step(self, clock: pimm.Clock) -> None: time.sleep(0) self._take(self._future, clock) - def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> bool: - """Install the call's trajectory once the world has paid for it; True once the future is consumed. + def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: + """Install the call's trajectory once the world has paid for it. Under a constant charge the world holds still until the call answers — blocking here blocks the loop thread, which is what advances a virtual clock. Until a call answers there is no telling a skip from @@ -576,20 +583,20 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) elif not future.done(): ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) if ahead <= 0.0: - return False + return _Answer.PENDING concurrent.futures.wait([future], timeout=ahead) if not future.done(): - return False + return _Answer.PENDING actions = future.result() # taken on the loop thread, so a failing call still seals the episode if actions and self._charge is not None: # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP # and slips the install a full round. if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): - return False # the schedule already playing carries the world to the release instant + return _Answer.PENDING # the schedule already playing carries the world to the release instant self._future = None if actions is not None: self._install(actions, clock) - return True + return _Answer.CONSUMED def _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: """Replace the schedule being played with the session's trajectory. Every channel it names gets that From 8bf7047cb551d950106f32b9d57690ba96ca5a20 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 00:27:04 +0200 Subject: [PATCH 33/64] Move `_Answer` into `Harness`, above the step that reads it --- positronic/policy/harness.py | 25 ++++++++++++------------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 3e32980ba..1cdfdc11a 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -83,14 +83,6 @@ def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: ) -class _Answer(Enum): - """What became of the call ``_take`` was handed: its trajectory is installed and the future spent, or the - world has not yet paid the charge and the same future comes back next round.""" - - CONSUMED = 'consumed' - PENDING = 'pending' - - class DirectiveType(Enum): RUN = 'run' FINISH = 'finish' @@ -536,6 +528,13 @@ 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()} + class _Answer(Enum): + """What became of the call ``_take`` was handed: its trajectory is installed and the future spent, or + the world has not yet paid the charge and the same future comes back next round.""" + + CONSUMED = 'consumed' + PENDING = 'pending' + def _step(self, clock: pimm.Clock) -> None: """Keep one session call in flight and install the trajectory it returns. @@ -544,7 +543,7 @@ def _step(self, clock: pimm.Clock) -> None: """ session, executor = self._policy_session, self._executor assert session is not None and executor is not None # only a live episode steps - if self._future is not None and self._take(self._future, clock) is _Answer.PENDING: + if self._future is not None and self._take(self._future, clock) is Harness._Answer.PENDING: return obs = self._build_obs(clock) if obs is None: @@ -583,20 +582,20 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) elif not future.done(): ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) if ahead <= 0.0: - return _Answer.PENDING + return Harness._Answer.PENDING concurrent.futures.wait([future], timeout=ahead) if not future.done(): - return _Answer.PENDING + return Harness._Answer.PENDING actions = future.result() # taken on the loop thread, so a failing call still seals the episode if actions and self._charge is not None: # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP # and slips the install a full round. if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): - return _Answer.PENDING # the schedule already playing carries the world to the release instant + return Harness._Answer.PENDING # the schedule already playing carries the world to the release instant self._future = None if actions is not None: self._install(actions, clock) - return _Answer.CONSUMED + return Harness._Answer.CONSUMED def _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: """Replace the schedule being played with the session's trajectory. Every channel it names gets that From c2cd8b181029d24f229c440d53937b6a6a9e5a34 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 00:39:12 +0200 Subject: [PATCH 34/64] Wait on the future instead of spinning, and name the charge `_fixed_latency` --- positronic/policy/harness.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 1cdfdc11a..0e7f147ba 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -240,9 +240,10 @@ def __init__( # was submitted. self._t0_ns = 0 self._wall_t0 = 0.0 - # Seconds each model call costs the world clock this episode, or ``None`` to charge the call's own - # wall duration (hardware pace, and the sim's ``inference_latency=True``). - self._charge: float | None = None + # Seconds each model call costs the world clock this episode, the same figure for every call. + # ``None`` when no figure is fixed and the world is charged the call's own wall duration (hardware + # pace, and the sim's ``inference_latency=True``). + self._fixed_latency: float | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None # Whether this episode's first observation has landed. Until it does the deadline stands where the @@ -385,10 +386,10 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: # every model call. latency = self.context.setdefault(keys.INFERENCE_LATENCY, False) else: - # The charge is a device for simulating a trial, so a real rig ignores it and pays the wall time - # its calls really take. + # A fixed latency is a device for simulating a trial, so a real rig ignores it and pays the wall + # time its calls really take. latency = True - self._charge = None if latency is True else float(latency) + self._fixed_latency = None if latency is True else float(latency) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -514,7 +515,7 @@ def _effect_time(self) -> float: charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; the loop thread writes the call's start fields before submitting it. """ - charge = time.monotonic() - self._wall_t0 if self._charge is None else self._charge + charge = time.monotonic() - self._wall_t0 if self._fixed_latency is None else self._fixed_latency return self._t0_ns / 1e9 + charge @staticmethod @@ -559,10 +560,8 @@ def _step(self, clock: pimm.Clock) -> None: self._t0_ns = clock.now_ns() self._wall_t0 = time.monotonic() self._future = executor.submit(session, frozen_view(self._owned(obs))) - if self._charge is None: - # Sleeping zero hands the worker the GIL without adding a wake-up granularity to the handshake. - while not self._future.done() and time.monotonic() - self._wall_t0 < SKIP_REPLY_SEC: - time.sleep(0) + if self._fixed_latency is None: + concurrent.futures.wait([self._future], timeout=SKIP_REPLY_SEC) self._take(self._future, clock) def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: @@ -577,7 +576,7 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) stops what is executing — has no such instant and lands at once. A charge measured in wall time can hold nothing still, so there the world runs no further ahead of the call's start than wall time has. """ - if self._charge is not None: + if self._fixed_latency is not None: concurrent.futures.wait([future]) elif not future.done(): ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) @@ -587,10 +586,10 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) if not future.done(): return Harness._Answer.PENDING actions = future.result() # taken on the loop thread, so a failing call still seals the episode - if actions and self._charge is not None: + if actions and self._fixed_latency is not None: # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP # and slips the install a full round. - if clock.now_ns() < self._t0_ns + round(self._charge * 1e9): + if clock.now_ns() < self._t0_ns + round(self._fixed_latency * 1e9): return Harness._Answer.PENDING # the schedule already playing carries the world to the release instant self._future = None if actions is not None: From 888b345acefb43db86951c4c6f0ff6837320ca64 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 01:06:24 +0200 Subject: [PATCH 35/64] Make the in-scope call in `address-review`, not the human --- .claude/skills/address-review/SKILL.md | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index 65e264895..04a388d91 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -111,8 +111,8 @@ For each open comment, decide and note severity if the bot tagged one (e.g. Code resolve only if you land a concrete change, else leave open. - **Decline** — wrong, not applicable, or contradicts a deliberate decision → reasoned reply, **leave open**. -- **Defer** — valid but out of scope for this PR → reply (note where it's tracked), - **leave open**. +- **Defer** — valid but out of scope for this PR → file the tracking issue unless something + already tracks it, reply naming it, **leave open**. - **Discuss** — the reviewer is asking a question or opening a design discussion, not requesting a change → answer it, **leave open** for them to respond. @@ -122,8 +122,15 @@ Present the triage as a short numbered list: comment → verdict → planned fix declines without prompting. "Unclear" is a verdict you reach after the research above, never before it — a question you could have answered by reading the code costs the user a round trip and comes back as *go and find out*, which is the same work plus a delay. Pause and confirm -only for what research cannot settle: a preference, a cost, a scope call, or a change whose -blast radius is genuinely the user's to accept. The user may drop or override any item. +only for what research cannot settle: a preference, a cost, or a change whose blast radius is +genuinely the user's to accept. The user may drop or override any item. + +**Whether a finding belongs in this PR is your call.** "Fix it here or track it for later" is +the verdict this skill exists to reach, and by the time you can pose it the research is done, +so the human has nothing to add that you don't already have. Reach it, file the issue, reply, +and put it in the Step 6 report where they can overrule it. Bring a deferral back only when +leaving it undone changes what the PR ships — a gap in the behaviour the user asked for, not a +follow-up carrying an issue number. When you do bring something back, explain it in full and in plain words: what the thing is and how it works, what the reviewer asked for, what you did or did not do and why, and what each From 85422714d761608777bf29842aa5b0f547ca22e1 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 01:21:12 +0200 Subject: [PATCH 36/64] Join the abandoned call at shutdown, and drop two complexity waivers --- .basedpyright/baseline.json | 16 ---------- positronic/dataset/ds_writer_agent.py | 38 +++++++++++++---------- positronic/policy/harness.py | 33 +++++++++++--------- positronic/policy/tests/test_harness.py | 41 +++++++++++++++++++++++++ positronic/simulator/mujoco/sim.py | 23 ++++++-------- 5 files changed, 91 insertions(+), 60 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index ecfd23cd5..ef5e98a09 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -1529,14 +1529,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 56, - "endColumn": 58, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { @@ -1545,14 +1537,6 @@ "lineCount": 1 } }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 57, - "endColumn": 59, - "lineCount": 1 - } - }, { "code": "reportOptionalMemberAccess", "range": { diff --git a/positronic/dataset/ds_writer_agent.py b/positronic/dataset/ds_writer_agent.py index 54b204594..8e2079e16 100644 --- a/positronic/dataset/ds_writer_agent.py +++ b/positronic/dataset/ds_writer_agent.py @@ -131,7 +131,24 @@ def add_signal(self, name: str, serializer: Serializer | StatefulSerializer | No def inputs(self) -> dict[str, pimm.ControlSystemReceiver[Any]]: return frozen_keys_dict(self._inputs) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C901 + def _record(self, ep_writer: EpisodeWriter, name: str, msg: pimm.Message, clock: pimm.Clock) -> None: + """Append one input's sample, stamped as ``time_mode`` selects and carrying every clock beside it.""" + world_time_ns, message_time_ns = clock.now_ns(), msg.ts + primary_ts = world_time_ns if self._time_mode == TimeMode.CLOCK else message_time_ns + + extra_ts = {'message': message_time_ns, 'system': pimm.world.SystemClock().now_ns()} + # Only add 'world' if clock is not system clock + if not isinstance(clock, pimm.world.SystemClock): + extra_ts['world'] = world_time_ns + + with self._telemetry_span(): + serializer = self._serializers.get(name) + value = msg.data + if serializer is not None: + value = serializer(value) + _append(ep_writer, name, value, primary_ts, extra_ts) + + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): """Main loop: process commands and append updated inputs to the episode.""" limiter = pimm.utils.RateLimiter(clock, hz=self._poll_hz) pace = (lambda: pimm.Yield()) if self._virtual_time else limiter.wait @@ -155,6 +172,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 if ep_writer is not None: for name, reader in self._inputs.items(): msg = reader.read() + if msg is None: + continue # Scope a command turn's drain to the episode window: the open turn keeps only # samples after START (dropping the inter-episode home command and any pre-reset # frame), the closing turn keeps only samples at or before STOP (dropping post-STOP @@ -163,21 +182,8 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): # noqa: C90 # normally. after_start = not opened or msg.ts > cmd_msg.ts before_stop = not closing or msg.ts <= cmd_msg.ts - if msg is not None and msg.updated and after_start and before_stop: - world_time_ns, message_time_ns = clock.now_ns(), msg.ts - primary_ts = world_time_ns if self._time_mode == TimeMode.CLOCK else message_time_ns - - extra_ts = {'message': message_time_ns, 'system': pimm.world.SystemClock().now_ns()} - # Only add 'world' if clock is not system clock - if not isinstance(clock, pimm.world.SystemClock): - extra_ts['world'] = world_time_ns - - with self._telemetry_span(): - serializer = self._serializers.get(name) - value = msg.data - if serializer is not None: - value = serializer(value) - _append(ep_writer, name, value, primary_ts, extra_ts) + if msg.updated and after_start and before_stop: + self._record(ep_writer, name, msg, clock) if closing: ep_writer, ep_counter = self._handle_command(cmd_msg.data, ep_writer, ep_counter) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 0e7f147ba..4c2695eb9 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -343,13 +343,13 @@ def _retire_worker(self) -> None: self._retiring, self._executor = self._executor, None def _reap_worker(self) -> None: - """Wait out the previous episode's abandoned call before this one opens a session. + """Wait out an abandoned call before anything else touches the policy. - An in-process policy is a single model across episodes, so ``new_session`` resets the very object an - abandoned call may still be inside — a running thread survives ``shutdown(cancel_futures=True)``, - which cancels only what is still queued. Waiting here rather than at the end of the episode it - belongs to keeps a hung model from holding up that episode's recording and home, and the task reset - this follows usually covers the wait. + An in-process policy is a single model across episodes and across runs, so ``new_session`` and + ``close`` reach the very object an abandoned call may still be inside — a running thread survives + ``shutdown(cancel_futures=True)``, which cancels only what is still queued. The wait sits at the next + episode's start and at the harness's own shutdown, never at the end of the episode the call belongs + to, so a hung model cannot hold up that episode's recording and home. """ if self._retiring is not None: self._retiring.shutdown(wait=True) @@ -371,6 +371,14 @@ def _finalize_recording( # period) to the closing episode — the cooperative scheduler cannot give the recorder a turn alone. self._telemetry.end(virtual_now) + def _effect_time(self) -> float: + """The trial instant the in-flight call's output takes effect: its observation instant plus the + charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; + the loop thread writes the call's start fields before submitting it. + """ + charge = time.monotonic() - self._wall_t0 if self._fixed_latency is None else self._fixed_latency + return self._t0_ns / 1e9 + charge + def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. @@ -510,14 +518,6 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it return inputs - def _effect_time(self) -> float: - """The trial instant the in-flight call's output takes effect: its observation instant plus the - charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; - the loop thread writes the call's start fields before submitting it. - """ - charge = time.monotonic() - self._wall_t0 if self._fixed_latency is None else self._fixed_latency - return self._t0_ns / 1e9 + charge - @staticmethod def _owned(obs: dict[str, Any]) -> dict[str, Any]: """The observation with its arrays copied, so nothing rewrites what the worker is still reading. @@ -653,9 +653,12 @@ def _shutdown(self) -> None: dropped: the run is over and nothing is left to install it. The harness does not own the policy's lifetime: the caller may run several harnesses over one policy - (a multi-eval sweep), so it closes the policy once, after the last run. + (a multi-eval sweep), so it closes the policy once, after the last run. That leaves no later + ``_begin_episode`` to reap on, and the next harness reaches the shared policy through a session of + its own, so the call is waited out here. """ self._retire_worker() + self._reap_worker() if self._policy_session is not None: self._policy_session.close() self._policy_session = None diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 9f5e9dd90..995cac07a 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,3 +1,4 @@ +import threading import time from contextlib import contextmanager from functools import partial @@ -2058,6 +2059,46 @@ def new_session(self, context=None, now=None): assert isinstance(cmd_recorder.emitted[-1][1], Reset) +@pytest.mark.timeout(20.0) +def test_the_run_ends_only_once_the_call_it_abandoned_is_out_of_the_policy(): + """A sweep runs a harness per eval over one shared policy, so a call still inside the model when a run + ends would meet the next run's ``new_session`` — or ``policy.close()``. The run outlives its own call.""" + hang_sec = 1.0 + left_the_model = threading.Event() + + class _HangingSession(Session): + def __call__(self, obs): + time.sleep(hang_sec) + left_the_model.set() + return None + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver]), steps=40) + + assert left_the_model.is_set(), 'the run returned with a call still inside the shared policy' + + @pytest.mark.timeout(3.0) def test_installed_trajectory_clears_the_channels_it_omits(world): """A trajectory naming only one channel replaces the whole schedule: the omitted channel stops being diff --git a/positronic/simulator/mujoco/sim.py b/positronic/simulator/mujoco/sim.py index b92dd02ed..6d41d3077 100644 --- a/positronic/simulator/mujoco/sim.py +++ b/positronic/simulator/mujoco/sim.py @@ -156,12 +156,14 @@ def __init__( # replays these states through it (``mj_setState`` + ``mj_forward``). self.sim_state: pimm.SignalEmitter[dict[str, np.ndarray]] = pimm.ControlSystemEmitter(self) - def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: # noqa: C901 + def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Sleep]: self._emit_robot_meta() - state_due = _Cadence(self._state_fps) - grip_due = _Cadence(self._grip_fps) - sim_state_due = _Cadence(self._sim_state_fps) - cameras_due = _Cadence(self._camera_fps) + streams = [ + (_Cadence(self._state_fps), self._emit_state), + (_Cadence(self._grip_fps), self._emit_grip), + (_Cadence(self._sim_state_fps), self._emit_sim_state), + (_Cadence(self._camera_fps), self._emit_cameras), + ] while not should_stop.value: yield pimm.Sleep(self.model.opt.timestep) @@ -192,14 +194,9 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p with telemetry.span(telemetry_keys.SPAN_ENV_STEP): self.step() self.fps_counter.tick() - if state_due(now): - self._emit_state() - if grip_due(now): - self._emit_grip() - if sim_state_due(now): - self._emit_sim_state() - if cameras_due(now): - self._emit_cameras() + for due, emit in streams: + if due(now): + emit() if self._renderer is not None: self._renderer.close() From a3a98955a76c26234d51330efe55eadf66018054 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 01:50:30 +0200 Subject: [PATCH 37/64] Close a retired session only once its call has left it --- positronic/policy/harness.py | 52 +++++++++++++------------ positronic/policy/tests/test_harness.py | 48 +++++++++++++++++++++++ 2 files changed, 76 insertions(+), 24 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 4c2695eb9..eef2e7f25 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -234,7 +234,9 @@ def __init__( # worker belongs to the episode: ending one abandons the call in flight rather than waiting for it, # so the next episode must not queue behind it. self._executor: ThreadPoolExecutor | None = None - self._retiring: ThreadPoolExecutor | None = None + # The retired worker and the session its abandoned call is still inside, held as one because closing + # the session is what the join makes safe. + self._retiring: tuple[ThreadPoolExecutor, Session] | None = None self._future: Future[list[dict[str, Any]] | None] | None = None # The in-flight call's start: the world instant its observation was built, and the wall instant it # was submitted. @@ -332,36 +334,42 @@ def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: logging.error(f'Inference failed after the episode that asked for it ended: {exc}') def _retire_worker(self) -> None: - """Let go of this episode's worker and the call it is running: the answer lands nowhere and the - failure only reaches the log. The worker is kept for ``_reap_worker`` to join at the next episode's - start, since ending an episode must not wait for a model that hangs.""" + """Let go of this episode's worker, the call it is running and the session that call is inside: the + answer lands nowhere and the failure only reaches the log. All three are kept for ``_reap_worker``, + since ending an episode must not wait for a model that hangs.""" if self._future is not None: self._future.add_done_callback(self._report_abandoned) self._future = None - if self._executor is not None: + if self._executor is not None and self._policy_session is not None: self._executor.shutdown(wait=False, cancel_futures=True) - self._retiring, self._executor = self._executor, None + self._retiring = (self._executor, self._policy_session) + self._executor, self._policy_session = None, None def _reap_worker(self) -> None: - """Wait out an abandoned call before anything else touches the policy. - - An in-process policy is a single model across episodes and across runs, so ``new_session`` and - ``close`` reach the very object an abandoned call may still be inside — a running thread survives - ``shutdown(cancel_futures=True)``, which cancels only what is still queued. The wait sits at the next - episode's start and at the harness's own shutdown, never at the end of the episode the call belongs - to, so a hung model cannot hold up that episode's recording and home. + """Wait out an abandoned call, then close the session it was inside. + + A running thread survives ``shutdown(cancel_futures=True)``, which cancels only what is still queued, + so until the join returns the call still holds the session's resources — a ``RemoteSession``'s + websocket is the one ``close`` would pull out from under it — and, for an in-process policy, the one + model that every session across episodes and runs shares. The wait sits at the next episode's start + and at the harness's own shutdown, never at the end of the episode the call belongs to, so a hung + model cannot hold up that episode's recording and home. """ if self._retiring is not None: - self._retiring.shutdown(wait=True) + executor, session = self._retiring self._retiring = None + executor.shutdown(wait=True) + session.close() 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 episode's full static meta (plus any terminal payload) — then close its span.""" + # Stamped before the session is retired with the worker: the meta overlays what the session reports. + stop = DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})}) self._cancel_session() - self.ds_command.emit(DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})})) + self.ds_command.emit(stop) virtual_now = clock.now() # before the round below, whose sim-clock advance belongs to no rollout # Give the recorder a round to commit the STOP before the next START (they share ``ds_command``, where # last-value-wins would drop one) and before the home command, so homing stays out of the recording. @@ -423,10 +431,12 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: def _end_episode( self, clock: pimm.Clock, payload: dict[str, Any] | None = None, *, abort: bool = False ) -> Generator[pimm.Command, None, None]: - """Close the live episode: finalize (or abort) the recording, release the session, home devices. + """Close the live episode: finalize (or abort) the recording, retire the session, home devices. - Releasing the session here, not only at shutdown, closes a ``RemoteSession``'s websocket promptly, so - the offboard server's per-session cleanup (active-session decrement, idle watchdog) runs now. + The session is retired with the worker rather than closed here, so a ``RemoteSession``'s websocket + outlives the call still using it; ``_reap_worker`` closes it at the next episode's start or at + shutdown, and the offboard server's per-session cleanup (active-session decrement, idle watchdog) + runs then. """ if self._running: if abort: @@ -436,9 +446,6 @@ def _end_episode( self._telemetry.abort() else: yield from self._finalize_recording(clock, payload) - if self._policy_session: - self._policy_session.close() - self._policy_session = None self._home() self._running = False @@ -659,9 +666,6 @@ def _shutdown(self) -> None: """ self._retire_worker() self._reap_worker() - if self._policy_session is not None: - self._policy_session.close() - self._policy_session = None def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: while not should_stop.value: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 995cac07a..df85e31ba 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -2099,6 +2099,54 @@ def new_session(self, context=None, now=None): assert left_the_model.is_set(), 'the run returned with a call still inside the shared policy' +@pytest.mark.timeout(20.0) +def test_the_session_is_closed_only_once_its_call_has_left_it(): + """``RemoteSession.close`` shuts the websocket the call in flight is talking over, and ``Session`` asks + for no thread safety, so the session is retired with its worker and closed once that worker is joined.""" + inside_at_close = [] + + class _HangingSession(Session): + def __init__(self): + self.inside = False + + def __call__(self, obs): + self.inside = True + try: + time.sleep(1.0) + return None + finally: + self.inside = False + + def close(self): + inside_at_close.append(self.inside) + + class _HangingPolicy(Policy): + def new_session(self, context=None, now=None): + return _HangingSession() + + with pimm.World() as world: + harness = Harness(_HangingPolicy(), make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.01), + (partial(directive_em.emit, Directive.FINISH()), 0.05), + (None, 0.05), + ]) + drive_scheduler(world.start([harness, driver]), steps=40) + + assert inside_at_close == [False], 'the session was closed while its own call was still inside it' + + @pytest.mark.timeout(3.0) def test_installed_trajectory_clears_the_channels_it_omits(world): """A trajectory naming only one channel replaces the whole schedule: the omitted channel stops being From 1df58de2e70ee059f87e594a78f668f6c488bb81 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 02:02:53 +0200 Subject: [PATCH 38/64] Drop a chunk whose call outlived the trial deadline --- positronic/policy/harness.py | 5 ++++- positronic/policy/tests/test_harness.py | 18 ++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index eef2e7f25..b6529a541 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -599,7 +599,10 @@ def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) if clock.now_ns() < self._t0_ns + round(self._fixed_latency * 1e9): return Harness._Answer.PENDING # the schedule already playing carries the world to the release instant self._future = None - if actions is not None: + # The world reached the deadline while the call 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. + expired = self._deadline is not None and clock.now() >= self._deadline + if actions is not None and not expired: self._install(actions, clock) return Harness._Answer.CONSUMED diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index df85e31ba..9e61e312b 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,5 +1,6 @@ import threading import time +from concurrent.futures import Future from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -1689,6 +1690,23 @@ def test_anchored_chunk_passes(): _assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) +@pytest.mark.parametrize(('expired', 'installed'), [(True, False), (False, True)]) +def test_a_reply_is_installed_only_while_the_trial_still_has_budget(world, expired, installed): + """A trial advertises the instant it stops at. A call whose rounds in flight carried the world past that + instant has its chunk dropped instead of placed, and ``_run`` finishes the trial on the next round.""" + harness = Harness(StubPolicy(), make_embodiment()) + now = world.clock.now() + harness._deadline = now - 1.0 if expired else now + 1.0 + pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) + future = Future() + future.set_result([{keys.ROBOT_COMMAND: CartesianPosition(pose=pose), keys.ACTION_TIMESTAMP: now}]) + + harness._take(future, world.clock) + + played = harness._players[keys.ROBOT_COMMAND].advance(world.clock.now_ns()) + assert (played is not None) is installed + + class _SlowSession(Session): """A session whose inference costs ``wall_sec`` of real time and returns a fixed-length chunk.""" From 9c2afd628c81f5b8d6579d96b6ad3778d8142b0c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 02:13:23 +0200 Subject: [PATCH 39/64] Let each wrapper own its wire name, and key the table by it --- positronic/policy/action.py | 16 +++++++--- positronic/policy/base.py | 8 +++-- positronic/policy/codec.py | 28 ++++++++++++----- positronic/policy/observation.py | 4 ++- positronic/policy/spec.py | 39 +++++++++++------------- positronic/policy/tests/test_wrappers.py | 4 ++- positronic/policy/wrappers.py | 12 ++++++-- 7 files changed, 72 insertions(+), 39 deletions(-) diff --git a/positronic/policy/action.py b/positronic/policy/action.py index cb9a801d6..7af571de9 100644 --- a/positronic/policy/action.py +++ b/positronic/policy/action.py @@ -26,6 +26,8 @@ def _relative_rot_vec(q_current: np.ndarray, q_target: np.ndarray, representatio class AbsolutePositionAction(Codec): + WIRE_NAME = 'absolute_position_action' + def __init__(self, tgt_ee_pose_key: str, tgt_grip_key: str, rotation_rep: RotRep | str = RotRep.QUAT): self.rot_rep = RotRep(rotation_rep) self.tgt_ee_pose_key = tgt_ee_pose_key @@ -54,7 +56,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'absolute_position_action', + 'name': self.WIRE_NAME, 'args': { 'tgt_ee_pose_key': self.tgt_ee_pose_key, 'tgt_grip_key': self.tgt_grip_key, @@ -64,6 +66,8 @@ def to_spec(self): class AbsoluteJointsAction(Codec): + WIRE_NAME = 'absolute_joints_action' + def __init__(self, tgt_joints_key: str, tgt_grip_key: str, num_joints: int = 7): self.tgt_joints_key = tgt_joints_key self.tgt_grip_key = tgt_grip_key @@ -92,7 +96,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'absolute_joints_action', + 'name': self.WIRE_NAME, 'args': { 'tgt_joints_key': self.tgt_joints_key, 'tgt_grip_key': self.tgt_grip_key, @@ -137,6 +141,8 @@ def training_encoder(self): class RelativePositionAction(Codec): + WIRE_NAME = 'relative_position_action' + def __init__( self, rotation_rep: RotRep | str = RotRep.QUAT, @@ -194,7 +200,7 @@ def training_encoder(self): def to_spec(self): return { - 'name': 'relative_position_action', + 'name': self.WIRE_NAME, 'args': { 'rotation_rep': self.rot_rep.value, 'robot_pose_key': self.robot_pose_key, @@ -211,6 +217,8 @@ class JointDeltaAction(Codec): into a ``JointDelta`` command; the driver integrates each delta onto the live measured joints. """ + WIRE_NAME = 'joint_delta_action' + # General DROID form scales each normalized velocity by its own per-joint delta limit, then # renorms the velocity vector so no joint exceeds its limit: # RELATIVE_MAX_JOIN_DELTA = np.array([0.2, 0.2, 0.2, 0.2, 0.2, 0.2, 0.2]) @@ -241,4 +249,4 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return {keys.ROBOT_COMMAND: command.JointDelta(velocities=velocities), keys.TARGET_GRIP: grip} def to_spec(self): - return {'name': 'joint_delta_action', 'args': {'num_joints': self.num_joints}} + return {'name': self.WIRE_NAME, 'args': {'num_joints': self.num_joints}} diff --git a/positronic/policy/base.py b/positronic/policy/base.py index 46dd2eb6d..ed16a804e 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 +from typing import Any, ClassVar Now = Callable[[], float] @@ -148,11 +148,15 @@ def wrap_session(self, inner: Session, context: dict[str, Any] | None, now: Now """Wrap a single session. Subclasses override this for per-session wrapping.""" raise NotImplementedError('Override wrap_session or wrap') + # The name this wrapper travels under, set by every deliverable subclass. ``WIRE_WRAPPERS`` is keyed by + # it, so the name is written once and both sides of the wire read the same attribute. + WIRE_NAME: ClassVar[str] + def to_spec(self) -> dict[str, Any]: """Plain-data wire spec of this wrapper, for a server's local-stack declaration. Only wrappers registered in ``positronic.policy.spec.WIRE_WRAPPERS`` are deliverable to a rig. - The spec is ``{'name': <wire name>}`` plus ``{'args': {...}}`` when the wrapper takes any; + The spec is ``{'name': WIRE_NAME}`` plus ``{'args': {...}}`` when the wrapper takes any; ``args`` are constructor keywords, since the rig rebuilds by calling the constructor with them. """ raise NotImplementedError(f'{type(self).__name__} is not deliverable to a rig (no wire spec)') diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 16ae3a000..bba31fd04 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -240,6 +240,8 @@ class ActionTimestamp(Codec): At training time, surfaces ``action_fps`` as transform metadata. """ + WIRE_NAME = 'action_timestamp' + def __init__(self, *, fps: float): self._fps = fps self._dt = 1.0 / fps @@ -266,7 +268,7 @@ def meta(self): return {'action_fps': self._fps} def to_spec(self): - return {'name': 'action_timestamp', 'args': {'fps': self._fps}} + return {'name': self.WIRE_NAME, 'args': {'fps': self._fps}} class ActionHorizon(Codec): @@ -285,6 +287,8 @@ class ActionHorizon(Codec): At training time, surfaces ``action_horizon_sec`` as transform metadata. """ + WIRE_NAME = 'action_horizon' + def __init__(self, horizon_sec: float): self._horizon_sec = horizon_sec @@ -310,7 +314,7 @@ def meta(self): return {'action_horizon_sec': self._horizon_sec} def to_spec(self): - return {'name': 'action_horizon', 'args': {'horizon_sec': self._horizon_sec}} + return {'name': self.WIRE_NAME, 'args': {'horizon_sec': self._horizon_sec}} def ActionTiming(*, fps: float, horizon_sec: float | None = None) -> Codec: @@ -335,6 +339,8 @@ class BinarizeGripTraining(Codec): timing | BinarizeGripTraining(('grip', 'target_grip')) | BinarizeGripInference() | obs & action """ + WIRE_NAME = 'binarize_grip_training' + def __init__(self, keys: tuple[str, ...], threshold: float = 0.5): self._keys = keys self._threshold = threshold @@ -361,7 +367,7 @@ def _derive(episode): return Group(Derive(**transforms), Identity()) def to_spec(self): - return {'name': 'binarize_grip_training', 'args': {'keys': list(self._keys), 'threshold': self._threshold}} + return {'name': self.WIRE_NAME, 'args': {'keys': list(self._keys), 'threshold': self._threshold}} class BinarizeGripInference(Codec): @@ -372,6 +378,8 @@ class BinarizeGripInference(Codec): timing | BinarizeGripInference() | obs & action """ + WIRE_NAME = 'binarize_grip_inference' + def __init__(self, threshold: float = 0.5, key: str = obs_keys.TARGET_GRIP): self._threshold = threshold self._key = key @@ -389,7 +397,7 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return data def to_spec(self): - return {'name': 'binarize_grip_inference', 'args': {'threshold': self._threshold, 'key': self._key}} + return {'name': self.WIRE_NAME, 'args': {'threshold': self._threshold, 'key': self._key}} class FlipGrip(Codec): @@ -404,6 +412,8 @@ class FlipGrip(Codec): timing | FlipGrip() | obs & action """ + WIRE_NAME = 'flip_grip' + def encode(self, data): # Copy: the original dict is also the decode ``context`` and the raw recording tap's input. if obs_keys.GRIP in data: @@ -420,7 +430,7 @@ def _decode_single(self, data: dict, context: dict | None) -> dict: return data def to_spec(self): - return {'name': 'flip_grip'} + return {'name': self.WIRE_NAME} def _scaled(image: np.ndarray, width: int, height: int) -> np.ndarray: @@ -444,6 +454,8 @@ class RestrictImageSize(Codec): ChunkedSchedule() | RestrictImageSize() | remote | codec | source """ + WIRE_NAME = 'restrict_image_size' + def __init__(self, width: int = 640, height: int = 640): self._width = width self._height = height @@ -475,7 +487,7 @@ def training_encoder(self) -> EpisodeTransform: ) def to_spec(self): - return {'name': 'restrict_image_size', 'args': {'width': self._width, 'height': self._height}} + return {'name': self.WIRE_NAME, 'args': {'width': self._width, 'height': self._height}} class ChangeEEFrame(Codec): @@ -491,6 +503,8 @@ class ChangeEEFrame(Codec): of the observation/action codecs. """ + WIRE_NAME = 'change_ee_frame' + @staticmethod def _move(value: Any, transform: geom.Transform3D) -> Any: """A pose vector or an arm command, re-expressed through ``transform``.""" @@ -565,6 +579,6 @@ def meta(self): def to_spec(self): # Lists, not tuples, so the spec is identical before and after a wire round-trip. return { - 'name': 'change_ee_frame', + 'name': self.WIRE_NAME, 'args': {'transform': self._transform.as_vector(_QUAT).tolist(), 'keys': list(self._keys)}, } diff --git a/positronic/policy/observation.py b/positronic/policy/observation.py index ba621f2e4..b62c7da0e 100644 --- a/positronic/policy/observation.py +++ b/positronic/policy/observation.py @@ -25,6 +25,8 @@ class ObservationCodec(Codec): task_field: output key carrying the language prompt at inference. """ + WIRE_NAME = 'observation_codec' + def __init__( self, state: dict[str, dict[str, int]], @@ -98,6 +100,6 @@ def to_spec(self): # Normalized to lists so the spec is identical before and after a wire round-trip. images = {name: [key, list(size)] for name, (key, size) in self._image_configs.items()} return { - 'name': 'observation_codec', + 'name': self.WIRE_NAME, 'args': {'state': self._state, 'images': images, 'task_field': self._task_field}, } diff --git a/positronic/policy/spec.py b/positronic/policy/spec.py index fbd9aa99e..f9447a388 100644 --- a/positronic/policy/spec.py +++ b/positronic/policy/spec.py @@ -142,28 +142,25 @@ def load(self, model_id: str, on_progress: Callable[[str], None] | None = None) return self._policy -# TODO(hardcoded-keys): every name here is written twice — once as a literal in the wrapper's own ``to_spec`` -# and once as a key below — so a rename in one place silently desyncs the wire. Let each wrapper own its wire -# name as a class attribute and build this mapping from it. ``test_wire_names_match_table`` is what catches a -# desync until then. -# rules-allow: hardcoded-keys — fixing one entry leaves the rest on the old pattern; the TODO above names the -# whole-table fix, and the test above catches a desync meanwhile WIRE_WRAPPERS: dict[str, type[PolicyWrapper]] = { - 'chunked_schedule': ChunkedSchedule, - 'stop_on_fault': StopOnFault, - 'temporal_stack': TemporalStack, - 'action_timestamp': ActionTimestamp, - 'action_horizon': ActionHorizon, - 'binarize_grip_training': BinarizeGripTraining, - 'binarize_grip_inference': BinarizeGripInference, - 'flip_grip': FlipGrip, - 'restrict_image_size': RestrictImageSize, - 'change_ee_frame': ChangeEEFrame, - 'observation_codec': ObservationCodec, - 'absolute_position_action': AbsolutePositionAction, - 'absolute_joints_action': AbsoluteJointsAction, - 'relative_position_action': RelativePositionAction, - 'joint_delta_action': JointDeltaAction, + wrapper.WIRE_NAME: wrapper + for wrapper in ( + ChunkedSchedule, + StopOnFault, + TemporalStack, + ActionTimestamp, + ActionHorizon, + BinarizeGripTraining, + BinarizeGripInference, + FlipGrip, + RestrictImageSize, + ChangeEEFrame, + ObservationCodec, + AbsolutePositionAction, + AbsoluteJointsAction, + RelativePositionAction, + JointDeltaAction, + ) } diff --git a/positronic/policy/tests/test_wrappers.py b/positronic/policy/tests/test_wrappers.py index 8e041c611..6b30b4637 100644 --- a/positronic/policy/tests/test_wrappers.py +++ b/positronic/policy/tests/test_wrappers.py @@ -409,7 +409,9 @@ def test_non_deliverable_wrapper_fails_loudly(self): with pytest.raises(NotImplementedError, match='not deliverable'): IKJointsAction(solver_cls=None).to_spec() - def test_wire_names_match_table(self): + 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(), 'stop_on_fault': StopOnFault(), diff --git a/positronic/policy/wrappers.py b/positronic/policy/wrappers.py index 27faedb97..abbbabf07 100644 --- a/positronic/policy/wrappers.py +++ b/positronic/policy/wrappers.py @@ -32,6 +32,8 @@ class StopOnFault(PolicyWrapper): observation from anywhere else — a probe replaying a recording — carries no arm to fault. """ + WIRE_NAME = 'stop_on_fault' + class _Session(DelegatingSession): def __call__(self, obs): if not obs.get(keys.ROBOT_FAULT, False): @@ -43,7 +45,7 @@ def wrap_session(self, inner: Session, context, now: Now | None): return StopOnFault._Session(inner) def to_spec(self): - return {'name': 'stop_on_fault'} + return {'name': self.WIRE_NAME} class ChunkedSchedule(PolicyWrapper): @@ -55,6 +57,8 @@ class ChunkedSchedule(PolicyWrapper): until the last action's timestamp is reached, then calls the inner policy. """ + WIRE_NAME = 'chunked_schedule' + class _Session(DelegatingSession): """Skips inner calls while the current trajectory plays; stamps absolute on emit.""" @@ -94,7 +98,7 @@ def wrap_session(self, inner: Session, context, now: Now | None): return ChunkedSchedule._Session(inner, now) def to_spec(self): - return {'name': 'chunked_schedule'} + return {'name': self.WIRE_NAME} class _StackBuffer: @@ -160,6 +164,8 @@ class TemporalStack(PolicyWrapper): chunk-0 empty prefix) never engage it on a padded full-length stack. """ + WIRE_NAME = 'temporal_stack' + class _Session(DelegatingSession): def __init__(self, inner: Session, keys: tuple[str, ...], offsets_sec: tuple[float, ...], pad_start: bool): super().__init__(inner) @@ -189,6 +195,6 @@ def wrap_session(self, inner: Session, context, now: Now | None): def to_spec(self): return { - 'name': 'temporal_stack', + 'name': self.WIRE_NAME, 'args': {'keys': list(self._keys), 'offsets_sec': list(self._offsets_sec), 'pad_start': self._pad_start}, } From 37381ef1f115df0d61bc711ab6f1ea1ef28df15f Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 02:28:08 +0200 Subject: [PATCH 40/64] Name the replay command port once --- positronic/replay_record.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/positronic/replay_record.py b/positronic/replay_record.py index a490b06ca..f3fee43c5 100644 --- a/positronic/replay_record.py +++ b/positronic/replay_record.py @@ -24,6 +24,14 @@ from positronic.utils import package_assets_path from positronic.utils.logging import init_logging +# The port ``Replay`` plays the arm's commands from. A player matches a column to a port by name, so +# ``RestoreCommand`` must derive its column under this same name. +# TODO(#632): it spells the name a third time as a literal instead of reading it here, because ``Derive`` +# takes its columns as ``**kwargs`` and unpacking a dict into them cannot type-check. +# The name reaches no dataset: ``wire`` connects the port through the property below, and the recorder keys +# it ``keys.ROBOT_COMMAND``. +_COMMANDS_CHANNEL = 'robot_commands' + class Replay(DsPlayerAgent): """Adapts `DsPlayerAgent` to be used as a policy control system.""" @@ -35,12 +43,12 @@ def __init__(self, poll_hz: float = 100.0): self.gripper_state = pimm.FakeReceiver(self) self.robot_meta_in = pimm.FakeReceiver(self) self.frames = pimm.ReceiverDict(self, fake=True) - self.outputs['robot_commands'] = pimm.ControlSystemEmitter(self) + self.outputs[_COMMANDS_CHANNEL] = pimm.ControlSystemEmitter(self) self.outputs[keys.TARGET_GRIP] = pimm.ControlSystemEmitter(self) @property def robot_commands(self) -> pimm.ControlSystemEmitter: - return self.outputs['robot_commands'] + return self.outputs[_COMMANDS_CHANNEL] @property def target_grip(self) -> pimm.ControlSystemEmitter: @@ -49,7 +57,7 @@ def target_grip(self) -> pimm.ControlSystemEmitter: class RestoreCommand(Derive): def __init__(self): - super().__init__(robot_commands=self._commands_from_episode) + super().__init__(robot_commands=self._commands_from_episode) # TODO(#632): _COMMANDS_CHANNEL @staticmethod def _commands_from_episode(episode: Episode) -> Any: From 0a2b477ad55d368a68604565f67af8ae4e712b8a Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 02:36:04 +0200 Subject: [PATCH 41/64] Keep `_reap_worker` to its own contract --- positronic/policy/harness.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index b6529a541..bd43ac2a9 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -351,9 +351,7 @@ def _reap_worker(self) -> None: A running thread survives ``shutdown(cancel_futures=True)``, which cancels only what is still queued, so until the join returns the call still holds the session's resources — a ``RemoteSession``'s websocket is the one ``close`` would pull out from under it — and, for an in-process policy, the one - model that every session across episodes and runs shares. The wait sits at the next episode's start - and at the harness's own shutdown, never at the end of the episode the call belongs to, so a hung - model cannot hold up that episode's recording and home. + model that every session across episodes and runs shares. """ if self._retiring is not None: executor, session = self._retiring From 72f373c67088ac5fafced2f2df2ac760ec660339 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 02:44:39 +0200 Subject: [PATCH 42/64] Rank the round the harness read above the trial context --- positronic/policy/harness.py | 7 +++++-- positronic/policy/tests/test_harness.py | 28 +++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index bd43ac2a9..65031effb 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -516,11 +516,14 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: return None if self._awaiting_obs: return None + # The trial's context goes under what the harness read and stamped this round, never over it. A + # context carries whatever keys the RUN directive or the trial plan puts in it, so overlaying it last + # would let a ``robot_state.fault`` in an eval config tell ``StopOnFault`` that a faulted arm is sound. + inputs = {**self.context, **inputs} inputs[keys.ROBOT_FAULT] = faulted inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() - inputs.update(self.context) - inputs['descriptor'] = self._descriptor # last, so a context key can't shadow it + inputs['descriptor'] = self._descriptor return inputs @staticmethod diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 9e61e312b..5cd98219f 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1999,6 +1999,34 @@ def emit_states(): assert policy.last_obs[keys.ROBOT_FAULT] is True +def test_the_trial_context_cannot_stand_in_for_what_the_harness_read(world): + """A context carries whatever the trial plan puts in it. What the harness read and stamped this round + outranks it, so a config key can neither report the arm sound nor supply its state.""" + policy = SpyPolicy() + harness = Harness(policy, make_embodiment()) + harness.commands[keys.ROBOT_COMMAND]._bind(RecordingEmitter()) + harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) + harness.ds_command._bind(RecordingEmitter()) + frame_em = world.pair(harness.observations[CAM]) + robot_em = world.pair(harness.observations[keys.ROBOT_STATE]) + grip_em = world.pair(harness.observations[keys.GRIP]) + directive_em = cast(pimm.SignalEmitter, world.pair(harness.directive)) + + faulted = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.ERROR) + context = {'task': 'test', keys.ROBOT_FAULT: False, keys.GRIP: 'from the config'} + driver = ManualDriver([ + (partial(directive_em.emit, Directive.RUN(**context)), 0.0), + (partial(emit_ready_payload, frame_em, robot_em, grip_em, faulted), 0.01), + (None, 0.02), + ]) + + drive_scheduler(world.start([harness, driver]), steps=40) + + assert policy.last_obs is not None + assert policy.last_obs[keys.ROBOT_FAULT] is True, 'a context key reported the faulted arm sound' + assert policy.last_obs[keys.GRIP] != 'from the config', 'a context key stood in for a channel' + + @pytest.mark.timeout(20.0) def test_a_stop_lands_without_waiting_out_the_charge(world): """A charge places a trajectory's waypoints, and a stop has none: an arm that faults mid-chunk stops in From c81a9738d60dde4a72d6ec3866fc60ccedc7e6c6 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 10:58:09 +0200 Subject: [PATCH 43/64] Declare each command channel's collapse on the channel itself The harness picked a reducer from the channel's name, so an embodiment naming its arm anything else silently dropped delta motion. `Command` now carries it, defaulting to last-wins. --- positronic/cfg/embodiment.py | 14 ++++++--- positronic/eval.py | 14 +++++++-- positronic/policy/harness.py | 28 ++++-------------- positronic/policy/tests/test_harness.py | 37 ++++++++++++++++++++++-- positronic/simulator/env_server/proxy.py | 5 +++- 5 files changed, 65 insertions(+), 33 deletions(-) diff --git a/positronic/cfg/embodiment.py b/positronic/cfg/embodiment.py index 42876b3e7..511557d86 100644 --- a/positronic/cfg/embodiment.py +++ b/positronic/cfg/embodiment.py @@ -30,7 +30,9 @@ def droid(robot_arm, gripper, cameras): **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { - keys.ROBOT_COMMAND: Command(robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command), + keys.ROBOT_COMMAND: Command( + robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce + ), keys.TARGET_GRIP: Command(gripper.target_grip, 0.0, None), } return Embodiment( @@ -53,7 +55,9 @@ def yam(robot_arm, cameras): **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { - keys.ROBOT_COMMAND: Command(robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command), + keys.ROBOT_COMMAND: Command( + robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce + ), keys.TARGET_GRIP: Command(robot_arm.target_grip, 0.0, None), } return Embodiment( @@ -102,7 +106,9 @@ def yam_bimanual(left_channel: str, right_channel: str, mounts: dict[str, list[f } commands = { **{ - f'{keys.ROBOT_COMMAND}.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) + f'{keys.ROBOT_COMMAND}.{s}': Command( + arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce + ) for s, arm in arms.items() }, **{f'{keys.TARGET_GRIP}.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, @@ -142,7 +148,7 @@ def mujoco_franka(sim, camera_dict): # trial's end state right when the operator reviews it. home = roboarm_command.JointPosition(np.array(sim.initial_ctrl[:7])) commands = { - keys.ROBOT_COMMAND: Command(sim.commands, home, Serializers.robot_command), + keys.ROBOT_COMMAND: Command(sim.commands, home, Serializers.robot_command, roboarm_command.reduce), keys.TARGET_GRIP: Command(sim.target_grip, 0.0, None), } return Embodiment( diff --git a/positronic/eval.py b/positronic/eval.py index 039ef2241..4a11ef1dd 100644 --- a/positronic/eval.py +++ b/positronic/eval.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from typing import Any @@ -29,17 +29,25 @@ class Observation: serializer: Serializer | None +def keep_last(due: Sequence[Any]) -> Any: + """The trailing value wins — the collapse for a channel whose values state where to be.""" + return due[-1] + + @dataclass class Command: - """A policy action channel: where its waypoints go and how it homes/records. + """A policy action channel: where its waypoints go, how it homes/records, and how the waypoints one + round finds due collapse into the single value to emit. ``home`` is the value emitted to send this channel to its safe state; ``serializer`` - serializes the channel's values, recorded under the channel's own key. + serializes the channel's values, recorded under the channel's own key. A channel carrying deltas + declares a ``reduce`` that sums their motion, since keeping the last would drop it. """ dest: pimm.SignalReceiver home: Any serializer: Serializer | None + reduce: Callable[[Sequence[Any]], Any] = keep_last @dataclass diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 65031effb..c21137ada 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -39,21 +39,12 @@ Trajectory: TypeAlias = list[tuple[int, Any]] -def _last(due: Sequence[Any]) -> Any: - """The trailing value wins — the right collapse for absolute setpoints and gripper targets.""" - return due[-1] - - class TrajectoryPlayer: """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value - to emit. - - ``reduce`` collapses the waypoints that came due together in one round. Keeping the last is right for a - value that states where to be; a channel carrying deltas passes ``roboarm.command.reduce`` so their - motion is summed rather than dropped. + to emit. ``reduce`` is the channel's own collapse for the waypoints one round finds due together. """ - def __init__(self, reduce: Callable[[Sequence[Any]], Any] = _last): + def __init__(self, reduce: Callable[[Sequence[Any]], Any]): self._pending: deque[tuple[int, Any]] = deque() self._reduce = reduce @@ -264,10 +255,7 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] - self._players = { - name: TrajectoryPlayer(roboarm.command.reduce if keys.is_robot_command(name) else _last) - for name in embodiment.commands - } + self._players = {name: TrajectoryPlayer(cmd.reduce) for name, cmd in embodiment.commands.items()} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) @@ -432,9 +420,7 @@ def _end_episode( """Close the live episode: finalize (or abort) the recording, retire the session, home devices. The session is retired with the worker rather than closed here, so a ``RemoteSession``'s websocket - outlives the call still using it; ``_reap_worker`` closes it at the next episode's start or at - shutdown, and the offboard server's per-session cleanup (active-session decrement, idle watchdog) - runs then. + outlives the call still using it. """ if self._running: if abort: @@ -663,10 +649,8 @@ def _shutdown(self) -> None: """Release the worker and the session. A call still in flight runs to completion and its result is dropped: the run is over and nothing is left to install it. - The harness does not own the policy's lifetime: the caller may run several harnesses over one policy - (a multi-eval sweep), so it closes the policy once, after the last run. That leaves no later - ``_begin_episode`` to reap on, and the next harness reaches the shared policy through a session of - its own, so the call is waited out here. + The join happens here rather than being deferred, since no later episode will do it and the policy + the call holds outlives this harness. """ self._retire_worker() self._reap_worker() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 5cd98219f..6e6f153b3 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -17,7 +17,7 @@ from positronic.drivers.roboarm import RobotStatus from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, JointDelta, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model -from positronic.eval import Command, Embodiment, Observation, Task +from positronic.eval import Command, Embodiment, Observation, Task, keep_last from positronic.geom import Rotation, Transform3D from positronic.offboard.client import InferenceSession from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session @@ -58,7 +58,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simu for cam in cameras: observations[cam] = Observation(pimm.NoOpEmitter(), Serializers.camera_images) commands = { - keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command), + keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command, roboarm.command.reduce), 'target_grip': Command(pimm.NoOpReceiver(), 0.0, None), } return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter(), simulated=simulated) @@ -1353,7 +1353,7 @@ def test_cartesian_delta_without_a_frame_is_rejected(): def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): - player = TrajectoryPlayer() + player = TrajectoryPlayer(keep_last) player.set([(10, 'a'), (20, 'b'), (30, 'c')]) assert player.next_due() == 10 assert player.advance(5) is None @@ -1387,6 +1387,37 @@ def test_trajectory_player_refuses_to_collapse_a_delta_onto_an_absolute(): player.advance(25) +def test_a_channel_collapses_by_what_it_declares_not_by_what_it_is_named(world): + """An embodiment is free to name its channels, so the collapse travels on the channel rather than on its + spelling: an arm reached through a name of its own still sums the deltas a round overtook, and a channel + declaring nothing keeps the last.""" + embodiment = Embodiment( + descriptor='', + observations={'x': Observation(pimm.NoOpEmitter(), None)}, + commands={ + 'arm': Command(pimm.NoOpReceiver(), Reset(), None, roboarm.command.reduce), + 'grip': Command(pimm.NoOpReceiver(), 0.0, None), + }, + static_meta={}, + meta_source=pimm.NoOpEmitter(), + ) + harness = Harness(StubPolicy(), embodiment) + now = world.clock.now() + harness._install( + [ + {'arm': JointDelta(np.array([0.1, 0.0])), 'grip': 0.2, keys.ACTION_TIMESTAMP: now}, + {'arm': JointDelta(np.array([0.2, 0.5])), 'grip': 0.7, keys.ACTION_TIMESTAMP: now}, + ], + world.clock, + ) + now_ns = world.clock.now_ns() + + arm = harness._players['arm'].advance(now_ns) + assert isinstance(arm, JointDelta) + np.testing.assert_allclose(arm.velocities, [0.3, 0.5]) + assert harness._players['grip'].advance(now_ns) == 0.7 + + def test_cartesian_delta_applies_in_world_frame(): current = Transform3D(np.array([0.5, 0.1, 0.3]), Rotation.from_rotvec(np.array([0.2, 0.1, 0.4]))) delta = Transform3D(np.array([0.02, -0.01, 0.05]), Rotation.from_rotvec(np.array([0.1, 0.0, 0.0]))) diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index 4896cf99c..2a61e7716 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -159,7 +159,10 @@ def remote_franka_embodiment( } commands = { keys.ROBOT_COMMAND: Command( - proxy.commands[keys.ROBOT_COMMAND], roboarm_command.Reset(), Serializers.robot_command + proxy.commands[keys.ROBOT_COMMAND], + roboarm_command.Reset(), + Serializers.robot_command, + roboarm_command.reduce, ), keys.TARGET_GRIP: Command(proxy.commands[keys.TARGET_GRIP], 0.0, None), } From b6909ed025f89856ac792c01d69f13c41a6c0c6a Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 18:37:42 +0200 Subject: [PATCH 44/64] Collapse every channel to its last due waypoint Reverts the per-channel reducer. A run of deltas due in one round arrives as its final step alone; `TrajectoryPlayer.advance` states that cost where it is paid. --- positronic/cfg/embodiment.py | 14 ++---- positronic/drivers/roboarm/command.py | 30 ------------ positronic/eval.py | 14 ++---- positronic/policy/harness.py | 35 +++++++------ positronic/policy/tests/test_harness.py | 62 ++---------------------- positronic/simulator/env_server/proxy.py | 5 +- 6 files changed, 33 insertions(+), 127 deletions(-) diff --git a/positronic/cfg/embodiment.py b/positronic/cfg/embodiment.py index 511557d86..42876b3e7 100644 --- a/positronic/cfg/embodiment.py +++ b/positronic/cfg/embodiment.py @@ -30,9 +30,7 @@ def droid(robot_arm, gripper, cameras): **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { - keys.ROBOT_COMMAND: Command( - robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce - ), + keys.ROBOT_COMMAND: Command(robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command), keys.TARGET_GRIP: Command(gripper.target_grip, 0.0, None), } return Embodiment( @@ -55,9 +53,7 @@ def yam(robot_arm, cameras): **{name: Observation(cam.frame, Serializers.camera_images) for name, cam in cameras.items()}, } commands = { - keys.ROBOT_COMMAND: Command( - robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce - ), + keys.ROBOT_COMMAND: Command(robot_arm.commands, roboarm_command.Reset(), Serializers.robot_command), keys.TARGET_GRIP: Command(robot_arm.target_grip, 0.0, None), } return Embodiment( @@ -106,9 +102,7 @@ def yam_bimanual(left_channel: str, right_channel: str, mounts: dict[str, list[f } commands = { **{ - f'{keys.ROBOT_COMMAND}.{s}': Command( - arm.commands, roboarm_command.Reset(), Serializers.robot_command, roboarm_command.reduce - ) + f'{keys.ROBOT_COMMAND}.{s}': Command(arm.commands, roboarm_command.Reset(), Serializers.robot_command) for s, arm in arms.items() }, **{f'{keys.TARGET_GRIP}.{s}': Command(arm.target_grip, 0.0, None) for s, arm in arms.items()}, @@ -148,7 +142,7 @@ def mujoco_franka(sim, camera_dict): # trial's end state right when the operator reviews it. home = roboarm_command.JointPosition(np.array(sim.initial_ctrl[:7])) commands = { - keys.ROBOT_COMMAND: Command(sim.commands, home, Serializers.robot_command, roboarm_command.reduce), + keys.ROBOT_COMMAND: Command(sim.commands, home, Serializers.robot_command), keys.TARGET_GRIP: Command(sim.target_grip, 0.0, None), } return Embodiment( diff --git a/positronic/drivers/roboarm/command.py b/positronic/drivers/roboarm/command.py index b6372627b..f7c62e016 100644 --- a/positronic/drivers/roboarm/command.py +++ b/positronic/drivers/roboarm/command.py @@ -1,6 +1,5 @@ """Collection of commands that can be sent to the robot.""" -from collections.abc import Sequence from dataclasses import dataclass, field from typing import Any @@ -120,32 +119,3 @@ def from_wire(wire: dict[str, Any]) -> CommandType: ) case _: raise ValueError(f'Unknown command type: {wire["type"]}') - - -def _combine(acc: CommandType, cmd: CommandType) -> CommandType: - match (acc, cmd): - case (CartesianDelta(a, frame_a), CartesianDelta(b, frame_b)): - if not np.allclose(frame_a.as_matrix, frame_b.as_matrix): - raise ValueError('Cannot accumulate cartesian deltas expressed in different frames') - return CartesianDelta(_compose_delta(a, b), frame_a) - case (JointDelta(a), JointDelta(b)): - return JointDelta(a + b) - case (CartesianDelta() | JointDelta(), _) | (_, CartesianDelta() | JointDelta()): - raise ValueError(f'Cannot reduce {type(acc).__name__} then {type(cmd).__name__} in one tick') - case _: - return cmd - - -def reduce(due: Sequence[CommandType]) -> CommandType: - """Collapse the commands due in one round into the single command to execute. - - Folds the batch in timestamp order. A run of same-space deltas accumulates (their motion is summed, so a - round spanning several waypoints catches up rather than dropping them); a run of absolute commands keeps - the last. Mixing an absolute with a delta, or two delta spaces, has no faithful single-command form — a - delta binds to the pose measured when it is consumed, which an absolute target or a foreign space cannot - supply — and raises. - """ - result = due[0] - for cmd in due[1:]: - result = _combine(result, cmd) - return result diff --git a/positronic/eval.py b/positronic/eval.py index 4a11ef1dd..039ef2241 100644 --- a/positronic/eval.py +++ b/positronic/eval.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Sequence +from collections.abc import Callable from dataclasses import dataclass, field from typing import Any @@ -29,25 +29,17 @@ class Observation: serializer: Serializer | None -def keep_last(due: Sequence[Any]) -> Any: - """The trailing value wins — the collapse for a channel whose values state where to be.""" - return due[-1] - - @dataclass class Command: - """A policy action channel: where its waypoints go, how it homes/records, and how the waypoints one - round finds due collapse into the single value to emit. + """A policy action channel: where its waypoints go and how it homes/records. ``home`` is the value emitted to send this channel to its safe state; ``serializer`` - serializes the channel's values, recorded under the channel's own key. A channel carrying deltas - declares a ``reduce`` that sums their motion, since keeping the last would drop it. + serializes the channel's values, recorded under the channel's own key. """ dest: pimm.SignalReceiver home: Any serializer: Serializer | None - reduce: Callable[[Sequence[Any]], Any] = keep_last @dataclass diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index c21137ada..75d7f65a4 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -2,7 +2,7 @@ import logging import time from collections import deque -from collections.abc import Callable, Generator, Iterable, Iterator, Sequence +from collections.abc import Generator, Iterable, Iterator from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum @@ -41,12 +41,10 @@ class TrajectoryPlayer: """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value - to emit. ``reduce`` is the channel's own collapse for the waypoints one round finds due together. - """ + to emit.""" - def __init__(self, reduce: Callable[[Sequence[Any]], Any]): + def __init__(self): self._pending: deque[tuple[int, Any]] = deque() - self._reduce = reduce def set(self, trajectory: Trajectory): self._pending = deque(trajectory) @@ -56,12 +54,17 @@ def next_due(self) -> int | None: return self._pending[0][0] if self._pending else None def advance(self, current_time: int): - """The single value due at ``current_time``, collapsed by ``reduce`` when several came due since the - previous call, or ``None`` when none did.""" - due = [] + """The single value due at ``current_time`` — the last, when several came due since the previous + call — or ``None`` when none did. + + Collapsing to the last is exact for an absolute setpoint and lossy for a relative one: a run of + deltas due together arrives as its final step alone. Pacing keeps one waypoint due per round + wherever a round is shorter than the spacing between waypoints. + """ + value = None while self._pending and self._pending[0][0] <= current_time: - due.append(self._pending.popleft()[1]) - return self._reduce(due) if due else None + value = self._pending.popleft()[1] + return value def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: @@ -255,7 +258,7 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] - self._players = {name: TrajectoryPlayer(cmd.reduce) for name, cmd in embodiment.commands.items()} + self._players = {name: TrajectoryPlayer() for name in embodiment.commands} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) @@ -420,7 +423,9 @@ def _end_episode( """Close the live episode: finalize (or abort) the recording, retire the session, home devices. The session is retired with the worker rather than closed here, so a ``RemoteSession``'s websocket - outlives the call still using it. + outlives the call still using it; ``_reap_worker`` closes it at the next episode's start or at + shutdown, and the offboard server's per-session cleanup (active-session decrement, idle watchdog) + runs then. """ if self._running: if abort: @@ -649,8 +654,10 @@ def _shutdown(self) -> None: """Release the worker and the session. A call still in flight runs to completion and its result is dropped: the run is over and nothing is left to install it. - The join happens here rather than being deferred, since no later episode will do it and the policy - the call holds outlives this harness. + The harness does not own the policy's lifetime: the caller may run several harnesses over one policy + (a multi-eval sweep), so it closes the policy once, after the last run. That leaves no later + ``_begin_episode`` to reap on, and the next harness reaches the shared policy through a session of + its own, so the call is waited out here. """ self._retire_worker() self._reap_worker() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 6e6f153b3..792b566f3 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -15,9 +15,9 @@ from positronic.dataset.serializers import Serializers from positronic.drivers import roboarm from positronic.drivers.roboarm import RobotStatus -from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, JointDelta, Reset, from_wire, to_wire +from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model -from positronic.eval import Command, Embodiment, Observation, Task, keep_last +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, Policy, PolicyWrapper, Session @@ -58,7 +58,7 @@ def make_embodiment(descriptor: str = '', cameras=(CAM,), static_meta=None, simu for cam in cameras: observations[cam] = Observation(pimm.NoOpEmitter(), Serializers.camera_images) commands = { - keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command, roboarm.command.reduce), + keys.ROBOT_COMMAND: Command(pimm.NoOpReceiver(), Reset(), Serializers.robot_command), 'target_grip': Command(pimm.NoOpReceiver(), 0.0, None), } return Embodiment(descriptor, observations, commands, static_meta or {}, pimm.NoOpEmitter(), simulated=simulated) @@ -1353,7 +1353,7 @@ def test_cartesian_delta_without_a_frame_is_rejected(): def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): - player = TrajectoryPlayer(keep_last) + player = TrajectoryPlayer() player.set([(10, 'a'), (20, 'b'), (30, 'c')]) assert player.next_due() == 10 assert player.advance(5) is None @@ -1364,60 +1364,6 @@ def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): assert player.advance(40) is None -def test_trajectory_player_sums_the_deltas_a_late_round_overtook(): - """A delta states how far to move, so the ones a round overtook are motion still owed. Keeping only the - last would silently shorten the trajectory.""" - player = TrajectoryPlayer(roboarm.command.reduce) - player.set([(10, JointDelta(np.array([0.1, 0.0]))), (20, JointDelta(np.array([0.2, 0.5])))]) - - caught_up = player.advance(25) - - assert isinstance(caught_up, JointDelta) - np.testing.assert_allclose(caught_up.velocities, [0.3, 0.5]) - - -def test_trajectory_player_refuses_to_collapse_a_delta_onto_an_absolute(): - """No single command carries both: a delta binds to the pose measured when it is consumed, which an - absolute target cannot supply.""" - player = TrajectoryPlayer(roboarm.command.reduce) - pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) - player.set([(10, CartesianPosition(pose=pose)), (20, JointDelta(np.array([0.1, 0.0])))]) - - with pytest.raises(ValueError, match='Cannot reduce'): - player.advance(25) - - -def test_a_channel_collapses_by_what_it_declares_not_by_what_it_is_named(world): - """An embodiment is free to name its channels, so the collapse travels on the channel rather than on its - spelling: an arm reached through a name of its own still sums the deltas a round overtook, and a channel - declaring nothing keeps the last.""" - embodiment = Embodiment( - descriptor='', - observations={'x': Observation(pimm.NoOpEmitter(), None)}, - commands={ - 'arm': Command(pimm.NoOpReceiver(), Reset(), None, roboarm.command.reduce), - 'grip': Command(pimm.NoOpReceiver(), 0.0, None), - }, - static_meta={}, - meta_source=pimm.NoOpEmitter(), - ) - harness = Harness(StubPolicy(), embodiment) - now = world.clock.now() - harness._install( - [ - {'arm': JointDelta(np.array([0.1, 0.0])), 'grip': 0.2, keys.ACTION_TIMESTAMP: now}, - {'arm': JointDelta(np.array([0.2, 0.5])), 'grip': 0.7, keys.ACTION_TIMESTAMP: now}, - ], - world.clock, - ) - now_ns = world.clock.now_ns() - - arm = harness._players['arm'].advance(now_ns) - assert isinstance(arm, JointDelta) - np.testing.assert_allclose(arm.velocities, [0.3, 0.5]) - assert harness._players['grip'].advance(now_ns) == 0.7 - - def test_cartesian_delta_applies_in_world_frame(): current = Transform3D(np.array([0.5, 0.1, 0.3]), Rotation.from_rotvec(np.array([0.2, 0.1, 0.4]))) delta = Transform3D(np.array([0.02, -0.01, 0.05]), Rotation.from_rotvec(np.array([0.1, 0.0, 0.0]))) diff --git a/positronic/simulator/env_server/proxy.py b/positronic/simulator/env_server/proxy.py index 2a61e7716..4896cf99c 100644 --- a/positronic/simulator/env_server/proxy.py +++ b/positronic/simulator/env_server/proxy.py @@ -159,10 +159,7 @@ def remote_franka_embodiment( } commands = { keys.ROBOT_COMMAND: Command( - proxy.commands[keys.ROBOT_COMMAND], - roboarm_command.Reset(), - Serializers.robot_command, - roboarm_command.reduce, + proxy.commands[keys.ROBOT_COMMAND], roboarm_command.Reset(), Serializers.robot_command ), keys.TARGET_GRIP: Command(proxy.commands[keys.TARGET_GRIP], 0.0, None), } From ccf66d149fbc7b9a34390896f2b7aa18fdd88e32 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 18:46:54 +0200 Subject: [PATCH 45/64] Finish the last episode's business before the next one's span opens Its abandoned call was joined after `_telemetry.begin`, so the wait landed inside the new episode. Restores two docstring trims an over-broad revert undid. --- positronic/policy/harness.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 75d7f65a4..23f020a2d 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -385,6 +385,9 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: armed here and moved to that first observation once it lands, so an episode that never gets one is still bounded. """ + # Before the span opens, so the wait for a call the last episode abandoned is inter-episode wall + # rather than overhead the timing reducer attributes to this one. + self._reap_worker() self.context = dict(context) if self._embodiment.simulated: # A sim trial that doesn't ask for latency simulation runs free of it: the world holds still for @@ -406,7 +409,6 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._reap_worker() # Arm the clock before handing it out: a session reading it before its first call must see this # episode's start, not the release time of the last episode's final call. self._t0_ns = clock.now_ns() @@ -423,9 +425,7 @@ def _end_episode( """Close the live episode: finalize (or abort) the recording, retire the session, home devices. The session is retired with the worker rather than closed here, so a ``RemoteSession``'s websocket - outlives the call still using it; ``_reap_worker`` closes it at the next episode's start or at - shutdown, and the offboard server's per-session cleanup (active-session decrement, idle watchdog) - runs then. + outlives the call still using it. """ if self._running: if abort: @@ -654,10 +654,8 @@ def _shutdown(self) -> None: """Release the worker and the session. A call still in flight runs to completion and its result is dropped: the run is over and nothing is left to install it. - The harness does not own the policy's lifetime: the caller may run several harnesses over one policy - (a multi-eval sweep), so it closes the policy once, after the last run. That leaves no later - ``_begin_episode`` to reap on, and the next harness reaches the shared policy through a session of - its own, so the call is waited out here. + The join happens here rather than being deferred, since no later episode will do it and the policy + the call holds outlives this harness. """ self._retire_worker() self._reap_worker() From 4e77288e2d8b3ad9290829a06b5b093328b96654 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 18:56:06 +0200 Subject: [PATCH 46/64] State which observation keys a faulted arm withholds --- docs/connect-your-model.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index de7a59a58..81e155f8b 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -109,7 +109,7 @@ The client sends the full raw robot state as a dict. Keys are flat strings (the | `task` | str | — | Language instruction for the episode | | `descriptor` | str | — | Embodiment the observation came from (e.g. `mujoco.franka`); empty string when unset. Lets a multi-embodiment policy adapt to the current robot | -Your server receives every key each step. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. +Your server receives every key each step, except that an arm with no sound sample to give — faulted, or mid-reset — omits its `robot_state.*` measurements and sends only `robot_state.fault`. The standard stack puts `StopOnFault` ahead of the model, which answers such a step itself and never forwards it; a stack without it reaches the model with those keys missing. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. ### Actions (server → client) From 661fa6a6f9b213e37338e64799e22e57c890d945 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 19:20:22 +0200 Subject: [PATCH 47/64] Finish #418's rename in the replay port The port and its derived column kept the pre-rename plural. Both ends are internal, so nothing broke and nothing flagged it; they now read `keys.ROBOT_COMMAND` like every other site. --- positronic/replay_record.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/positronic/replay_record.py b/positronic/replay_record.py index f3fee43c5..04014a3c0 100644 --- a/positronic/replay_record.py +++ b/positronic/replay_record.py @@ -24,14 +24,6 @@ from positronic.utils import package_assets_path from positronic.utils.logging import init_logging -# The port ``Replay`` plays the arm's commands from. A player matches a column to a port by name, so -# ``RestoreCommand`` must derive its column under this same name. -# TODO(#632): it spells the name a third time as a literal instead of reading it here, because ``Derive`` -# takes its columns as ``**kwargs`` and unpacking a dict into them cannot type-check. -# The name reaches no dataset: ``wire`` connects the port through the property below, and the recorder keys -# it ``keys.ROBOT_COMMAND``. -_COMMANDS_CHANNEL = 'robot_commands' - class Replay(DsPlayerAgent): """Adapts `DsPlayerAgent` to be used as a policy control system.""" @@ -43,12 +35,12 @@ def __init__(self, poll_hz: float = 100.0): self.gripper_state = pimm.FakeReceiver(self) self.robot_meta_in = pimm.FakeReceiver(self) self.frames = pimm.ReceiverDict(self, fake=True) - self.outputs[_COMMANDS_CHANNEL] = pimm.ControlSystemEmitter(self) + self.outputs[keys.ROBOT_COMMAND] = pimm.ControlSystemEmitter(self) self.outputs[keys.TARGET_GRIP] = pimm.ControlSystemEmitter(self) @property def robot_commands(self) -> pimm.ControlSystemEmitter: - return self.outputs[_COMMANDS_CHANNEL] + return self.outputs[keys.ROBOT_COMMAND] @property def target_grip(self) -> pimm.ControlSystemEmitter: @@ -56,8 +48,13 @@ def target_grip(self) -> pimm.ControlSystemEmitter: class RestoreCommand(Derive): + """Derives the column ``Replay`` plays the arm's commands from: a player matches a column to a port by + name, so this must carry the name that port is registered under.""" + def __init__(self): - super().__init__(robot_commands=self._commands_from_episode) # TODO(#632): _COMMANDS_CHANNEL + # ``Derive`` takes its columns as ``**kwargs`` beside a keyword-only ``meta``, so a checker reading an + # unpacked dict cannot rule out ``meta`` and matches the transform against its type. + super().__init__(**{keys.ROBOT_COMMAND: self._commands_from_episode}) # pyright: ignore[reportArgumentType] @staticmethod def _commands_from_episode(episode: Episode) -> Any: From 265ce41dbae92fb904582f8c9e501b9aeaab1a6f Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 19:56:40 +0200 Subject: [PATCH 48/64] Condense the `Harness` docstring --- .claude/skills/address-review/SKILL.md | 9 ++++++--- positronic/drivers/gripper/dh.py | 2 +- positronic/policy/harness.py | 27 +++++++++----------------- 3 files changed, 16 insertions(+), 22 deletions(-) diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index 04a388d91..f5144ad7c 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -111,8 +111,11 @@ For each open comment, decide and note severity if the bot tagged one (e.g. Code resolve only if you land a concrete change, else leave open. - **Decline** — wrong, not applicable, or contradicts a deliberate decision → reasoned reply, **leave open**. -- **Defer** — valid but out of scope for this PR → file the tracking issue unless something - already tracks it, reply naming it, **leave open**. +- **Defer** — valid but out of scope for this PR → reply saying so, **leave open**, and name it in + the Step 6 report as wanting a tracking issue. **Never open an issue without asking first.** An + issue is filed under the user's own account, so it reads to their team as their judgement about + what matters; that call is theirs. Propose the title and the substance and wait for a yes. If an + existing issue already covers it, cite that instead — citing is not filing. - **Discuss** — the reviewer is asking a question or opening a design discussion, not requesting a change → answer it, **leave open** for them to respond. @@ -313,7 +316,7 @@ Their verdicts are terse. What each one means: | "next" / "resolve, next" | reply, resolve, present the next thread — same shape, no re-asking | | "leave it" | reply that it is accepted as a known limit, resolve | | "fix it" | fix, run the gates, commit, push, reply with the SHA, resolve | -| "file it" / "open a bug" | file the issue, reply linking it, resolve | +| "file it" / "open a bug" | file the issue, reply linking it, resolve — this verdict is the only thing that authorizes filing one | | "why declined?" | the reasoning only | This is the one path on which a declined thread gets resolved: the human closed it, thread by thread. diff --git a/positronic/drivers/gripper/dh.py b/positronic/drivers/gripper/dh.py index 6a27f5864..9b3a6b524 100644 --- a/positronic/drivers/gripper/dh.py +++ b/positronic/drivers/gripper/dh.py @@ -34,7 +34,7 @@ def _state_r(): last_grip = 0.0 - # TODO: We must translate these to physical units (N and m/s) + # TODO: Should we translate these to physical units (N and m/s)? while not should_stop.value: try: grip_msg = self.target_grip.read() diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 23f020a2d..dae7da534 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -39,6 +39,7 @@ Trajectory: TypeAlias = list[tuple[int, Any]] +# TODO(624): This class is likely to go away when this issue is resolved. class TrajectoryPlayer: """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value to emit.""" @@ -184,24 +185,14 @@ 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. - Handles directives (RUN/FINISH/ABORT) and dataset recording. Inference intelligence — scheduling, - error recovery, blending, absolute time stamping — lives in the policy/session layer: the wrapper owns - the plan, the harness plays it, one command per channel per round. The session call runs on a worker - thread so playing continues while the model does; the harness withholds the trajectory, and the world - clock, until the trial's inference charge (``inference_latency``) is paid, and the ``now`` it hands - ``new_session`` reads time the same way — so wrappers stamp chunks for the paid instant without - knowing the mode. The RUN context is handed whole to the task's scene reset, which reads the - per-trial keys it needs (e.g. ``eval.seed``). - - A ``trials`` plan (a sequence of RUN contexts) makes the harness self-driving: it starts the next trial - whenever idle and returns once the plan is exhausted, so the unattended path needs no driver. A task's - ``timeout`` bounds every trial, self-driven or operator-driven, so an attended episode still terminates - at the deadline if the operator never sends FINISH. A bounded trial also ends early on a truthy - privileged ``done``, recording ``eval.terminated`` True and the delivered payload in its static data; a - timeout records False. A task-less session has neither deadline nor budget and ends only on directives. - - The ``Embodiment`` supplies the observation serializers (which own the canonical key names), the command - channels and the home action. The policy owns its wrapper stack; the harness runs what it is given. + The wrapper owns the plan, the harness plays it, one command per channel per round. The session call + runs on a worker so playing continues while the model does, and its trajectory — with the world clock — + is withheld until the trial's ``inference_latency`` is paid; the ``now`` handed to ``new_session`` reads + time the same way, so wrappers stamp for the paid instant without knowing the mode. + + A ``trials`` plan makes the harness self-driving: it starts the next trial whenever idle and returns + once the plan is exhausted. A task's ``timeout`` bounds every trial either way, and a truthy privileged + ``done`` ends one early — ``eval.terminated`` records which. A task-less session ends only on directives. """ def __init__( From 321d7946553beb8b03ad2c169cdee5bc763758eb Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:03:45 +0200 Subject: [PATCH 49/64] Spell out what an inference charge does to the world clock --- positronic/policy/harness.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index dae7da534..3c7447021 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -227,9 +227,10 @@ def __init__( # was submitted. self._t0_ns = 0 self._wall_t0 = 0.0 - # Seconds each model call costs the world clock this episode, the same figure for every call. - # ``None`` when no figure is fixed and the world is charged the call's own wall duration (hardware - # pace, and the sim's ``inference_latency=True``). + # What one model call costs this episode's world clock, in seconds. A fixed figure — ``0.0`` or more — + # holds the world still until the call answers, then releases its chunk that many seconds after the + # call began, so the trial runs the same however fast the machine is. ``None`` charges the call's own + # wall duration instead: what hardware pays, and what the sim asks for with ``inference_latency=True``. self._fixed_latency: float | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None From a7061c3ccef6c89148bede9477852a7c0a28da1c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:05:37 +0200 Subject: [PATCH 50/64] Spell out what an inference charge does to the world clock --- positronic/policy/harness.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 3c7447021..c34d61e49 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -227,10 +227,8 @@ def __init__( # was submitted. self._t0_ns = 0 self._wall_t0 = 0.0 - # What one model call costs this episode's world clock, in seconds. A fixed figure — ``0.0`` or more — - # holds the world still until the call answers, then releases its chunk that many seconds after the - # call began, so the trial runs the same however fast the machine is. ``None`` charges the call's own - # wall duration instead: what hardware pays, and what the sim asks for with ``inference_latency=True``. + # Seconds of world clock a model call costs. A number: world frozen for the call, chunk released that + # long after it began. ``None``: charged the call's real wall time (hardware; ``inference_latency=True``). self._fixed_latency: float | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None From a97fb1ae831105ca46b88f1cec96f2eb9e531d84 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:21:08 +0200 Subject: [PATCH 51/64] Read the trial's inference charge without writing it back --- positronic/policy/harness.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index c34d61e49..428f6454d 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -378,15 +378,9 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: # Before the span opens, so the wait for a call the last episode abandoned is inter-episode wall # rather than overhead the timing reducer attributes to this one. self._reap_worker() - self.context = dict(context) - if self._embodiment.simulated: - # A sim trial that doesn't ask for latency simulation runs free of it: the world holds still for - # every model call. - latency = self.context.setdefault(keys.INFERENCE_LATENCY, False) - else: - # A fixed latency is a device for simulating a trial, so a real rig ignores it and pays the wall - # time its calls really take. - latency = True + self.context = context + # A sim trial without the key runs free; a real rig ignores this sim-only knob and pays wall time. + latency = self.context.get(keys.INFERENCE_LATENCY, False) if self._embodiment.simulated else True self._fixed_latency = None if latency is True else float(latency) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False From 4f49190e5ef85fbe8f567c1503381fd5de87da91 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:24:40 +0200 Subject: [PATCH 52/64] Arm the call clock only where a call is submitted --- positronic/policy/harness.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 428f6454d..956bd366a 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -393,10 +393,6 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - # Arm the clock before handing it out: a session reading it before its first call must see this - # episode's start, not the release time of the last episode's final call. - self._t0_ns = clock.now_ns() - self._wall_t0 = time.monotonic() self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') self._policy_session = self.policy.new_session(self.context, self._effect_time) self._running = True From d5b4401668b9dfbbc8ea23f0d9bcbc3c764af609 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:36:42 +0200 Subject: [PATCH 53/64] Cut harness prose and the duplicates it was hiding `seal` closes through `end` instead of repeating it, `_take` asks `_effect_time` for the instant it was already recomputing, and the cached `_descriptor` gives way to the embodiment that owns it. --- positronic/policy/harness.py | 175 ++++++++++++++--------------------- 1 file changed, 72 insertions(+), 103 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 956bd366a..e0f96b69c 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -21,18 +21,16 @@ from positronic.policy.base import Policy, Session from positronic.utils import flatten_dict, frozen_view -# How far from now an action may be scheduled. A chunk spans seconds, so this is loose enough that no real -# trajectory approaches it, and tight enough to catch a rig-side stack that left timestamps relative to the -# chunk (decades behind) or anchored them twice (decades ahead). +# 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 directive is # noticed, and with it the granularity every command timestamp is quantized to. POLL_PERIOD_SEC = 0.01 -# How long a submitted session call may take to answer and still resolve within its round. A wrapper that -# skips inference answers in microseconds; a real model call runs far past this and is then paced by -# ``_take`` across rounds. +# How long a submitted call may take and still resolve within its round: a wrapper that skips inference +# answers in microseconds, a real model call runs far past this and is paced across rounds by ``_take``. SKIP_REPLY_SEC = 0.001 # One channel's schedule: waypoints stamped with absolute clock ns, ascending. @@ -55,12 +53,11 @@ def next_due(self) -> int | None: return self._pending[0][0] if self._pending else None def advance(self, current_time: int): - """The single value due at ``current_time`` — the last, when several came due since the previous - call — or ``None`` when none did. + """The value due at ``current_time`` — the last, when several came due since the previous call — or + ``None`` when none did. - Collapsing to the last is exact for an absolute setpoint and lossy for a relative one: a run of - deltas due together arrives as its final step alone. Pacing keeps one waypoint due per round - wherever a round is shorter than the spacing between waypoints. + Collapsing to the last is exact for an absolute setpoint and lossy for a relative one. Pacing keeps + one waypoint due per round wherever a round is shorter than the spacing between waypoints. """ value = None while self._pending and self._pending[0][0] <= current_time: @@ -108,12 +105,9 @@ def ABORT(cls) -> 'Directive': class _EpisodeTelemetry: - """The live rollout's wall-clock telemetry: the episode span, its index, its step count and the virtual - instant it began. Inert while telemetry is unbound, so the harness calls it unconditionally. - - The span stays anchored while open, so the rollout's phase spans (reset, env.step, policy.infer, - record.io) parent to it rather than to the pass. - """ + """The live rollout's episode span, with the index, step count and virtual start it is stamped with. + Inert while telemetry is unbound, so the harness calls it unconditionally. The span stays anchored while + open, so the rollout's phase spans parent to it rather than to the pass.""" def __init__(self) -> None: self._span: Span | None = None @@ -141,8 +135,8 @@ def step(self) -> None: self._steps += 1 def end(self, virtual_now: float) -> None: - """Close a finished rollout, stamped with its step count and its virtual duration up to - ``virtual_now`` — captured when the rollout ended, before the flush round advances the sim clock.""" + """Close the rollout, stamped with its step count and its virtual duration up to ``virtual_now`` — + captured when the rollout ended, before the flush round advances the sim clock.""" if self._span is None: return self._close(virtual_now) @@ -156,15 +150,13 @@ def abort(self) -> None: self._end_span() def seal(self, virtual_now: float) -> None: - """Close a rollout abandoned mid-flight by a raising ``reset`` / ``new_session`` / session call, stamped - like a clean end and marked ``episode.partial``. Ending it is what exports it — the batch processor drops - an unended span, orphaning the finished children and losing their phases. Partial rather than aborted so - the reduce keeps it. Inert when no span is open.""" + """Close a rollout abandoned mid-flight by a raising ``reset`` / ``new_session`` / session call, marked + ``episode.partial`` so the reduce keeps it. Ending it is what exports it: the batch processor drops an + unended span, orphaning the finished children and losing their phases.""" if self._span is None: return telemetry.set_attrs(self._span, **{telemetry_keys.ATTR_EPISODE_PARTIAL: True}) - self._close(virtual_now) - telemetry.force_flush() + self.end(virtual_now) def _close(self, virtual_now: float) -> None: # A rollout whose first observation never landed — a reset that raised, or a task already done before @@ -215,16 +207,13 @@ def __init__( self._policy_session: Session | None = None # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. self._running = False - # One session call at a time, on a worker so the harness keeps playing while the model runs. The - # worker belongs to the episode: ending one abandons the call in flight rather than waiting for it, - # so the next episode must not queue behind it. + # One call at a time, on a worker so the harness keeps playing while the model runs. The worker + # belongs to the episode: ending one abandons the call in flight, so the next must not queue behind it. self._executor: ThreadPoolExecutor | None = None - # The retired worker and the session its abandoned call is still inside, held as one because closing - # the session is what the join makes safe. + # A retired worker and the session its abandoned call is inside: the join is what makes closing safe. self._retiring: tuple[ThreadPoolExecutor, Session] | None = None self._future: Future[list[dict[str, Any]] | None] | None = None - # The in-flight call's start: the world instant its observation was built, and the wall instant it - # was submitted. + # The in-flight call's start: the world instant of its observation, the wall instant of its submit. self._t0_ns = 0 self._wall_t0 = 0.0 # Seconds of world clock a model call costs. A number: world frozen for the call, chunk released that @@ -232,8 +221,8 @@ def __init__( self._fixed_latency: float | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None - # Whether this episode's first observation has landed. Until it does the deadline stands where the - # reset put it, which bounds an episode whose first observation never arrives. + # False until this episode's first observation lands; until then the deadline stands where the reset + # put it, which bounds an episode that never gets one. self._rollout_started = False # Wall-clock telemetry for the live rollout, opened under ``--timing`` and inert otherwise. self._telemetry = _EpisodeTelemetry() @@ -241,7 +230,6 @@ def __init__( # emptying this set is what keeps the first inference off the previous episode's final frame. self._awaiting_obs: set[str] = set() - self._descriptor = embodiment.descriptor self.observations = pimm.ReceiverDict(self) self.commands = pimm.EmitterDict(self) for name in embodiment.observations: @@ -290,17 +278,16 @@ def _pace(self, clock: pimm.Clock) -> pimm.Command: waypoint is emitted at its own time and a round rarely finds more than one due.""" if self._embodiment.simulated: return pimm.Yield() - due = [ts for player in self._players.values() if (ts := player.next_due()) is not None] - if not due: + due = min((ts for player in self._players.values() if (ts := player.next_due()) is not None), default=None) + if due is None: return pimm.Sleep(POLL_PERIOD_SEC) - return pimm.Sleep(min(POLL_PERIOD_SEC, max(min(due) - clock.now_ns(), 1) / 1e9)) + return pimm.Sleep(min(POLL_PERIOD_SEC, max(due - clock.now_ns(), 1) / 1e9)) def _cancel_session(self) -> None: """Drop everything the episode has going: the schedule being played, and the call on the worker. - The call is let go of rather than waited for — a model that hangs must not hold up the recording's - stop or the home — so its worker is retired with it and whatever it eventually answers, or raises, - lands nowhere. Devices hold their last commanded position; nothing is buffered downstream to clear. + The call is let go of rather than waited for, so a model that hangs cannot hold up the recording's + stop or the home. Devices hold their last commanded position; nothing downstream is buffered. """ for player in self._players.values(): player.set([]) @@ -329,9 +316,8 @@ def _retire_worker(self) -> None: def _reap_worker(self) -> None: """Wait out an abandoned call, then close the session it was inside. - A running thread survives ``shutdown(cancel_futures=True)``, which cancels only what is still queued, - so until the join returns the call still holds the session's resources — a ``RemoteSession``'s - websocket is the one ``close`` would pull out from under it — and, for an in-process policy, the one + ``shutdown(cancel_futures=True)`` cancels only what is still queued, so until the join returns the + call still holds the session's resources: a ``RemoteSession``'s websocket, or the one in-process model that every session across episodes and runs shares. """ if self._retiring is not None: @@ -353,15 +339,14 @@ def _finalize_recording( # Give the recorder a round to commit the STOP before the next START (they share ``ds_command``, where # last-value-wins would drop one) and before the home command, so homing stays out of the recording. yield self._pace(clock) - # After that round, so the recorder's STOP-time record.io span is still in flight and parents to the - # episode. Accepted skew: a producer stepping in that shared round charges one span (≤ one control - # period) to the closing episode — the cooperative scheduler cannot give the recorder a turn alone. + # After that round, so the recorder's STOP-time record.io span still parents to the episode. Skew: a + # producer stepping in that shared round charges ≤ one control period to the closing episode. self._telemetry.end(virtual_now) def _effect_time(self) -> float: """The trial instant the in-flight call's output takes effect: its observation instant plus the - charge — the declared constant whole, or the wall time elapsed so far. Read on the worker thread; - the loop thread writes the call's start fields before submitting it. + charge — the declared constant whole, or the wall time elapsed so far. The loop thread writes the + call's start fields before the submit, so a worker reading this sees its own call's. """ charge = time.monotonic() - self._wall_t0 if self._fixed_latency is None else self._fixed_latency return self._t0_ns / 1e9 + charge @@ -369,11 +354,9 @@ def _effect_time(self) -> float: def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. - A resettable task's ``reset`` only arms the producer, which publishes the first observation on a later - round. The recorder drains its channels the turn it opens, so the pre-reset frame and the - inter-episode home command drop out and its first sample is the post-reset scene. The deadline is - armed here and moved to that first observation once it lands, so an episode that never gets one is - still bounded. + A resettable task's ``reset`` only arms the producer; the first observation lands a later round. The + recorder drains its channels the turn it opens, so the pre-reset frame and the inter-episode home + command drop out. The deadline is armed here and moved to that first observation once it lands. """ # Before the span opens, so the wait for a call the last episode abandoned is inter-episode wall # rather than overhead the timing reducer attributes to this one. @@ -462,13 +445,11 @@ def _read_channel(self, name: str, obs: Observation) -> tuple[dict[str, Any] | N def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: """Read every observation channel and assemble the policy input dict. - Raises ``NoValueException`` if any channel has no value yet. Returns ``None`` while a serializer - reports a sample is not ready (``robot_state`` during a ``RESETTING`` arm) or a channel still holds a - pre-reset value — either way the harness skips inference rather than feed a partial or stale obs. - - A faulted arm is the exception: it has no sample either, but the plan being played was made for an - arm that is now somewhere else, so the observation goes to the policy stack carrying - ``keys.ROBOT_FAULT`` and without the arm's own entries. + Raises ``NoValueException`` if any channel has no value yet. Returns ``None`` while a sample is not + ready (``robot_state`` during a ``RESETTING`` arm) or a channel still holds a pre-reset value, rather + than feed a partial or stale obs. A faulted arm is the exception: the plan being played was made for + an arm that is now somewhere else, so its observation reaches the stack carrying ``keys.ROBOT_FAULT`` + and without the arm's own entries. """ # Against the live model, not the one known at episode start: a remote env publishes its ``robot_meta`` # a turn after the reset that produced it, so at episode start there is no model to check. @@ -483,18 +464,16 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs.update(entries or {}) # Every channel is read before this decision, so a bimanual rig cannot hide one arm's fault behind # another arm's not-ready sample: whichever channel comes first, the fault still reaches the stack. - if not_ready and not faulted: + if (not_ready and not faulted) or self._awaiting_obs: return None - if self._awaiting_obs: - return None - # The trial's context goes under what the harness read and stamped this round, never over it. A - # context carries whatever keys the RUN directive or the trial plan puts in it, so overlaying it last - # would let a ``robot_state.fault`` in an eval config tell ``StopOnFault`` that a faulted arm is sound. + # The trial's context goes under what the harness read this round, never over it: a context carries + # whatever keys the RUN directive puts in it, and a ``robot_state.fault`` among them must not tell + # ``StopOnFault`` that a faulted arm is sound. inputs = {**self.context, **inputs} inputs[keys.ROBOT_FAULT] = faulted inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() - inputs['descriptor'] = self._descriptor + inputs['descriptor'] = self._embodiment.descriptor return inputs @staticmethod @@ -502,9 +481,8 @@ def _owned(obs: dict[str, Any]) -> dict[str, Any]: """The observation with its arrays copied, so nothing rewrites what the worker is still reading. A producer may reuse one buffer for every sample it emits — a camera renders into the array behind - the adapter it re-emits each frame — and the loop thread yields while a call charged in wall time - runs, so that producer advances alongside the worker. Copying at dispatch pays once per call rather - than per round. + the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. + Copying at dispatch pays once per call rather than per round. """ return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} @@ -529,9 +507,8 @@ def _step(self, clock: pimm.Clock) -> None: if obs is None: return if not self._rollout_started: - # The rollout begins at its first observation, not when the reset returned: a reset only asks - # the producer for a scene, and the turns spent delivering it are neither the trial's budget - # nor its duration. + # The rollout begins at its first observation, not when the reset returned: the turns spent + # delivering the scene are neither the trial's budget nor its duration. self._rollout_started = True self._telemetry.start_rollout(clock.now()) if self._task is not None: @@ -546,19 +523,18 @@ def _step(self, clock: pimm.Clock) -> None: def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: """Install the call's trajectory once the world has paid for it. - Under a constant charge the world holds still until the call answers — blocking here blocks the loop - thread, which is what advances a virtual clock. Until a call answers there is no telling a skip from - a model call, so letting the world run meanwhile would spend trial time on whichever the machine - turned out to be slow at. What the charge then buys is the instant a trajectory takes effect: it is - stamped for ``t0`` plus the charge and withheld until the world reaches it, playing what is already - scheduled on the way. An answer with no waypoints to place — a skip, or the empty trajectory that - stops what is executing — has no such instant and lands at once. A charge measured in wall time can - hold nothing still, so there the world runs no further ahead of the call's start than wall time has. + Constant charge: blocking here holds the virtual clock still, so the trial is billed the declared + charge and not what the machine took — until a call answers, a skip and a model call look alike. The + trajectory is withheld until the world reaches ``t0`` plus the charge, playing what is already + scheduled on the way; an answer with no waypoints — a skip, or the empty trajectory that stops what + is executing — lands at once. + + Wall charge: nothing can be held still, so the world runs no further ahead of ``t0`` than wall time. """ if self._fixed_latency is not None: concurrent.futures.wait([future]) elif not future.done(): - ahead = clock.now() - (self._t0_ns / 1e9 + time.monotonic() - self._wall_t0) + ahead = clock.now() - self._effect_time() if ahead <= 0.0: return Harness._Answer.PENDING concurrent.futures.wait([future], timeout=ahead) @@ -586,8 +562,8 @@ def _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: _assert_anchored(actions, clock.now()) self._telemetry.step() for name, player in self._players.items(): - # Wrappers do action-timing math in float seconds; the schedule and every pimm channel are in ns. - # This is the single explicit seconds->ns seam. + # The single explicit seconds->ns seam: wrappers time actions in float seconds, the schedule and + # every pimm channel are in ns. player.set([(int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in actions if name in a]) def _play(self, clock: pimm.Clock) -> None: @@ -601,11 +577,10 @@ def _play(self, clock: pimm.Clock) -> None: def _trial_terminal(self, clock: pimm.Clock) -> dict[str, Any] | None: """The terminal static payload if a self-driven trial has ended this round, else ``None``. - The deadline is hard: a truthy ``done`` delivered within budget records ``eval.terminated`` True plus - its payload, the budget passing records False, and a terminal past the deadline is a timeout rather - than a late success. Only a freshly delivered ``done`` counts — the receiver latches its last value, - so a prior trial's terminal would otherwise re-fire; gating on delivery clears it without asking the - producer to republish. Reached only for a task with a deadline. + The deadline is hard: a truthy ``done`` within budget records ``eval.terminated`` True plus its + payload, the budget passing records False, and a terminal past the deadline is a timeout rather than + a late success. Only a freshly delivered ``done`` counts, or the receiver's latched value would + re-fire a prior trial's terminal. Reached only for a task with a deadline. """ done_msg = self.done.read() if done_msg.updated and done_msg.data and done_msg.ts <= self._deadline * 1e9: @@ -622,31 +597,25 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p try: yield from self._run(should_stop, clock) except BaseException: - # A failure mid-rollout unwinds past the normal span close. Seal the open span before the - # exception reaches ``bind``'s exit flush, or it never exports and its finished children orphan, - # losing their phases and charging the episode's wall to between_episodes. + # Seal the open span before the exception reaches ``bind``'s exit flush: an unended span never + # exports, orphaning its finished children and charging the episode's wall to between_episodes. self._telemetry.seal(clock.now()) raise finally: self._shutdown() def _shutdown(self) -> None: - """Release the worker and the session. A call still in flight runs to completion and its result is - dropped: the run is over and nothing is left to install it. - - The join happens here rather than being deferred, since no later episode will do it and the policy - the call holds outlives this harness. - """ + """Release the worker and the session: a call still in flight runs to completion and its result is + dropped. The join happens here rather than being deferred, since no later episode will do it and the + policy the call holds outlives this harness.""" self._retire_worker() self._reap_worker() def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: while not should_stop.value: # One action per round, mutually exclusive: handle a directive, start the next trial (or exit - # when the plan is done), finish a self-driven trial that is out of budget or done, or step the - # policy. Starting takes its own round so a begin never shares one with a step — inference waits - # for the producer's post-reset observation, which the recorder logs once its open-turn drain has - # cleared the channels. + # when the plan is done), finish one that is out of budget or done, or step the policy. Starting + # takes its own round, so inference waits for the producer's post-reset observation. directive_msg = self.directive.read() # Read every round so the flag clears mid-episode; a press during a trial is consumed, not replayed. manual_msg = self.manual_command.read() From 2a303d147dca4f09363c6b81195babdc0bdc7e67 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 20:42:18 +0200 Subject: [PATCH 54/64] Hold each channel's schedule as a plain deque `TrajectoryPlayer` wrapped one deque behind three methods; the harness now keeps `dict[str, deque]` and plays it in `_play`, where the collapse-to-last rule is stated. --- positronic/policy/harness.py | 63 ++++++++----------------- positronic/policy/tests/test_harness.py | 17 +------ 2 files changed, 22 insertions(+), 58 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index e0f96b69c..d98c9101b 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -6,7 +6,7 @@ from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum -from typing import Any, TypeAlias +from typing import Any import numpy as np from opentelemetry.trace import Span @@ -33,37 +33,6 @@ # answers in microseconds, a real model call runs far past this and is paced across rounds by ``_take``. SKIP_REPLY_SEC = 0.001 -# One channel's schedule: waypoints stamped with absolute clock ns, ascending. -Trajectory: TypeAlias = list[tuple[int, Any]] - - -# TODO(624): This class is likely to go away when this issue is resolved. -class TrajectoryPlayer: - """Plays one channel's schedule: ``set()`` a trajectory, then ``advance(now)`` each round for the value - to emit.""" - - def __init__(self): - self._pending: deque[tuple[int, Any]] = deque() - - def set(self, trajectory: Trajectory): - self._pending = deque(trajectory) - - def next_due(self) -> int | None: - """Timestamp of the earliest waypoint not yet played, or ``None`` once the schedule is exhausted.""" - return self._pending[0][0] if self._pending else None - - def advance(self, current_time: int): - """The value due at ``current_time`` — the last, when several came due since the previous call — or - ``None`` when none did. - - Collapsing to the last is exact for an absolute setpoint and lossy for a relative one. Pacing keeps - one waypoint due per round wherever a round is shorter than the spacing between waypoints. - """ - value = None - while self._pending and self._pending[0][0] <= current_time: - value = self._pending.popleft()[1] - return value - def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: """Reject a chunk whose timestamps are not times on the harness clock.""" @@ -236,7 +205,8 @@ def __init__( self.observations[name] # touch to allocate the port for name in embodiment.commands: self.commands[name] - self._players = {name: TrajectoryPlayer() for name in embodiment.commands} + # 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} self.directive = pimm.ControlSystemReceiver[Directive](self, default=None, maxsize=3) self.manual_command = pimm.ControlSystemReceiver(self, default=None) @@ -278,7 +248,7 @@ def _pace(self, clock: pimm.Clock) -> pimm.Command: waypoint is emitted at its own time and a round rarely finds more than one due.""" if self._embodiment.simulated: return pimm.Yield() - due = min((ts for player in self._players.values() if (ts := player.next_due()) is not None), default=None) + 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)) @@ -289,8 +259,8 @@ def _cancel_session(self) -> None: The call is let go of rather than waited for, so a model that hangs cannot hold up the recording's stop or the home. Devices hold their last commanded position; nothing downstream is buffered. """ - for player in self._players.values(): - player.set([]) + for schedule in self._schedules.values(): + schedule.clear() self._retire_worker() @staticmethod @@ -561,16 +531,23 @@ def _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: """ _assert_anchored(actions, clock.now()) self._telemetry.step() - for name, player in self._players.items(): - # The single explicit seconds->ns seam: wrappers time actions in float seconds, the schedule and - # every pimm channel are in ns. - player.set([(int(a[keys.ACTION_TIMESTAMP] * 1e9), a[name]) for a in actions if name in a]) + # The single explicit seconds->ns seam: wrappers 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 actions 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.""" + """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, player in self._players.items(): - value = player.advance(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) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 792b566f3..6242e5cd9 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -22,7 +22,7 @@ from positronic.offboard.client import InferenceSession from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import Directive, DirectiveType, Harness, TrajectoryPlayer, _assert_anchored +from positronic.policy.harness import Directive, DirectiveType, Harness, _assert_anchored from positronic.policy.remote import RemoteSession from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -1352,18 +1352,6 @@ def test_cartesian_delta_without_a_frame_is_rejected(): from_wire(wire) -def test_trajectory_player_collapses_several_due_waypoints_to_the_last(): - player = TrajectoryPlayer() - player.set([(10, 'a'), (20, 'b'), (30, 'c')]) - assert player.next_due() == 10 - assert player.advance(5) is None - assert player.advance(25) == 'b' # a late round overtakes 'a'; the trailing setpoint is the live one - assert player.next_due() == 30 - assert player.advance(30) == 'c' - assert player.next_due() is None - assert player.advance(40) is None - - def test_cartesian_delta_applies_in_world_frame(): current = Transform3D(np.array([0.5, 0.1, 0.3]), Rotation.from_rotvec(np.array([0.2, 0.1, 0.4]))) delta = Transform3D(np.array([0.02, -0.01, 0.05]), Rotation.from_rotvec(np.array([0.1, 0.0, 0.0]))) @@ -1680,8 +1668,7 @@ def test_a_reply_is_installed_only_while_the_trial_still_has_budget(world, expir harness._take(future, world.clock) - played = harness._players[keys.ROBOT_COMMAND].advance(world.clock.now_ns()) - assert (played is not None) is installed + assert bool(harness._schedules[keys.ROBOT_COMMAND]) is installed class _SlowSession(Session): From f73f545d1e02b39019f53115a971cafdc36046ac Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 21:30:35 +0200 Subject: [PATCH 55/64] Give the harness a charge and a worker instead of six timing fields The declared-vs-wall latency mode was re-tested in four places; it is now the choice of one `_Charge` at episode start. `_InferenceWorker` takes the session, the executor and the future, so `_take` drops from 22 lines to 9 and stops carrying a future it always had. --- positronic/policy/harness.py | 294 ++++++++++++++---------- positronic/policy/tests/test_harness.py | 19 +- 2 files changed, 186 insertions(+), 127 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index d98c9101b..f1303a28a 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,8 +1,9 @@ import concurrent.futures import logging import time +from abc import ABC, abstractmethod from collections import deque -from collections.abc import Generator, Iterable, Iterator +from collections.abc import Generator, Iterable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum @@ -34,16 +35,6 @@ SKIP_REPLY_SEC = 0.001 -def _assert_anchored(actions: 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 actions), 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 rig-side ' - f'stack is not anchoring chunks to the harness clock' - ) - - class DirectiveType(Enum): RUN = 'run' FINISH = 'finish' @@ -73,6 +64,140 @@ def ABORT(cls) -> 'Directive': return cls(DirectiveType.ABORT) +class _Answer(Enum): + """Whether the in-flight call's answer is the harness's to act on this round.""" + + READY = 'ready' + PENDING = 'pending' + + +class _Charge(ABC): + """What one session call costs the trial, in world-clock seconds.""" + + def __init__(self) -> None: + self._t0_ns = 0 + + def begin(self, t0_ns: int) -> None: + """Stamp the start of the call about to be submitted.""" + self._t0_ns = t0_ns + + @abstractmethod + def effect_time(self) -> float: + """The trial instant the in-flight call's output takes effect.""" + + @abstractmethod + def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: + """Block the loop for the call as far as this charge allows, then say where its answer stands.""" + + +class _DeclaredCharge(_Charge): + """A sim trial's ``inference_latency``: the call costs this, whatever the machine took.""" + + def __init__(self, seconds: float) -> None: + super().__init__() + self._ns = round(seconds * 1e9) + + def effect_time(self) -> float: + return (self._t0_ns + self._ns) / 1e9 + + def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: + # Blocking the loop thread holds a virtual clock still, which is what bills the trial the declared + # charge and not what the machine took. + concurrent.futures.wait([future]) + # An answer with nothing to place — a skip, or the empty trajectory that stops what is executing — + # has no effect instant to wait for. Integer ns, the world's own timeline: a float compare misses + # the instant by one ULP and slips the install a full round. + due = not future.result() or clock.now_ns() >= self._t0_ns + self._ns + return _Answer.READY if due else _Answer.PENDING + + +class _WallCharge(_Charge): + """A call costs the wall time it took: a real rig, and a sim asked for real latency.""" + + def __init__(self) -> None: + super().__init__() + self._wall_t0 = 0.0 + + def begin(self, t0_ns: int) -> None: + super().begin(t0_ns) + self._wall_t0 = time.monotonic() + + def effect_time(self) -> float: + return self._t0_ns / 1e9 + time.monotonic() - self._wall_t0 + + def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: + # The timeout is the sim's overrun over wall time, handed back: wall time cannot be held still, so + # the world runs no further ahead of the call's start than wall time has. An answer that is in has + # nothing left to wait for. + concurrent.futures.wait([future], timeout=max(clock.now() - self.effect_time(), 0.0)) + return _Answer.READY if future.done() else _Answer.PENDING + + +def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: + """Report the failure of a call nobody is waiting for any more.""" + # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no + # caller left to raise to; the log is the only place its failure can go. + if not future.cancelled() and (exc := future.exception()) is not None: + logging.error(f'Inference failed after the episode that asked for it ended: {exc}') + + +class _InferenceWorker: + """One episode's policy session, called one at a time on a thread of its own so the harness keeps + playing while the model runs. Ending an episode ``abandon``s the call in flight rather than waiting for + it, so the next episode never queues behind a model that hangs.""" + + def __init__(self, session: Session, charge: _Charge) -> None: + self._session = session + self._charge = charge + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') + self._future: Future[list[dict[str, Any]] | None] | None = None + + @property + def meta(self) -> dict[str, Any]: + return self._session.meta + + @property + def idle(self) -> bool: + """Whether a call could be submitted: none is in flight.""" + return self._future is None + + def submit(self, obs: Mapping[str, Any], clock: pimm.Clock) -> None: + """Start a call on ``obs``, charged from this instant. The moment's wait lets a wrapper that skips + inference resolve in the round it was asked.""" + self._charge.begin(clock.now_ns()) + self._future = self._executor.submit(self._session, obs) + concurrent.futures.wait([self._future], timeout=SKIP_REPLY_SEC) + + def wait(self, clock: pimm.Clock) -> _Answer: + """Block for the in-flight call as far as its charge allows, then say where its answer stands.""" + assert self._future is not None + return self._charge.wait(self._future, clock) + + def collect(self) -> list[dict[str, Any]] | None: + """The ready call's trajectory — ``None`` when it had nothing to place — leaving the worker idle.""" + assert self._future is not None + # Read on the loop thread, so a failing call still seals the episode. + actions, self._future = self._future.result(), None + return actions + + def abandon(self) -> None: + """Let go of the call in flight: its answer lands nowhere and its failure only reaches the log.""" + if self._future is not None: + self._future.add_done_callback(_report_abandoned) + self._future = None + self._executor.shutdown(wait=False, cancel_futures=True) + + def join(self) -> None: + """Wait out an abandoned call, then close the session it was inside. + + ``shutdown(cancel_futures=True)`` cancels only what is still queued, so until this returns the call + still holds the session's resources: a ``RemoteSession``'s websocket, or the in-process model every + session shares. + """ + self._executor.shutdown(wait=True) + self._session.close() + + class _EpisodeTelemetry: """The live rollout's episode span, with the index, step count and virtual start it is stamped with. Inert while telemetry is unbound, so the harness calls it unconditionally. The span stays anchored while @@ -173,21 +298,12 @@ def __init__( self.policy: Policy = policy self.context: dict[str, Any] = {} self._static_meta = static_meta or {} - self._policy_session: Session | None = None # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. self._running = False - # One call at a time, on a worker so the harness keeps playing while the model runs. The worker - # belongs to the episode: ending one abandons the call in flight, so the next must not queue behind it. - self._executor: ThreadPoolExecutor | None = None - # A retired worker and the session its abandoned call is inside: the join is what makes closing safe. - self._retiring: tuple[ThreadPoolExecutor, Session] | None = None - self._future: Future[list[dict[str, Any]] | None] | None = None - # The in-flight call's start: the world instant of its observation, the wall instant of its submit. - self._t0_ns = 0 - self._wall_t0 = 0.0 - # Seconds of world clock a model call costs. A number: world frozen for the call, chunk released that - # long after it began. ``None``: charged the call's real wall time (hardware; ``inference_latency=True``). - self._fixed_latency: float | None = None + # This episode's session and the thread it runs on; ``None`` while no episode is live. + self._worker: _InferenceWorker | None = None + # A worker let go of mid-call, kept until the join that makes closing its session safe. + self._retiring: _InferenceWorker | None = None # ``task.timeout``, set per episode; a task-less session has no deadline and ends on directives. self._deadline: float | None = None # False until this episode's first observation lands; until then the deadline stands where the reset @@ -229,7 +345,7 @@ def _build_episode_meta(self, context: dict[str, Any]) -> dict[str, Any]: meta['eval.timeout'] = self._task.timeout # ``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._policy_session.meta if self._policy_session else {}) + session_meta = self.policy.meta | (self._worker.meta if self._worker else {}) for k, v in flatten_dict(session_meta).items(): meta[f'{keys.POLICY_META}.{k}'] = v meta.update(context) @@ -263,45 +379,25 @@ def _cancel_session(self) -> None: schedule.clear() self._retire_worker() - @staticmethod - def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: - """Report the failure of a call nobody is waiting for any more.""" - # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no - # caller left to raise to; the log is the only place its failure can go. - if not future.cancelled() and (exc := future.exception()) is not None: - logging.error(f'Inference failed after the episode that asked for it ended: {exc}') - def _retire_worker(self) -> None: - """Let go of this episode's worker, the call it is running and the session that call is inside: the - answer lands nowhere and the failure only reaches the log. All three are kept for ``_reap_worker``, - since ending an episode must not wait for a model that hangs.""" - if self._future is not None: - self._future.add_done_callback(self._report_abandoned) - self._future = None - if self._executor is not None and self._policy_session is not None: - self._executor.shutdown(wait=False, cancel_futures=True) - self._retiring = (self._executor, self._policy_session) - self._executor, self._policy_session = None, None + """Let go of this episode's worker and the call it is running, keeping it for ``_reap_worker``: + ending an episode must not wait for a model that hangs.""" + if self._worker is not None: + self._worker.abandon() + self._retiring, self._worker = self._worker, None def _reap_worker(self) -> None: - """Wait out an abandoned call, then close the session it was inside. - - ``shutdown(cancel_futures=True)`` cancels only what is still queued, so until the join returns the - call still holds the session's resources: a ``RemoteSession``'s websocket, or the one in-process - model that every session across episodes and runs shares. - """ + """Join the retired worker and close the session its abandoned call was inside.""" if self._retiring is not None: - executor, session = self._retiring + self._retiring.join() self._retiring = None - executor.shutdown(wait=True) - session.close() 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 episode's full static meta (plus any terminal payload) — then close its span.""" - # Stamped before the session is retired with the worker: the meta overlays what the session reports. + # Stamped before the worker is retired: the meta overlays what its session reports. stop = DsWriterCommand.STOP({**self._build_episode_meta(self.context), **(payload or {})}) self._cancel_session() self.ds_command.emit(stop) @@ -313,14 +409,6 @@ def _finalize_recording( # producer stepping in that shared round charges ≤ one control period to the closing episode. self._telemetry.end(virtual_now) - def _effect_time(self) -> float: - """The trial instant the in-flight call's output takes effect: its observation instant plus the - charge — the declared constant whole, or the wall time elapsed so far. The loop thread writes the - call's start fields before the submit, so a worker reading this sees its own call's. - """ - charge = time.monotonic() - self._wall_t0 if self._fixed_latency is None else self._fixed_latency - return self._t0_ns / 1e9 + charge - def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: """Open a fresh episode: reset the scene, fix the task context and session, and open the recording. @@ -334,7 +422,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self.context = context # A sim trial without the key runs free; a real rig ignores this sim-only knob and pays wall time. latency = self.context.get(keys.INFERENCE_LATENCY, False) if self._embodiment.simulated else True - self._fixed_latency = None if latency is True else float(latency) + charge = _WallCharge() if latency is True else _DeclaredCharge(float(latency)) self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -346,8 +434,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') - self._policy_session = self.policy.new_session(self.context, self._effect_time) + self._worker = _InferenceWorker(self.policy.new_session(self.context, charge.effect_time), charge) self._running = True self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -357,8 +444,8 @@ def _end_episode( ) -> Generator[pimm.Command, None, None]: """Close the live episode: finalize (or abort) the recording, retire the session, home devices. - The session is retired with the worker rather than closed here, so a ``RemoteSession``'s websocket - outlives the call still using it. + The worker is retired rather than joined here, so a ``RemoteSession``'s websocket outlives the call + still using it. """ if self._running: if abort: @@ -456,22 +543,10 @@ 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()} - class _Answer(Enum): - """What became of the call ``_take`` was handed: its trajectory is installed and the future spent, or - the world has not yet paid the charge and the same future comes back next round.""" - - CONSUMED = 'consumed' - PENDING = 'pending' - def _step(self, clock: pimm.Clock) -> None: - """Keep one session call in flight and install the trajectory it returns. - - The call goes to the worker so the harness keeps playing while the model runs; a wrapper that - answers without inference still resolves in the round it was asked. - """ - session, executor = self._policy_session, self._executor - assert session is not None and executor is not None # only a live episode steps - if self._future is not None and self._take(self._future, clock) is Harness._Answer.PENDING: + """Keep one session call in flight and install the trajectory it returns.""" + assert self._worker is not None # only a live episode steps + if not self._worker.idle and self._take(clock) is _Answer.PENDING: return obs = self._build_obs(clock) if obs is None: @@ -483,53 +558,38 @@ def _step(self, clock: pimm.Clock) -> None: self._telemetry.start_rollout(clock.now()) if self._task is not None: self._deadline = clock.now() + self._task.timeout - self._t0_ns = clock.now_ns() - self._wall_t0 = time.monotonic() - self._future = executor.submit(session, frozen_view(self._owned(obs))) - if self._fixed_latency is None: - concurrent.futures.wait([self._future], timeout=SKIP_REPLY_SEC) - self._take(self._future, clock) - - def _take(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: - """Install the call's trajectory once the world has paid for it. - - Constant charge: blocking here holds the virtual clock still, so the trial is billed the declared - charge and not what the machine took — until a call answers, a skip and a model call look alike. The - trajectory is withheld until the world reaches ``t0`` plus the charge, playing what is already - scheduled on the way; an answer with no waypoints — a skip, or the empty trajectory that stops what - is executing — lands at once. - - Wall charge: nothing can be held still, so the world runs no further ahead of ``t0`` than wall time. - """ - if self._fixed_latency is not None: - concurrent.futures.wait([future]) - elif not future.done(): - ahead = clock.now() - self._effect_time() - if ahead <= 0.0: - return Harness._Answer.PENDING - concurrent.futures.wait([future], timeout=ahead) - if not future.done(): - return Harness._Answer.PENDING - actions = future.result() # taken on the loop thread, so a failing call still seals the episode - if actions and self._fixed_latency is not None: - # Integer ns, the world's own timeline: a float compare misses the release instant by one ULP - # and slips the install a full round. - if clock.now_ns() < self._t0_ns + round(self._fixed_latency * 1e9): - return Harness._Answer.PENDING # the schedule already playing carries the world to the release instant - self._future = None + self._worker.submit(frozen_view(self._owned(obs)), clock) + self._take(clock) + + def _take(self, clock: pimm.Clock) -> _Answer: + """Install the in-flight call's trajectory once the world has paid its charge.""" + assert self._worker is not None + if self._worker.wait(clock) is _Answer.PENDING: + return _Answer.PENDING # the schedule already playing carries the world to the effect instant + actions = self._worker.collect() # The world reached the deadline while the call 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. expired = self._deadline is not None and clock.now() >= self._deadline if actions is not None and not expired: self._install(actions, clock) - return Harness._Answer.CONSUMED + return _Answer.READY + + @staticmethod + def _assert_anchored(actions: 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 actions), 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 _install(self, actions: 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 wrapper for the instant its charge is paid. """ - _assert_anchored(actions, clock.now()) + self._assert_anchored(actions, clock.now()) self._telemetry.step() # The single explicit seconds->ns seam: wrappers time actions in float seconds, the schedules and # every pimm channel are in ns. diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 6242e5cd9..e4c764e30 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -1,6 +1,5 @@ import threading import time -from concurrent.futures import Future from contextlib import contextmanager from functools import partial from types import SimpleNamespace @@ -22,7 +21,7 @@ from positronic.offboard.client import InferenceSession from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import Directive, DirectiveType, Harness, _assert_anchored +from positronic.policy.harness import Directive, DirectiveType, Harness, _InferenceWorker, _WallCharge from positronic.policy.remote import RemoteSession from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -1641,32 +1640,32 @@ def episode(gap_steps: int) -> float: 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'): - _assert_anchored([{'timestamp': 0.0}], now=1.7e9) + 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'): - _assert_anchored([{'timestamp': 3.4e9}], now=1.7e9) + 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.""" - _assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) + Harness._assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) @pytest.mark.parametrize(('expired', 'installed'), [(True, False), (False, True)]) def test_a_reply_is_installed_only_while_the_trial_still_has_budget(world, expired, installed): """A trial advertises the instant it stops at. A call whose rounds in flight carried the world past that instant has its chunk dropped instead of placed, and ``_run`` finishes the trial on the next round.""" - harness = Harness(StubPolicy(), make_embodiment()) + policy = StubPolicy() + harness = Harness(policy, make_embodiment()) now = world.clock.now() harness._deadline = now - 1.0 if expired else now + 1.0 - pose = Transform3D(translation=np.array([0.4, 0.5, 0.6], dtype=np.float32), rotation=Rotation.identity) - future = Future() - future.set_result([{keys.ROBOT_COMMAND: CartesianPosition(pose=pose), keys.ACTION_TIMESTAMP: now}]) + harness._worker = _InferenceWorker(policy.new_session(), _WallCharge()) + harness._worker.submit({}, world.clock) - harness._take(future, world.clock) + harness._take(world.clock) assert bool(harness._schedules[keys.ROBOT_COMMAND]) is installed From 042cf041f33b0f0a3fc9af918c77a7c29f41e2b9 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 21:44:55 +0200 Subject: [PATCH 56/64] Take the harness and its tests off the type-check baseline `_trial_terminal` narrows `_deadline` itself instead of trusting the caller's guard, so the task-less case is stated where it applies. The test stubs get the annotations their `None`-initialised fields always needed. --- .basedpyright/baseline.json | 94 +------------------------ positronic/policy/harness.py | 12 ++-- positronic/policy/tests/test_harness.py | 13 ++-- 3 files changed, 18 insertions(+), 101 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index ef5e98a09..cd984ef26 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -5632,56 +5632,7 @@ } } ], - "./positronic/policy/harness.py": [ - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 20, - "endColumn": 27, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 41, - "endColumn": 45, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 59, - "endColumn": 61, - "lineCount": 1 - } - }, - { - "code": "reportOptionalOperand", - "range": { - "startColumn": 65, - "endColumn": 79, - "lineCount": 1 - } - }, - { - "code": "reportOptionalMemberAccess", - "range": { - "startColumn": 31, - "endColumn": 35, - "lineCount": 1 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 11, - "endColumn": 40, - "lineCount": 1 - } - } - ], + "./positronic/policy/harness.py": [], "./positronic/policy/observation.py": [ { "code": "reportArgumentType", @@ -5786,48 +5737,7 @@ } } ], - "./positronic/policy/tests/test_harness.py": [ - { - "code": "reportIncompatibleMethodOverride", - "range": { - "startColumn": 8, - "endColumn": 19, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 65, - "endColumn": 73, - "lineCount": 1 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 15, - "endColumn": 32, - "lineCount": 1 - } - }, - { - "code": "reportOptionalSubscript", - "range": { - "startColumn": 11, - "endColumn": 36, - "lineCount": 1 - } - }, - { - "code": "reportArgumentType", - "range": { - "startColumn": 15, - "endColumn": 45, - "lineCount": 1 - } - } - ], + "./positronic/policy/tests/test_harness.py": [], "./positronic/policy/tests/test_policy_io.py": [ { "code": "reportArgumentType", diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index f1303a28a..a6634528b 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -617,12 +617,16 @@ def _trial_terminal(self, clock: pimm.Clock) -> dict[str, Any] | None: The deadline is hard: a truthy ``done`` within budget records ``eval.terminated`` True plus its payload, the budget passing records False, and a terminal past the deadline is a timeout rather than a late success. Only a freshly delivered ``done`` counts, or the receiver's latched value would - re-fire a prior trial's terminal. Reached only for a task with a deadline. + re-fire a prior trial's terminal. """ + deadline = self._deadline + if deadline is None: # a task-less session has no budget and ends on directives alone + return None done_msg = self.done.read() - if done_msg.updated and done_msg.data and done_msg.ts <= self._deadline * 1e9: + assert done_msg is not None # the receiver carries a default, so ``read`` always yields a message + if done_msg.updated and done_msg.data and done_msg.ts <= deadline * 1e9: return {**done_msg.data, keys.EVAL_TERMINATED: True} - if clock.now() >= self._deadline: + if clock.now() >= deadline: return {keys.EVAL_TERMINATED: False} return None @@ -669,7 +673,7 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ yield pimm.Sleep(0.5) break self._begin_episode(trial, clock) - elif self._deadline is not None and (terminal := self._trial_terminal(clock)) is not None: + elif (terminal := self._trial_terminal(clock)) is not None: yield from self._end_episode(clock, terminal) else: try: diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index e4c764e30..90fcfb80e 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -119,16 +119,16 @@ def __init__( self.command = command self.target_grip = float(target_grip) self.last_obs: dict[str, Any] | None = None - self.observations: list[dict[str, object]] = [] + self.observations: list[dict[str, Any]] = [] self.reset_calls = 0 - self.last_reset_context = None + self.last_reset_context: dict[str, Any] | None = None self._meta: dict[str, object] = meta or {} @property def meta(self) -> dict[str, object]: return self._meta - def new_session(self, context=None, now=None): + def new_session(self, context=None, now=None) -> Session: self.reset_calls += 1 self.last_reset_context = context return _StubSession(self) @@ -234,7 +234,7 @@ def make_robot_state(translation, joints, status=RobotStatus.AVAILABLE) -> FakeR def emit_ready_payload(frame_emitter, robot_emitter, grip_emitter, robot_state): - frame_adapter = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.uint8) + frame_adapter = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) frame_adapter.array[:] = np.zeros((2, 2, 3), dtype=np.uint8) frame_emitter.emit(frame_adapter) robot_emitter.emit(robot_state) @@ -1146,6 +1146,7 @@ def reset(_context): scheduler = world.start([harness]) drive_scheduler(scheduler, steps=200) + assert policy.last_reset_context is not None assert policy.last_reset_context[keys.TASK] == 'resolved-on-reset' @@ -1370,7 +1371,9 @@ def test_robot_state_serializer_drops_not_ready(status): def test_robot_state_serializer_available_has_no_error_key(): state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6], status=RobotStatus.AVAILABLE) - assert set(Serializers.robot_state(state)) == {'.q', '.dq', '.ee_pose'} + serialized = Serializers.robot_state(state) + assert serialized is not None + assert set(serialized) == {'.q', '.dq', '.ee_pose'} @pytest.mark.timeout(3.0) From 3f592edb3b738bffc05054fcc41c749fda73549d Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 21:48:41 +0200 Subject: [PATCH 57/64] Fold `_shutdown` into the `finally` that calls it The body named what the name did: two already-named calls. --- positronic/policy/harness.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index a6634528b..fe35a216a 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -643,14 +643,10 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self._telemetry.seal(clock.now()) raise finally: - self._shutdown() - - def _shutdown(self) -> None: - """Release the worker and the session: a call still in flight runs to completion and its result is - dropped. The join happens here rather than being deferred, since no later episode will do it and the - policy the call holds outlives this harness.""" - self._retire_worker() - self._reap_worker() + # A call still in flight runs to completion and its result is dropped. The join is not deferred: + # no later episode will do it, and the policy the call holds outlives this harness. + self._retire_worker() + self._reap_worker() def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[pimm.Command]: while not should_stop.value: From c43832d8d9a96de4799e5fca0690e396edfb7d99 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 21:58:10 +0200 Subject: [PATCH 58/64] Ask whether a single-user name earns its place, and stop steering the checkers `earn-its-place` now covers functions and states the single-user test across the shapes it applies to; `stranded-definition` cites it instead of carrying it, leaving existence to one rule and placement to the other. `check-rules` prompts are four fixed parts. Anything an author adds is drawn from knowing the change, which is the knowledge that hid the violation from them. --- .claude/skills/check-rules/SKILL.md | 8 ++++++++ CODE_RULES.md | 17 ++++++++++++----- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/.claude/skills/check-rules/SKILL.md b/.claude/skills/check-rules/SKILL.md index 4505a9992..be2b7bc9f 100644 --- a/.claude/skills/check-rules/SKILL.md +++ b/.claude/skills/check-rules/SKILL.md @@ -39,6 +39,14 @@ of one cleanly checked. Each needs its rule and a way to tell which code changed not know how that code came to be: a cold-headed reviewer answering one question, not a colleague who sat through the reasoning. +Every prompt is the same four parts: the rule pasted whole and verbatim, the scope Step 1 settled, the +return format below, and the instructions that follow it. Nothing else goes in. + +Whatever you would add comes from knowing the change, and that knowledge is what kept you from seeing +the violation yourself. Emphasis, an example, a "pay attention to", a scope trimmed to the lines you +touched — each hands the agent your reading of the rule in place of the rule, and the reading you did +not have is the one you spawned it for. + **Spawn them on Sonnet** — `model: 'sonnet'`, on every agent. Each one answers a narrow question: one rule against one diff. A fan-out of that width spends most of what it costs on being that wide rather than on the judging. The parameter is not optional: omitted, an agent inherits whatever model the diff --git a/CODE_RULES.md b/CODE_RULES.md index 88c5c0179..e7b2f7d11 100644 --- a/CODE_RULES.md +++ b/CODE_RULES.md @@ -190,8 +190,9 @@ inside the single entity that uses it. Distance costs the reader a search in bot use, to find out what it does; from the definition, to find out why it exists. A private name has every user in the file, so its place is determined. With a single user, first ask -whether the name is worth keeping: a body that says as much as its name is better inlined than moved. -Not touching `self` is no reason to stay at module level — that is what `@staticmethod` is for. +whether the name is worth keeping at all (`earn-its-place`); this rule only decides where a name that +is worth keeping goes. Not touching `self` is no reason to stay at module level — that is what +`@staticmethod` is for. A public name is looser: most of its users are elsewhere, and a module may order its surface deliberately. Group it with its in-file callers where that ordering does not say otherwise. Module @@ -285,14 +286,20 @@ except json.JSONDecodeError: ### earn-its-place -Don't add a class, file, field or parameter for a distinction the code already encodes. Check each of -the places one can already live — the dict an entry sits in, an enum member, the calling context, a -module that owns the subject — and use that instead. +Don't add a class, file, function, field or parameter for a distinction the code already encodes. Check +each of the places one can already live — the dict an entry sits in, an enum member, the calling +context, a module that owns the subject — and use that instead. The cost is not the lines. Every new type is another thing a reader holds, another place a value can live, and another edge to keep in sync; a field duplicating what its caller already knows goes stale the first time only one of the two is set. +With a single user, ask what the name adds over what already stands at that one use. A function whose +body is its name spelled out, a class whose fields the caller holds as locals anyway, a parameter only +ever passed one value — in each the use site already carries the distinction, and the name is a second +place to look. A name earns its place by holding what the use site cannot say: a diagnostic it raises, +an invariant it keeps, a domain predicate that takes more than one clause. + A new distinction lands beside its siblings, in the module that owns the subject. Put it there without asking. Raise it only when nothing suggests a home: that is a finding about the layout, and a human decides it. This rule does not apply to a repository still laying out its structure. From 7f620e2c809873ff0f6e1a80f0e07c0589f10496 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 23:13:02 +0200 Subject: [PATCH 59/64] Charge a call wall time or nothing, and read the loop as one step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `inference_latency` is a flag: a sim either charges each call the wall time it took or holds the world for it; a real rig always pays wall time. With the constant charge gone the release condition is "the call returned" in both modes, so `_Answer`, the `_Charge` classes and `_take` go, and the mode is one bool on `_InferenceWorker` — `throttle`'s timeout and `effect_time`'s wall term. `_step` is the algorithm in one place: throttle and reschedule on the call in flight; with none in flight, submit on a fresh observation and give it the rest of the round. Rollout-start bookkeeping moves to the first observation, the deadline drop into `_reschedule`, `_running` into `_worker`. The golden keeps its fixed 50 ms latency through a test-local `_SimulatedLatency` policy that holds each chunk on the world clock; the golden file is unchanged. --- ARCHITECTURE.md | 14 +- positronic/cli/eval/run.py | 5 +- positronic/policy/harness.py | 268 +++++++----------- .../policy/tests/test_golden_pipeline.py | 61 +++- positronic/policy/tests/test_harness.py | 117 +++----- 5 files changed, 206 insertions(+), 259 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index a20792fc5..a7cea9daa 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -132,13 +132,13 @@ episode lifecycle — nothing else. Scheduling, blending, history stacking and e the wrapper stack around the policy; a session returning `None` means "keep executing the current trajectory". -**Inference cost is a fact of the trial, owned by the harness.** The trial context carries -`inference_latency` — a constant charge in seconds, or the call's own wall duration — and the harness -alone reads it: it withholds a returned trajectory, and the world clock, until the charge is paid, -and the clock it hands the policy stack (`now`) reads the instant the in-flight call's output takes -effect. An answer with no waypoints to place — a skip, or the empty trajectory that stops what is -executing — lands at once. A scheduling wrapper stamps its chunk at `now()` and never learns the -mode, so the charge is deterministic wherever a constant one is asked for. +**Inference cost is a fact of the trial, owned by the harness.** A call costs the trial either the +wall time it took or nothing: the trial context's `inference_latency` flag asks a sim for the former, +and a real rig pays it regardless. Only the harness reads the flag. Paying nothing means holding the +world for the call, which holds a virtual clock still; paying wall time means letting the world run, +though no further ahead of the call's start than wall time has. The clock the harness hands the +policy stack (`now`) reads the instant the in-flight call's output takes effect, so a scheduling +wrapper stamps its chunk at `now()` and never learns the mode. **Recordings are canonical; codecs bind the dialect late.** The dataset records every run in the canonical conventions (frames, key names, absolute time) — never in a model's dialect. Every diff --git a/positronic/cli/eval/run.py b/positronic/cli/eval/run.py index 435314e92..3d6e22dca 100644 --- a/positronic/cli/eval/run.py +++ b/positronic/cli/eval/run.py @@ -194,7 +194,7 @@ def run( eval: Eval | str, policy, output_dir=None, - inference_latency=False, + inference_latency: bool = False, timing=False, policy_image: str | None = None, alias: str | None = None, @@ -222,7 +222,8 @@ def run( if not isinstance(eval, Eval): raise SystemExit(f'--eval={eval!r} is a name, not a config: pass --policy-image to run it on the platform') # The eval config owns the trial sweep (seed, task range); ``inference_latency`` is the CLI's per-run knob - # (sim inference-cost simulation). Overlay it onto every trial context, then self-drive the eval. + # (whether a sim charges each call its wall time). Overlay it onto every trial context, then self-drive + # the eval. eval = replace(eval, trials=[{**trial, keys.INFERENCE_LATENCY: inference_latency} for trial in eval.trials]) main(policy=policy, evals=[eval], output_dir=output_dir, timing=timing) return None diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index fe35a216a..f494f010c 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -1,9 +1,8 @@ import concurrent.futures import logging import time -from abc import ABC, abstractmethod from collections import deque -from collections.abc import Generator, Iterable, Iterator, Mapping +from collections.abc import Generator, Iterable, Iterator from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from enum import Enum @@ -19,7 +18,7 @@ from positronic.drivers import roboarm from positronic.drivers.roboarm.ik import assert_default_frame from positronic.eval import Embodiment, Observation, Task -from positronic.policy.base import Policy, Session +from positronic.policy.base import Policy 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 @@ -31,7 +30,7 @@ POLL_PERIOD_SEC = 0.01 # How long a submitted call may take and still resolve within its round: a wrapper that skips inference -# answers in microseconds, a real model call runs far past this and is paced across rounds by ``_take``. +# answers in microseconds, a real model call runs far past this and is throttled across rounds. SKIP_REPLY_SEC = 0.001 @@ -64,75 +63,6 @@ def ABORT(cls) -> 'Directive': return cls(DirectiveType.ABORT) -class _Answer(Enum): - """Whether the in-flight call's answer is the harness's to act on this round.""" - - READY = 'ready' - PENDING = 'pending' - - -class _Charge(ABC): - """What one session call costs the trial, in world-clock seconds.""" - - def __init__(self) -> None: - self._t0_ns = 0 - - def begin(self, t0_ns: int) -> None: - """Stamp the start of the call about to be submitted.""" - self._t0_ns = t0_ns - - @abstractmethod - def effect_time(self) -> float: - """The trial instant the in-flight call's output takes effect.""" - - @abstractmethod - def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: - """Block the loop for the call as far as this charge allows, then say where its answer stands.""" - - -class _DeclaredCharge(_Charge): - """A sim trial's ``inference_latency``: the call costs this, whatever the machine took.""" - - def __init__(self, seconds: float) -> None: - super().__init__() - self._ns = round(seconds * 1e9) - - def effect_time(self) -> float: - return (self._t0_ns + self._ns) / 1e9 - - def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: - # Blocking the loop thread holds a virtual clock still, which is what bills the trial the declared - # charge and not what the machine took. - concurrent.futures.wait([future]) - # An answer with nothing to place — a skip, or the empty trajectory that stops what is executing — - # has no effect instant to wait for. Integer ns, the world's own timeline: a float compare misses - # the instant by one ULP and slips the install a full round. - due = not future.result() or clock.now_ns() >= self._t0_ns + self._ns - return _Answer.READY if due else _Answer.PENDING - - -class _WallCharge(_Charge): - """A call costs the wall time it took: a real rig, and a sim asked for real latency.""" - - def __init__(self) -> None: - super().__init__() - self._wall_t0 = 0.0 - - def begin(self, t0_ns: int) -> None: - super().begin(t0_ns) - self._wall_t0 = time.monotonic() - - def effect_time(self) -> float: - return self._t0_ns / 1e9 + time.monotonic() - self._wall_t0 - - def wait(self, future: Future[list[dict[str, Any]] | None], clock: pimm.Clock) -> _Answer: - # The timeout is the sim's overrun over wall time, handed back: wall time cannot be held still, so - # the world runs no further ahead of the call's start than wall time has. An answer that is in has - # nothing left to wait for. - concurrent.futures.wait([future], timeout=max(clock.now() - self.effect_time(), 0.0)) - return _Answer.READY if future.done() else _Answer.PENDING - - def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: """Report the failure of a call nobody is waiting for any more.""" # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no @@ -141,16 +71,32 @@ def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: logging.error(f'Inference failed after the episode that asked for it ended: {exc}') +def _owned(obs: dict[str, Any]) -> dict[str, Any]: + """The observation with its arrays copied, so nothing rewrites what the worker is still reading. + + A producer may reuse one buffer for every sample it emits — a camera renders into the array behind + the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. + Copying at dispatch pays once per call rather than per round. + """ + return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} + + class _InferenceWorker: """One episode's policy session, called one at a time on a thread of its own so the harness keeps playing while the model runs. Ending an episode ``abandon``s the call in flight rather than waiting for - it, so the next episode never queues behind a model that hangs.""" + it, so the next episode never queues behind a model that hangs. + + ``charge_wall`` is what a call costs the trial: the wall time it took, or nothing — the loop is held + for the call, which holds a virtual clock still. + """ - def __init__(self, session: Session, charge: _Charge) -> None: - self._session = session - self._charge = charge + def __init__(self, policy: Policy, context: dict[str, Any], charge_wall: bool) -> None: + self._charge_wall = charge_wall + self._session = policy.new_session(context, self.effect_time) self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') - self._future: Future[list[dict[str, Any]] | None] | None = None + self._call: Future[list[dict[str, Any]] | None] | None = None + self._t0_ns = 0 # world clock at the in-flight call's submit + self._wall_t0 = 0.0 # ``time.monotonic()`` at that submit @property def meta(self) -> dict[str, Any]: @@ -158,33 +104,46 @@ def meta(self) -> dict[str, Any]: @property def idle(self) -> bool: - """Whether a call could be submitted: none is in flight.""" - return self._future is None - - def submit(self, obs: Mapping[str, Any], clock: pimm.Clock) -> None: - """Start a call on ``obs``, charged from this instant. The moment's wait lets a wrapper that skips - inference resolve in the round it was asked.""" - self._charge.begin(clock.now_ns()) - self._future = self._executor.submit(self._session, obs) - concurrent.futures.wait([self._future], timeout=SKIP_REPLY_SEC) - - def wait(self, clock: pimm.Clock) -> _Answer: - """Block for the in-flight call as far as its charge allows, then say where its answer stands.""" - assert self._future is not None - return self._charge.wait(self._future, clock) - - def collect(self) -> list[dict[str, Any]] | None: - """The ready call's trajectory — ``None`` when it had nothing to place — leaving the worker idle.""" - assert self._future is not None + return self._call is None + + @property + def done(self) -> bool: + """Whether the call in flight has returned.""" + return self._call is not None and self._call.done() + + def effect_time(self) -> float: + """The trial instant the in-flight call's output takes effect: its submit, plus its wall duration so + far when the trial pays wall time.""" + wall = time.monotonic() - self._wall_t0 if self._charge_wall else 0.0 + return self._t0_ns / 1e9 + wall + + def submit(self, obs: dict[str, Any], clock: pimm.Clock) -> None: + """Start a call on ``obs``. The moment's wait lets a wrapper that skips inference resolve in the round + it was asked.""" + self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() + self._call = self._executor.submit(self._session, frozen_view(_owned(obs))) + concurrent.futures.wait([self._call], timeout=SKIP_REPLY_SEC) + + def throttle(self, clock: pimm.Clock) -> None: + """Slow the loop for the call in flight as the trial's mode requires: until the call returns when the + world is held for it, else only while the world is ahead of the call's own wall clock.""" + assert self._call is not None + # Wall time cannot be held still, so the world runs no further ahead of the call's start than it has. + timeout = max(clock.now() - self.effect_time(), 0.0) if self._charge_wall else None + concurrent.futures.wait([self._call], timeout=timeout) + + def result(self) -> list[dict[str, Any]] | None: + """The returned call's trajectory — ``None`` when it had nothing to place — leaving the worker idle.""" + assert self._call is not None # Read on the loop thread, so a failing call still seals the episode. - actions, self._future = self._future.result(), None + actions, self._call = self._call.result(), None return actions def abandon(self) -> None: """Let go of the call in flight: its answer lands nowhere and its failure only reaches the log.""" - if self._future is not None: - self._future.add_done_callback(_report_abandoned) - self._future = None + if self._call is not None: + self._call.add_done_callback(_report_abandoned) + self._call = None self._executor.shutdown(wait=False, cancel_futures=True) def join(self) -> None: @@ -272,9 +231,9 @@ class Harness(pimm.ControlSystem): """Control system that runs the episode lifecycle and plays the policy's trajectory to the drivers. The wrapper owns the plan, the harness plays it, one command per channel per round. The session call - runs on a worker so playing continues while the model does, and its trajectory — with the world clock — - is withheld until the trial's ``inference_latency`` is paid; the ``now`` handed to ``new_session`` reads - time the same way, so wrappers stamp for the paid instant without knowing the mode. + runs on a worker so playing continues while the model does. A call costs the trial either the wall time + it took or nothing — the world held still for it — and the ``now`` handed to ``new_session`` reads the + instant the call's output takes effect, so wrappers stamp for it without knowing the mode. A ``trials`` plan makes the harness self-driving: it starts the next trial whenever idle and returns once the plan is exhausted. A task's ``timeout`` bounds every trial either way, and a truthy privileged @@ -298,9 +257,8 @@ def __init__( self.policy: Policy = policy self.context: dict[str, Any] = {} self._static_meta = static_meta or {} - # True between RUN and FINISH/ABORT: the trial is live — stepping and recording happen together. - self._running = False - # This episode's session and the thread it runs on; ``None`` while no episode is live. + # This episode's session and the thread it runs on. ``None`` while no episode is live: between RUN + # and FINISH/ABORT, stepping and recording happen together. self._worker: _InferenceWorker | None = None # A worker let go of mid-call, kept until the join that makes closing its session safe. self._retiring: _InferenceWorker | None = None @@ -420,9 +378,12 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: # rather than overhead the timing reducer attributes to this one. self._reap_worker() self.context = context - # A sim trial without the key runs free; a real rig ignores this sim-only knob and pays wall time. - latency = self.context.get(keys.INFERENCE_LATENCY, False) if self._embodiment.simulated else True - charge = _WallCharge() if latency is True else _DeclaredCharge(float(latency)) + latency = self.context.get(keys.INFERENCE_LATENCY, False) + if not isinstance(latency, bool): + raise ValueError(f'{keys.INFERENCE_LATENCY} is a flag, got {latency!r}') + # A real rig pays wall time whatever the trial asks: the knob is sim-only, and the eval CLI writes it + # into every trial. + charge_wall = latency or not self._embodiment.simulated self._awaiting_obs = set(self._embodiment.observations) self._rollout_started = False # Before the reset, so the reset and the rollout's other phase spans parent to the episode span. @@ -434,8 +395,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._worker = _InferenceWorker(self.policy.new_session(self.context, charge.effect_time), charge) - self._running = True + self._worker = _InferenceWorker(self.policy, self.context, charge_wall) self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -447,7 +407,7 @@ def _end_episode( The worker is retired rather than joined here, so a ``RemoteSession``'s websocket outlives the call still using it. """ - if self._running: + if self._worker is not None: if abort: self._cancel_session() # abort has no finalize to do it — stop the episode before the home self.ds_command.emit(DsWriterCommand.ABORT()) @@ -456,16 +416,15 @@ def _end_episode( else: yield from self._finalize_recording(clock, payload) self._home() - self._running = False def _handle_directive(self, directive: Directive, clock: pimm.Clock) -> Generator[pimm.Command, None, None]: - """Dispatch a directive to the episode lifecycle; updates ``_running``.""" + """Dispatch a directive to the episode lifecycle.""" match directive.type: case DirectiveType.RUN: - if not self._running: # a RUN mid-trial is ignored — the operator finishes before starting anew + if self._worker is None: # a RUN mid-trial is ignored — the operator finishes before starting anew self._begin_episode(directive.payload or {}, clock) case DirectiveType.FINISH: - if self._running: # a FINISH while idle is ignored — nothing to finalize + if self._worker is not None: # a FINISH while idle is ignored — nothing to finalize yield from self._end_episode(clock, directive.payload) case DirectiveType.ABORT: yield from self._end_episode(clock, abort=True) @@ -531,26 +490,6 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() inputs['descriptor'] = self._embodiment.descriptor - return inputs - - @staticmethod - def _owned(obs: dict[str, Any]) -> dict[str, Any]: - """The observation with its arrays copied, so nothing rewrites what the worker is still reading. - - A producer may reuse one buffer for every sample it emits — a camera renders into the array behind - the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. - Copying at dispatch pays once per call rather than per round. - """ - return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} - - def _step(self, clock: pimm.Clock) -> None: - """Keep one session call in flight and install the trajectory it returns.""" - assert self._worker is not None # only a live episode steps - if not self._worker.idle and self._take(clock) is _Answer.PENDING: - return - obs = self._build_obs(clock) - if obs is None: - return if not self._rollout_started: # The rollout begins at its first observation, not when the reset returned: the turns spent # delivering the scene are neither the trial's budget nor its duration. @@ -558,44 +497,53 @@ def _step(self, clock: pimm.Clock) -> None: self._telemetry.start_rollout(clock.now()) if self._task is not None: self._deadline = clock.now() + self._task.timeout - self._worker.submit(frozen_view(self._owned(obs)), clock) - self._take(clock) - - def _take(self, clock: pimm.Clock) -> _Answer: - """Install the in-flight call's trajectory once the world has paid its charge.""" - assert self._worker is not None - if self._worker.wait(clock) is _Answer.PENDING: - return _Answer.PENDING # the schedule already playing carries the world to the effect instant - actions = self._worker.collect() - # The world reached the deadline while the call 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. - expired = self._deadline is not None and clock.now() >= self._deadline - if actions is not None and not expired: - self._install(actions, clock) - return _Answer.READY + return inputs + + def _step(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: + """One round of inference: throttle the loop for the call in flight as the trial's mode requires and + reschedule on what it returns; with no call in flight, submit one on a fresh observation and give it + the rest of the round to return.""" + self._throttle_and_reschedule(worker, clock) + if not worker.idle: # the schedule already playing carries the world on + return + obs = self._build_obs(clock) + if obs is not None: + worker.submit(obs, clock) + self._throttle_and_reschedule(worker, clock) + + def _throttle_and_reschedule(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: + if worker.idle: + return + worker.throttle(clock) + if worker.done and (trajectory := worker.result()) is not None: + self._reschedule(trajectory, clock) @staticmethod - def _assert_anchored(actions: list[dict[str, Any]], now: float) -> None: + 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 actions), default=0.0) + 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 _install(self, actions: list[dict[str, Any]], clock: pimm.Clock) -> None: + 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 wrapper for the instant its charge is paid. + by the scheduling wrapper for the instant the call's output takes effect. """ - self._assert_anchored(actions, clock.now()) + if self._deadline is not None and clock.now() >= self._deadline: + # The world reached the deadline while the call 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()) self._telemetry.step() # The single explicit seconds->ns seam: wrappers 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 actions if name in a) + 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. @@ -660,7 +608,7 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ assert directive_msg is not None and manual_msg is not None if directive_msg.updated: yield from self._handle_directive(directive_msg.data, clock) - elif not self._running: + elif self._worker is None: if manual_msg.updated and manual_msg.data is not None: self._emit(manual_msg.data) elif self._trials is not None: @@ -673,11 +621,11 @@ def _run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[ yield from self._end_episode(clock, terminal) else: try: - self._step(clock) + self._step(self._worker, clock) except pimm.NoValueException: pass self._play(clock) yield self._pace(clock) - if self._running: + if self._worker is not None: yield from self._finalize_recording(clock) diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index a4f6bccee..5519d3d5a 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -13,8 +13,8 @@ * horizon/gating regression -> the state trajectory diverges Everything runs on CPU in a virtual-time world only: no GL/GPU/MuJoCo, no wall -clock in the asserted path (a fixed-float ``inference_latency`` is used so -inference latency is deterministic). +clock in the asserted path (``_SimulatedLatency`` charges a fixed world-time +latency, so inference latency is deterministic). Regenerate the golden after an intentional behavior change: @@ -41,7 +41,7 @@ from positronic.drivers.roboarm.command import CartesianPosition, CommandType, Reset from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation from positronic.geom import Rotation, Transform3D -from positronic.policy.base import Policy, Session +from positronic.policy.base import DelegatingPolicy, DelegatingSession, Now, Policy, Session from positronic.policy.codec import ActionTiming from positronic.policy.harness import Directive, Harness from positronic.policy.wrappers import ChunkedSchedule, StopOnFault @@ -53,7 +53,7 @@ INITIAL_Q = np.array([0.10, -0.20, 0.30, -0.40, 0.50, -0.60, 0.70], dtype=np.float32) TARGET_POS = np.array([0.50, 0.00, 0.45], dtype=np.float32) -# Fixed deterministic inference latency. Spans >1 control tick (harness loop is +# Fixed deterministic inference latency in world time. Spans >1 control tick (harness loop is # 0.01 s) so the post-inference anchoring effect is observable in recorded ts. INFERENCE_LATENCY_S = 0.05 ACTION_FPS = 15.0 @@ -87,6 +87,49 @@ def new_session(self, context=None, now=None): return _ScriptedSession() +class _SimulatedLatency(DelegatingPolicy): + """A fixed inference latency in world time: every chunk is stamped for, and reaches the harness at, + ``latency_sec`` after the call that produced it, whatever the machine took. Sits outermost, so nothing + below sees an observation while a chunk is held. A skip or a stop passes through at once.""" + + def __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, now: Now, latency_ns: int): + super().__init__(inner) + self._now = now + self._latency_ns = latency_ns + self._held: list | None = None + self._release_at_ns = 0 + + def __call__(self, obs): + # Integer ns, the world's own timeline: a float compare can miss the release instant by one ULP. + now_ns = round(self._now() * 1e9) + if self._held is not None: + if now_ns < self._release_at_ns: + return None + held, self._held = self._held, None + return held + result = self._inner(obs) + if not result: + return result + self._held, self._release_at_ns = result, now_ns + self._latency_ns + return None + + def cancel(self): + self._held = None + super().cancel() + + def new_session(self, context=None, now=None): + assert now is not None, 'the harness supplies the clock' + latency_ns = self._latency_ns + return _SimulatedLatency._Session( + self._inner.new_session(context, lambda: now() + latency_ns / 1e9), now, latency_ns + ) + + class _FakeRobotState(roboarm.State): """Lossless re-expression of the last applied command over sim-time.""" @@ -194,11 +237,13 @@ def _run_pipeline(tmp_path: Path) -> dict: }, static_meta=dict(ROBOT_STATIC_META), meta_source=robot.robot_meta, - # ``inference_latency`` is a sim-only knob, and the fake robot's control-period sleep is this - # world's sole time-master — the shape a sim eval runs in. + # The fake robot's control-period sleep is this world's sole time-master — the shape a sim eval + # runs in. simulated=True, ) - harness = Harness((StopOnFault() | ChunkedSchedule()).wrap(policy), embodiment) + harness = Harness( + _SimulatedLatency((StopOnFault() | ChunkedSchedule()).wrap(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) directive_em = world.pair(harness.directive) @@ -206,7 +251,7 @@ def _run_pipeline(tmp_path: Path) -> dict: # Robot/gripper emit state every tick, so the script only drives the # episode lifecycle and the one-shot error injection. script = [ - (partial(directive_em.emit, Directive.RUN(task='golden', inference_latency=INFERENCE_LATENCY_S)), 0.0), + (partial(directive_em.emit, Directive.RUN(task='golden')), 0.0), (None, 1.5), # several reactive inference + chunk/horizon cycles (robot.inject_error, 0.0), # one-shot error: that frame is dropped, then inference resumes (None, 0.5), diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 90fcfb80e..228de76ed 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -21,7 +21,7 @@ from positronic.offboard.client import InferenceSession from positronic.policy.base import DelegatingSession, Policy, PolicyWrapper, Session from positronic.policy.codec import ActionTimestamp -from positronic.policy.harness import Directive, DirectiveType, Harness, _InferenceWorker, _WallCharge +from positronic.policy.harness import Directive, DirectiveType, Harness, _InferenceWorker from positronic.policy.remote import RemoteSession from positronic.policy.wrappers import ChunkedSchedule, StopOnFault from positronic.tests.testing_coutils import ManualDriver, RecordingEmitter, drive_scheduler @@ -932,14 +932,13 @@ def test_trial_plan_self_drives(world): @pytest.mark.timeout(3.0) def test_timeout_during_inference_drops_the_chunk(world): - """A trial whose deadline lapses while the model is still owed its latency ends with the call in flight: + """A trial whose deadline lapses while the model is still inside its call ends with the call in flight: the trajectory it eventually returns is discarded, never emitted past the advertised termination point.""" - policy = StubPolicy() harness = Harness( - ChunkedSchedule().wrap(policy), + ChunkedSchedule().wrap(SlowPolicy(wall_sec=0.3)), # the call runs well past the deadline make_embodiment(simulated=True), task=Task(instruction='test', timeout=0.05), - trials=[{keys.INFERENCE_LATENCY: 0.2}], # the charge holds the answer well past the deadline + trials=[{keys.INFERENCE_LATENCY: True}], ) cmd_recorder = RecordingEmitter() grip_recorder = RecordingEmitter() @@ -1657,20 +1656,20 @@ def test_anchored_chunk_passes(): Harness._assert_anchored([{'timestamp': 1.7e9 - 0.2}, {'timestamp': 1.7e9 + 1.5}], now=1.7e9) -@pytest.mark.parametrize(('expired', 'installed'), [(True, False), (False, True)]) -def test_a_reply_is_installed_only_while_the_trial_still_has_budget(world, expired, installed): +@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 call whose rounds in flight carried the world past that instant has its chunk dropped instead of placed, and ``_run`` finishes the trial on the next round.""" policy = StubPolicy() harness = Harness(policy, make_embodiment()) now = world.clock.now() harness._deadline = now - 1.0 if expired else now + 1.0 - harness._worker = _InferenceWorker(policy.new_session(), _WallCharge()) + harness._worker = _InferenceWorker(policy, {}, charge_wall=True) harness._worker.submit({}, world.clock) - harness._take(world.clock) + harness._throttle_and_reschedule(harness._worker, world.clock) - assert bool(harness._schedules[keys.ROBOT_COMMAND]) is installed + assert bool(harness._schedules[keys.ROBOT_COMMAND]) is scheduled class _SlowSession(Session): @@ -1782,17 +1781,6 @@ def test_default_latency_pauses_the_world_for_the_call(world): assert played[0][0] < 0.01, f'the world advanced during the call: first command at {played[0][0]}s' -@pytest.mark.timeout(20.0) -@pytest.mark.parametrize('wall_sec', [0.0, 0.05]) -def test_constant_latency_ignores_what_the_call_really_took(world, wall_sec): - """The reproducible mode: the wrapper is released a constant delay after the call started, so the played - trace is the same against a fast server and a slow one.""" - played = _run_episode(world, SlowPolicy(wall_sec=wall_sec), ChunkedSchedule(), latency=0.3) - - assert played, 'no command was played' - assert played[0][0] == pytest.approx(0.3, abs=0.02), f'first command at {played[0][0]}s, expected the 0.3s delay' - - @pytest.mark.timeout(20.0) def test_measured_latency_charges_the_calls_own_wall_duration(world): """``inference_latency=True`` charges the world what the model really took, so a slow server is scored @@ -1804,33 +1792,13 @@ def test_measured_latency_charges_the_calls_own_wall_duration(world): @pytest.mark.timeout(20.0) -def test_a_real_rig_ignores_the_latency_a_trial_asks_for(world): - """The charge simulates a trial; a real rig pays what its calls take, so a context carrying a constant - (the eval CLI writes one into every trial) does not hold the chunk back for it.""" - played = _run_episode(world, SlowPolicy(), ChunkedSchedule(), latency=5.0, simulated=False) +def test_a_real_rig_pays_wall_time_whatever_the_trial_asks_for(world): + """The knob simulates a trial; a real rig pays what its calls take, so a context asking for the world to + be held (the eval CLI writes the flag into every trial) does not hold it.""" + played = _run_episode(world, SlowPolicy(wall_sec=0.2), ChunkedSchedule(), latency=False, simulated=False) assert played, 'no command was played' - assert played[0][0] < 1.0, f"first command at {played[0][0]}s: the trial's 5s charge was honoured" - - -class _WallCost(PolicyWrapper): - """Burns real time on every call it passes down, answering or skipping — a stand-in for a stack that - copies buffers before the model ever sees the observation.""" - - def __init__(self, sec: float): - self._sec = sec - - class _Session(DelegatingSession): - def __init__(self, inner: Session, sec: float): - super().__init__(inner) - self._sec = sec - - def __call__(self, obs): - time.sleep(self._sec) - return self._inner(obs) - - def wrap_session(self, inner: Session, context, now): - return _WallCost._Session(inner, self._sec) + assert played[0][0] >= 0.2, f'first command at {played[0][0]}s, under the 0.2s the call took' class _ObservedTicks(PolicyWrapper): @@ -1852,43 +1820,29 @@ def wrap_session(self, inner: Session, context, now): return _ObservedTicks._Session(inner, self.seen) -@pytest.mark.timeout(60.0) -def test_a_constant_charge_keeps_the_trace_off_the_machine_clock(): - """The reproducible mode's promise: what the machine spends — inside the model call, or above it in a - wrapper that ends up skipping — must not reach the trial. Two worlds, one slow, one not.""" - - def played(wall_sec: float) -> list[tuple[float, Any]]: - with pimm.World(virtual_time=True) as w: - policy = SlowPolicy(wall_sec=wall_sec, span_sec=0.3, steps=15) - return _run_episode(w, policy, _WallCost(wall_sec) | ChunkedSchedule(), latency=0.2, run_sec=0.8) - - assert played(0.0) == played(0.003) - - @pytest.mark.timeout(30.0) -def test_the_wrappers_see_every_tick_outside_the_charge(world): - """What a temporal stack records: the charge is the only thing that keeps an observation from the - wrappers above the scheduler — never the machine, and never the rounds in between.""" +def test_the_wrappers_see_every_tick(world): + """What a temporal stack records: nothing keeps an observation from the wrappers above the scheduler — + not the machine inside the model call, and not the rounds in between.""" ticks = _ObservedTicks() _run_episode( - world, SlowPolicy(wall_sec=0.01, span_sec=0.3, steps=15), ticks | ChunkedSchedule(), latency=0.2, run_sec=1.0 + world, SlowPolicy(wall_sec=0.01, span_sec=0.3, steps=15), ticks | ChunkedSchedule(), latency=False, run_sec=1.0 ) - period, charge = 0.005, 0.2 # the pacer's control period, and the trial's charge + period = 0.005 # the pacer's control period gaps = [round(b - a, 4) for a, b in zip(ticks.seen, ticks.seen[1:], strict=False)] - assert all(gap <= period or charge <= gap <= charge + 2 * period for gap in gaps), gaps - assert len([gap for gap in gaps if gap > period]) >= 2, f'the trial inferred too few times: {gaps}' + assert gaps and all(gap <= period for gap in gaps), gaps @pytest.mark.timeout(20.0) def test_harness_keeps_playing_while_a_call_is_in_flight(world): """A wrapper that replans before its chunk is exhausted leaves waypoints due during inference, and the harness emits them on time instead of standing still until the model answers.""" - played = _run_episode(world, SlowPolicy(span_sec=0.4, steps=20), _ReplanEarly(), latency=0.15) + played = _run_episode(world, SlowPolicy(wall_sec=0.15, span_sec=0.4, steps=20), _ReplanEarly(), latency=True) - # The second call starts halfway through the first chunk (0.2s in) and is owed 0.15s; the waypoints due - # in that window have to keep going out. - during = [t for t, _ in played if 0.2 <= t < 0.35] + # The first chunk lands at ~0.15s and spans 0.4s; the second call starts halfway through it (~0.28s) and + # runs for 0.15s the world runs through; the waypoints due in that window have to keep going out. + during = [t for t, _ in played if 0.29 <= t < 0.4] assert len(during) >= 3, f'the harness stopped playing during inference: {[t for t, _ in played]}' @@ -1994,10 +1948,9 @@ def test_the_trial_context_cannot_stand_in_for_what_the_harness_read(world): @pytest.mark.timeout(20.0) -def test_a_stop_lands_without_waiting_out_the_charge(world): - """A charge places a trajectory's waypoints, and a stop has none: an arm that faults mid-chunk stops in - the round its fault is seen, not a charge later.""" - charge, fault_at, period = 0.3, 0.5, 0.005 +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() harness = Harness(stack.wrap(SlowPolicy(span_sec=1.0, steps=50)), make_embodiment(simulated=True)) grip_recorder = _TimedRecorder(world.clock) @@ -2012,9 +1965,9 @@ def test_a_stop_lands_without_waiting_out_the_charge(world): pose, joints = [0.1, 0.2, 0.3], [0.4, 0.5, 0.6] driver = ManualDriver([ - (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: charge})), 0.0), + (partial(directive_em.emit, Directive.RUN(task='t')), 0.0), (partial(emit_ready_payload, frame_em, robot_em, grip_em, make_robot_state(pose, joints)), fault_at), - (partial(robot_em.emit, make_robot_state(pose, joints, status=RobotStatus.ERROR)), charge + 0.2), + (partial(robot_em.emit, make_robot_state(pose, joints, status=RobotStatus.ERROR)), 0.5), ]) drive_scheduler(world.start([harness, driver, _Pacer(period)]), steps=4000) @@ -2160,7 +2113,7 @@ def new_session(self, context=None, now=None): @pytest.mark.timeout(3.0) -def test_installed_trajectory_clears_the_channels_it_omits(world): +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.""" @@ -2231,9 +2184,9 @@ def test_home_and_manual_commands_are_emitted_as_plain_values(world): @pytest.mark.timeout(20.0) def test_abort_discards_a_call_that_is_still_in_flight(world): - """An ABORT while the gate is still holding the model's answer throws that answer away: the trajectory - it carries never reaches the devices.""" - harness = Harness(ChunkedSchedule().wrap(SlowPolicy()), make_embodiment(simulated=True)) + """An ABORT while the model is still inside its call throws that answer away: the trajectory it carries + never reaches the devices.""" + harness = Harness(ChunkedSchedule().wrap(SlowPolicy(wall_sec=1.0)), make_embodiment(simulated=True)) cmd_recorder = RecordingEmitter() harness.commands[keys.ROBOT_COMMAND]._bind(cmd_recorder) harness.commands[keys.TARGET_GRIP]._bind(RecordingEmitter()) @@ -2246,9 +2199,9 @@ def test_abort_discards_a_call_that_is_still_in_flight(world): robot_state = make_robot_state([0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) driver = ManualDriver([ - (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: 1.0})), 0.0), + (partial(directive_em.emit, Directive.RUN(task='t', **{keys.INFERENCE_LATENCY: True})), 0.0), (partial(emit_ready_payload, frame_em, robot_em, grip_em, robot_state), 0.001), - (None, 0.05), # well inside the 1.0s the gate owes the call + (None, 0.05), # well inside the 1.0s the call takes (partial(directive_em.emit, Directive.ABORT()), 0.0), (None, 0.05), ]) From 8eb0703ad9ea3206cfec0251e74c15785a044b4c Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 23:22:03 +0200 Subject: [PATCH 60/64] Guard the idle worker in `_step`, not in the helper both sites share --- positronic/policy/harness.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index f494f010c..b1bfd581e 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -503,17 +503,15 @@ def _step(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: """One round of inference: throttle the loop for the call in flight as the trial's mode requires and reschedule on what it returns; with no call in flight, submit one on a fresh observation and give it the rest of the round to return.""" - self._throttle_and_reschedule(worker, clock) - if not worker.idle: # the schedule already playing carries the world on - return - obs = self._build_obs(clock) - if obs is not None: - worker.submit(obs, clock) + if not worker.idle: self._throttle_and_reschedule(worker, clock) + if worker.idle: + obs = self._build_obs(clock) + if obs is not None: + worker.submit(obs, clock) + self._throttle_and_reschedule(worker, clock) def _throttle_and_reschedule(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: - if worker.idle: - return worker.throttle(clock) if worker.done and (trajectory := worker.result()) is not None: self._reschedule(trajectory, clock) From 588a459933f1743f9f0fe0608f6fe583f83fb131 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 23:25:59 +0200 Subject: [PATCH 61/64] Fold `_read_channel` back into `_build_obs` --- positronic/policy/harness.py | 45 +++++++++++++++--------------------- 1 file changed, 18 insertions(+), 27 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index b1bfd581e..659c7af11 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -17,7 +17,7 @@ from positronic.dataset.serializers import expand_suffixed from positronic.drivers import roboarm from positronic.drivers.roboarm.ik import assert_default_frame -from positronic.eval import Embodiment, Observation, Task +from positronic.eval import Embodiment, Task from positronic.policy.base import Policy from positronic.utils import flatten_dict, frozen_view @@ -436,28 +436,6 @@ def _is_faulted(value: Any) -> bool: """Whether a raw observation is an arm reporting a fault. Every other not-ready sample is simply absent.""" return isinstance(value, roboarm.State) and value.status is roboarm.RobotStatus.ERROR - def _read_channel(self, name: str, obs: Observation) -> tuple[dict[str, Any] | None, bool]: - """This channel's entries under their full names, and whether the arm behind it is faulted. - - The entries are ``None`` when the channel has no sample to give — a resetting or faulted arm alike. - Raises ``NoValueException`` before the channel has produced anything at all. - """ - message = self.observations[name].read() - if message is None: - raise pimm.NoValueException - if message.updated: - self._awaiting_obs.discard(name) - value = message.data - if obs.serializer is not None: - value = obs.serializer(value) - if value is None: - # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so the - # fault is recovered from the raw sample and stapled on by the caller as `keys.ROBOT_FAULT` — a - # name already claiming to be part of `robot_state`. Emit it from the serializer and this - # branch, the raw-type check and the flag all go, and the fault reaches the recording as well. - return None, self._is_faulted(message.data) - return {full: v for full, v in expand_suffixed(name, value) if v is not None}, False - def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: """Read every observation channel and assemble the policy input dict. @@ -474,10 +452,23 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: faulted = False not_ready = False for name, obs in self._embodiment.observations.items(): - entries, channel_faulted = self._read_channel(name, obs) - faulted = faulted or channel_faulted - not_ready = not_ready or entries is None - inputs.update(entries or {}) + message = self.observations[name].read() + if message is None: + raise pimm.NoValueException + if message.updated: + self._awaiting_obs.discard(name) + value = message.data + if obs.serializer is not None: + value = obs.serializer(value) + if value is None: # no sample to give: a resetting or faulted arm alike + # HACK(#619): a serializer answers `None` for a resetting arm and a faulted one alike, so + # the fault is recovered from the raw sample and stapled on as `keys.ROBOT_FAULT` — a name + # already claiming to be part of `robot_state`. Emit it from the serializer and this + # branch, the raw-type check and the flag all go, and the fault reaches the recording too. + faulted = faulted or self._is_faulted(message.data) + not_ready = True + continue + inputs.update({full: v for full, v in expand_suffixed(name, value) if v is not None}) # Every channel is read before this decision, so a bimanual rig cannot hide one arm's fault behind # another arm's not-ready sample: whichever channel comes first, the fault still reaches the stack. if (not_ready and not faulted) or self._awaiting_obs: From 2036bc0efc064881cf2cdbf2e561554e081c40a3 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Sun, 16 Aug 2026 23:47:45 +0200 Subject: [PATCH 62/64] Say what `_awaiting_obs` guards against, and why by `updated` --- positronic/policy/harness.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index 659c7af11..f7749f59e 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -269,8 +269,10 @@ def __init__( self._rollout_started = False # Wall-clock telemetry for the live rollout, opened under ``--timing`` and inert otherwise. self._telemetry = _EpisodeTelemetry() - # Channels that have not delivered since this episode's reset. A receiver latches its last value, so - # emptying this set is what keeps the first inference off the previous episode's final frame. + # Channels that have not delivered since this episode's RUN; the first inference waits until every one + # has. A receiver latches its last value, so a producer silent between episodes — a remote env emits + # only when stepped — would otherwise feed the previous episode's final frame. Delivery is judged by + # ``updated``, not ``ts``: some producers stamp ``ts`` on their own clock. self._awaiting_obs: set[str] = set() self.observations = pimm.ReceiverDict(self) From 7620f1a5b387cbb5ab9667a2dc98055e5f2c144b Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Mon, 17 Aug 2026 00:18:14 +0200 Subject: [PATCH 63/64] Anchor the worker's clock at construction, and keep each name beside its user `_InferenceWorker` takes the clock and stamps `_t0_ns`/`_wall_t0` before `new_session`, so the `now` it hands the session reads a trial instant from the moment the session exists rather than raising on an unset field. `submit` and `throttle` read the held clock instead of taking one per call. `_owned` and `_report_abandoned` become static methods above their sole callers; `_append` inlines into `_record`. `keys.DESCRIPTOR` replaces the `'descriptor'` literal at the harness and its readers. --- .claude/skills/address-review/SKILL.md | 9 ++- docs/connect-your-model.md | 2 +- positronic/dataset/ds_writer_agent.py | 11 +-- positronic/keys.py | 8 +- positronic/policy/harness.py | 75 +++++++++---------- positronic/policy/tests/test_harness.py | 10 +-- .../tests/test_inference_integration.py | 2 +- 7 files changed, 57 insertions(+), 60 deletions(-) diff --git a/.claude/skills/address-review/SKILL.md b/.claude/skills/address-review/SKILL.md index f5144ad7c..1d28d359a 100644 --- a/.claude/skills/address-review/SKILL.md +++ b/.claude/skills/address-review/SKILL.md @@ -130,10 +130,11 @@ genuinely the user's to accept. The user may drop or override any item. **Whether a finding belongs in this PR is your call.** "Fix it here or track it for later" is the verdict this skill exists to reach, and by the time you can pose it the research is done, -so the human has nothing to add that you don't already have. Reach it, file the issue, reply, -and put it in the Step 6 report where they can overrule it. Bring a deferral back only when -leaving it undone changes what the PR ships — a gap in the behaviour the user asked for, not a -follow-up carrying an issue number. +so the human has nothing to add that you don't already have. Reach it, reply, and propose the +tracking issue in the Step 6 report, where they can approve or overrule it — filing it is +theirs to authorize, never yours. Bring a deferral back for a decision only when leaving it +undone changes what the PR ships — a gap in the behaviour the user asked for, not a follow-up +the report already names. When you do bring something back, explain it in full and in plain words: what the thing is and how it works, what the reviewer asked for, what you did or did not do and why, and what each diff --git a/docs/connect-your-model.md b/docs/connect-your-model.md index 81e155f8b..122cad7f1 100644 --- a/docs/connect-your-model.md +++ b/docs/connect-your-model.md @@ -109,7 +109,7 @@ The client sends the full raw robot state as a dict. Keys are flat strings (the | `task` | str | — | Language instruction for the episode | | `descriptor` | str | — | Embodiment the observation came from (e.g. `mujoco.franka`); empty string when unset. Lets a multi-embodiment policy adapt to the current robot | -Your server receives every key each step, except that an arm with no sound sample to give — faulted, or mid-reset — omits its `robot_state.*` measurements and sends only `robot_state.fault`. The standard stack puts `StopOnFault` ahead of the model, which answers such a step itself and never forwards it; a stack without it reaches the model with those keys missing. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. +Your server receives every key each step, except that a faulted arm omits its `robot_state.*` measurements and sends only `robot_state.fault`. The standard stack puts `StopOnFault` ahead of the model, which answers such a step itself and never forwards it; a stack without it reaches the model with those keys missing. A mid-reset arm is not that case: it produces no step at all, so the client skips the round rather than sending a partial observation. Use what your model needs and ignore the rest. Image stream names are configuration-driven, so key off the names your deployment uses rather than assuming fixed ones. The table above is a single-arm rig; a multi-arm one names its state and grip channels per arm. ### Actions (server → client) diff --git a/positronic/dataset/ds_writer_agent.py b/positronic/dataset/ds_writer_agent.py index 8e2079e16..f5d9b286f 100644 --- a/positronic/dataset/ds_writer_agent.py +++ b/positronic/dataset/ds_writer_agent.py @@ -59,13 +59,6 @@ def ABORT(): return DsWriterCommand(DsWriterCommandType.ABORT_EPISODE) -def _append(ep_writer: EpisodeWriter, name: str, value: Any, ts_ns: int, extra_ts: dict[str, int]): - for full_name, v in expand_suffixed(name, value): - if v is None: - continue - ep_writer.append(full_name, v, ts_ns, extra_ts) - - class TimeMode(IntEnum): """Mode of timestamping for the dataset writer.""" @@ -146,7 +139,9 @@ def _record(self, ep_writer: EpisodeWriter, name: str, msg: pimm.Message, clock: value = msg.data if serializer is not None: value = serializer(value) - _append(ep_writer, name, value, primary_ts, extra_ts) + for full_name, v in expand_suffixed(name, value): + if v is not None: + ep_writer.append(full_name, v, primary_ts, extra_ts) def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): """Main loop: process commands and append updated inputs to the episode.""" diff --git a/positronic/keys.py b/positronic/keys.py index a9149e59b..587a07dd6 100644 --- a/positronic/keys.py +++ b/positronic/keys.py @@ -38,6 +38,8 @@ def is_robot_command(name: str) -> bool: ROBOT_FAULT = f'{ROBOT_STATE}.fault' GRIP = 'grip' TASK = 'task' +# The embodiment an observation came from, so a multi-embodiment policy can tell which robot it is driving. +DESCRIPTOR = 'descriptor' # The prefix that identifies a camera on the wire: an embodiment declares its cameras by naming # them this way, and every consumer picks them out of the observations by it. IMAGE_PREFIX = 'image.' @@ -96,7 +98,7 @@ def is_robot_command(name: str) -> bool: EVAL_SUCCESS = 'eval.success' EVAL_TERMINATED = 'eval.terminated' -# The inference cost each model call charges the world clock: a constant number of seconds (reproducible), -# or ``True`` for the call's own wall duration. A sim trial without the key charges nothing (the world holds -# still per call); hardware always pays wall. +# Whether a model call charges the world clock its own wall duration. A flag: any other type is rejected +# when the episode starts. A sim trial without it charges nothing (the world holds still per call); +# hardware always pays wall whatever the trial asks. INFERENCE_LATENCY = 'inference_latency' diff --git a/positronic/policy/harness.py b/positronic/policy/harness.py index f7749f59e..fd50adf49 100644 --- a/positronic/policy/harness.py +++ b/positronic/policy/harness.py @@ -63,40 +63,23 @@ def ABORT(cls) -> 'Directive': return cls(DirectiveType.ABORT) -def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: - """Report the failure of a call nobody is waiting for any more.""" - # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no - # caller left to raise to; the log is the only place its failure can go. - if not future.cancelled() and (exc := future.exception()) is not None: - logging.error(f'Inference failed after the episode that asked for it ended: {exc}') - - -def _owned(obs: dict[str, Any]) -> dict[str, Any]: - """The observation with its arrays copied, so nothing rewrites what the worker is still reading. - - A producer may reuse one buffer for every sample it emits — a camera renders into the array behind - the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. - Copying at dispatch pays once per call rather than per round. - """ - return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} - - class _InferenceWorker: """One episode's policy session, called one at a time on a thread of its own so the harness keeps - playing while the model runs. Ending an episode ``abandon``s the call in flight rather than waiting for - it, so the next episode never queues behind a model that hangs. + playing while the model runs. ``charge_wall`` is what a call costs the trial: the wall time it took, or nothing — the loop is held for the call, which holds a virtual clock still. """ - def __init__(self, policy: Policy, context: dict[str, Any], charge_wall: bool) -> None: + def __init__(self, policy: Policy, context: dict[str, Any], charge_wall: bool, clock: pimm.Clock) -> None: self._charge_wall = charge_wall + self._clock = clock + # World clock and ``time.monotonic()`` at the in-flight call's submit, anchored at the episode's + # start so ``effect_time`` reads a trial instant from the moment the session exists. + self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() self._session = policy.new_session(context, self.effect_time) self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix='harness-session') self._call: Future[list[dict[str, Any]] | None] | None = None - self._t0_ns = 0 # world clock at the in-flight call's submit - self._wall_t0 = 0.0 # ``time.monotonic()`` at that submit @property def meta(self) -> dict[str, Any]: @@ -117,19 +100,29 @@ def effect_time(self) -> float: wall = time.monotonic() - self._wall_t0 if self._charge_wall else 0.0 return self._t0_ns / 1e9 + wall - def submit(self, obs: dict[str, Any], clock: pimm.Clock) -> None: + @staticmethod + def _owned(obs: dict[str, Any]) -> dict[str, Any]: + """The observation with its arrays copied, so nothing rewrites what the worker is still reading. + + A producer may reuse one buffer for every sample it emits — a camera renders into the array behind + the adapter it re-emits each frame — and it keeps stepping while a call charged in wall time runs. + Copying at dispatch pays once per call rather than per round. + """ + return {name: value.copy() if isinstance(value, np.ndarray) else value for name, value in obs.items()} + + def submit(self, obs: dict[str, Any]) -> None: """Start a call on ``obs``. The moment's wait lets a wrapper that skips inference resolve in the round it was asked.""" - self._t0_ns, self._wall_t0 = clock.now_ns(), time.monotonic() - self._call = self._executor.submit(self._session, frozen_view(_owned(obs))) + self._t0_ns, self._wall_t0 = self._clock.now_ns(), time.monotonic() + self._call = self._executor.submit(self._session, frozen_view(self._owned(obs))) concurrent.futures.wait([self._call], timeout=SKIP_REPLY_SEC) - def throttle(self, clock: pimm.Clock) -> None: + def throttle(self) -> None: """Slow the loop for the call in flight as the trial's mode requires: until the call returns when the world is held for it, else only while the world is ahead of the call's own wall clock.""" assert self._call is not None # Wall time cannot be held still, so the world runs no further ahead of the call's start than it has. - timeout = max(clock.now() - self.effect_time(), 0.0) if self._charge_wall else None + timeout = max(self._clock.now() - self.effect_time(), 0.0) if self._charge_wall else None concurrent.futures.wait([self._call], timeout=timeout) def result(self) -> list[dict[str, Any]] | None: @@ -139,10 +132,18 @@ def result(self) -> list[dict[str, Any]] | None: actions, self._call = self._call.result(), None return actions + @staticmethod + def _report_abandoned(future: Future[list[dict[str, Any]] | None]) -> None: + """Report the failure of a call nobody is waiting for any more.""" + # rules-allow: swallowed-error — the call outlived the episode that asked for it, so there is no + # caller left to raise to; the log is the only place its failure can go. + if not future.cancelled() and (exc := future.exception()) is not None: + logging.error(f'Inference failed after the episode that asked for it ended: {exc}') + def abandon(self) -> None: """Let go of the call in flight: its answer lands nowhere and its failure only reaches the log.""" if self._call is not None: - self._call.add_done_callback(_report_abandoned) + self._call.add_done_callback(self._report_abandoned) self._call = None self._executor.shutdown(wait=False, cancel_futures=True) @@ -270,9 +271,9 @@ def __init__( # Wall-clock telemetry for the live rollout, opened under ``--timing`` and inert otherwise. self._telemetry = _EpisodeTelemetry() # Channels that have not delivered since this episode's RUN; the first inference waits until every one - # has. A receiver latches its last value, so a producer silent between episodes — a remote env emits - # only when stepped — would otherwise feed the previous episode's final frame. Delivery is judged by - # ``updated``, not ``ts``: some producers stamp ``ts`` on their own clock. + # has. A receiver latches its last value, so a producer silent between episodes would otherwise feed + # the previous episode's final frame. Delivery is judged by ``updated``, not ``ts``: some producers + # stamp ``ts`` on their own clock. self._awaiting_obs: set[str] = set() self.observations = pimm.ReceiverDict(self) @@ -397,7 +398,7 @@ def _begin_episode(self, context: dict[str, Any], clock: pimm.Clock) -> None: self._task.reset(self.context) if self._task is not None: self.context = {**self.context, keys.TASK: self._task.instruction} - self._worker = _InferenceWorker(self.policy, self.context, charge_wall) + self._worker = _InferenceWorker(self.policy, self.context, charge_wall, clock) self._deadline = clock.now() + self._task.timeout if self._task is not None else None self.ds_command.emit(DsWriterCommand.START()) @@ -482,7 +483,7 @@ def _build_obs(self, clock: pimm.Clock) -> dict[str, Any] | None: inputs[keys.ROBOT_FAULT] = faulted inputs[keys.WALL_TIME_NS] = time.time_ns() inputs[keys.OBS_TIME_NS] = clock.now_ns() - inputs['descriptor'] = self._embodiment.descriptor + inputs[keys.DESCRIPTOR] = self._embodiment.descriptor if not self._rollout_started: # The rollout begins at its first observation, not when the reset returned: the turns spent # delivering the scene are neither the trial's budget nor its duration. @@ -501,11 +502,11 @@ def _step(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: if worker.idle: obs = self._build_obs(clock) if obs is not None: - worker.submit(obs, clock) + worker.submit(obs) self._throttle_and_reschedule(worker, clock) def _throttle_and_reschedule(self, worker: _InferenceWorker, clock: pimm.Clock) -> None: - worker.throttle(clock) + worker.throttle() if worker.done and (trajectory := worker.result()) is not None: self._reschedule(trajectory, clock) @@ -582,8 +583,6 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock) -> Iterator[p self._telemetry.seal(clock.now()) raise finally: - # A call still in flight runs to completion and its result is dropped. The join is not deferred: - # no later episode will do it, and the policy the call holds outlives this harness. self._retire_worker() self._reap_worker() diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 228de76ed..3415e7f91 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -336,7 +336,7 @@ def test_harness_emits_cartesian_move(world): np.testing.assert_allclose(obs[keys.JOINT_VEL], np.zeros_like(robot_state.q)) assert obs[keys.GRIP] == pytest.approx(0.25) assert obs[keys.TASK] == 'stack-blocks' - assert obs['descriptor'] == '' # no descriptor passed -> empty string reaches the policy + assert obs[keys.DESCRIPTOR] == '' # no descriptor passed -> empty string reaches the policy # Recording == canonical policy I/O: the policy sees the same ``robot_state`` serializer # the dataset records. wall/obs timestamps carry volatile values, so lock the stable key set. assert set(obs) - {keys.WALL_TIME_NS, keys.OBS_TIME_NS} == { @@ -347,7 +347,7 @@ def test_harness_emits_cartesian_move(world): keys.GRIP, keys.ROBOT_FAULT, keys.TASK, - 'descriptor', + keys.DESCRIPTOR, } cmds = _emitted_commands(cmd_recorder) @@ -386,7 +386,7 @@ def test_harness_passes_descriptor_to_policy(world): drive_scheduler(scheduler, steps=20) assert policy.last_obs is not None - assert policy.last_obs['descriptor'] == 'mujoco.franka' + assert policy.last_obs[keys.DESCRIPTOR] == 'mujoco.franka' @pytest.mark.timeout(3.0) @@ -1664,8 +1664,8 @@ def test_a_reply_is_scheduled_only_while_the_trial_still_has_budget(world, expir harness = Harness(policy, make_embodiment()) now = world.clock.now() harness._deadline = now - 1.0 if expired else now + 1.0 - harness._worker = _InferenceWorker(policy, {}, charge_wall=True) - harness._worker.submit({}, world.clock) + harness._worker = _InferenceWorker(policy, {}, charge_wall=True, clock=world.clock) + harness._worker.submit({}) harness._throttle_and_reschedule(harness._worker, world.clock) diff --git a/positronic/tests/test_inference_integration.py b/positronic/tests/test_inference_integration.py index 8b3d1826b..c6a08a874 100644 --- a/positronic/tests/test_inference_integration.py +++ b/positronic/tests/test_inference_integration.py @@ -144,7 +144,7 @@ def close(self): assert keys.EE_POSE in last_obs # The task's instruction is injected by the harness. assert last_obs[keys.TASK] == 'integration-test' - assert last_obs['descriptor'] == 'mujoco.franka' + assert last_obs[keys.DESCRIPTOR] == 'mujoco.franka' class _CountdownProducer(pimm.ControlSystem): From 3152af1174b979b95eb6794ef887c742d13507f2 Mon Sep 17 00:00:00 2001 From: Sergey Arkhangelskiy <mc.vertix@gmail.com> Date: Mon, 17 Aug 2026 10:48:06 +0200 Subject: [PATCH 64/64] Give the three fake arm states one definition `test_harness.py`, `test_golden_pipeline.py` and `test_ds_writer_agent.py` each carried a `roboarm.State` fake over the same four properties. They now share `FakeRobotState` in `drivers/roboarm/tests/fakes.py`, beside the interface it fakes, with `make_robot_state` for the stationary identity-rotation case the two policy tests want. The merged fake copies on read, which the golden producer relied on and the other two are indifferent to. The golden file is unchanged. --- .../dataset/tests/test_ds_writer_agent.py | 31 ++------------ positronic/drivers/roboarm/tests/fakes.py | 41 +++++++++++++++++++ .../policy/tests/test_golden_pipeline.py | 29 +------------ positronic/policy/tests/test_harness.py | 30 +------------- 4 files changed, 48 insertions(+), 83 deletions(-) create mode 100644 positronic/drivers/roboarm/tests/fakes.py diff --git a/positronic/dataset/tests/test_ds_writer_agent.py b/positronic/dataset/tests/test_ds_writer_agent.py index fc2a45c87..8f3aaf997 100644 --- a/positronic/dataset/tests/test_ds_writer_agent.py +++ b/positronic/dataset/tests/test_ds_writer_agent.py @@ -12,8 +12,9 @@ from positronic.dataset.ds_writer_agent import DsWriterAgent, DsWriterCommand, DsWriterCommandType, TimeMode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers -from positronic.drivers.roboarm import RobotStatus, State +from positronic.drivers.roboarm import RobotStatus from positronic.drivers.roboarm import command as rcmd +from positronic.drivers.roboarm.tests.fakes import FakeRobotState from positronic.tests.testing_coutils import run_scripted_agent @@ -350,30 +351,6 @@ def test_transform_3d_serializer(world): np.testing.assert_allclose(names_vals[0][1][3:], q.as_quat) -class _FakeState(State): - def __init__(self, q, dq, ee_pose, status): - self._q = q - self._dq = dq - self._ee = ee_pose - self._status = status - - @property - def q(self): - return self._q - - @property - def dq(self): - return self._dq - - @property - def ee_pose(self): - return self._ee - - @property - def status(self): - return self._status - - def test_robot_state_serializer_drops_reset_and_emits_components(world): ds = FakeDatasetWriter() agent, cmd_em, emitters = build_agent_with_pipes({keys.ROBOT_STATE: Serializers.robot_state}, ds, world) @@ -385,8 +362,8 @@ def test_robot_state_serializer_drops_reset_and_emits_components(world): script = [ (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.START_EPISODE)), 0.001), - (lambda: emitters[keys.ROBOT_STATE].emit(_FakeState(q, dq, pose, RobotStatus.RESETTING)), 0.001), - (lambda: emitters[keys.ROBOT_STATE].emit(_FakeState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(FakeRobotState(q, dq, pose, RobotStatus.RESETTING)), 0.001), + (lambda: emitters[keys.ROBOT_STATE].emit(FakeRobotState(q, dq, pose, RobotStatus.AVAILABLE)), 0.001), (lambda: cmd_em.emit(DsWriterCommand(DsWriterCommandType.STOP_EPISODE)), 0.001), ] diff --git a/positronic/drivers/roboarm/tests/fakes.py b/positronic/drivers/roboarm/tests/fakes.py new file mode 100644 index 000000000..9440bd4a4 --- /dev/null +++ b/positronic/drivers/roboarm/tests/fakes.py @@ -0,0 +1,41 @@ +"""Fakes for the ``roboarm`` interfaces, for tests that drive an arm without a driver.""" + +import numpy as np + +from positronic import geom +from positronic.drivers.roboarm import RobotStatus, State + + +class FakeRobotState(State): + """A ``State`` over the four values it is given, copied on read so a caller cannot reach the arrays + a producer keeps emitting from.""" + + def __init__(self, q: np.ndarray, dq: np.ndarray, ee_pose: geom.Transform3D, status: RobotStatus) -> None: + self._q = q + self._dq = dq + self._ee_pose = ee_pose + self._status = status + + @property + def q(self) -> np.ndarray: + return self._q.copy() + + @property + def dq(self) -> np.ndarray: + return self._dq.copy() + + @property + def ee_pose(self) -> geom.Transform3D: + return geom.Transform3D(translation=self._ee_pose.translation.copy(), rotation=self._ee_pose.rotation) + + @property + def status(self) -> RobotStatus: + return self._status + + +def make_robot_state(translation, joints, status: RobotStatus = RobotStatus.AVAILABLE) -> FakeRobotState: + """A stationary arm at ``translation`` with identity rotation: zero joint velocity throughout.""" + translation = np.asarray(translation, dtype=np.float32) + joints = np.asarray(joints, dtype=np.float32) + pose = geom.Transform3D(translation=translation, rotation=geom.Rotation.identity) + return FakeRobotState(joints, np.zeros_like(joints), pose, status) diff --git a/positronic/policy/tests/test_golden_pipeline.py b/positronic/policy/tests/test_golden_pipeline.py index 5519d3d5a..302d5e554 100644 --- a/positronic/policy/tests/test_golden_pipeline.py +++ b/positronic/policy/tests/test_golden_pipeline.py @@ -36,9 +36,9 @@ from positronic.dataset.ds_writer_agent import TimeMode from positronic.dataset.local_dataset import LocalDataset, LocalDatasetWriter from positronic.dataset.serializers import Serializers -from positronic.drivers import roboarm from positronic.drivers.roboarm import RobotStatus from positronic.drivers.roboarm.command import CartesianPosition, CommandType, Reset +from positronic.drivers.roboarm.tests.fakes import make_robot_state from positronic.eval import ROBOT_STATIC_META, Command, Embodiment, Observation from positronic.geom import Rotation, Transform3D from positronic.policy.base import DelegatingPolicy, DelegatingSession, Now, Policy, Session @@ -130,31 +130,6 @@ def new_session(self, context=None, now=None): ) -class _FakeRobotState(roboarm.State): - """Lossless re-expression of the last applied command over sim-time.""" - - def __init__(self, pos: np.ndarray, q: np.ndarray, status: RobotStatus): - self._pos = pos - self._q = q - self._status = status - - @property - def status(self) -> RobotStatus: - return self._status - - @property - def q(self) -> np.ndarray: - return self._q.copy() - - @property - def dq(self) -> np.ndarray: - return np.zeros(7, dtype=np.float32) - - @property - def ee_pose(self) -> Transform3D: - return Transform3D(translation=self._pos.copy(), rotation=Rotation.identity) - - class FakeRobot(pimm.ControlSystem): """Deterministic closed-loop arm: applies each command as it arrives. @@ -197,7 +172,7 @@ def run(self, should_stop: pimm.SignalReceiver, clock: pimm.Clock): if self._error_pending: self._status = RobotStatus.ERROR self._error_pending = False - self.state.emit(_FakeRobotState(self._pos, self._q, self._status)) + self.state.emit(make_robot_state(self._pos, self._q, self._status)) yield pimm.Sleep(CONTROL_PERIOD_S) diff --git a/positronic/policy/tests/test_harness.py b/positronic/policy/tests/test_harness.py index 3415e7f91..df300c832 100644 --- a/positronic/policy/tests/test_harness.py +++ b/positronic/policy/tests/test_harness.py @@ -16,6 +16,7 @@ from positronic.drivers.roboarm import RobotStatus from positronic.drivers.roboarm.command import CartesianDelta, CartesianPosition, Reset, from_wire, to_wire from positronic.drivers.roboarm.models import DEFAULT_FRAME, EE_LINK, bundled_franka_model +from positronic.drivers.roboarm.tests.fakes import make_robot_state from positronic.eval import Command, Embodiment, Observation, Task from positronic.geom import Rotation, Transform3D from positronic.offboard.client import InferenceSession @@ -198,41 +199,12 @@ def new_session(self, context=None, now=None) -> RemoteSession: return RemoteSession(_FakeInferenceSession(action)) -class FakeRobotState(roboarm.State): - def __init__(self, translation: np.ndarray, joints: np.ndarray, status: RobotStatus) -> None: - self._ee_pose = Transform3D(translation=translation, rotation=Rotation.identity) - self._q = joints - self._status = status - - @property - def q(self) -> np.ndarray: - return self._q - - @property - def dq(self) -> np.ndarray: - return np.zeros_like(self._q) - - @property - def ee_pose(self) -> Transform3D: - return self._ee_pose - - @property - def status(self) -> RobotStatus: - return self._status - - @pytest.fixture def world(): with pimm.World(virtual_time=True) as w: yield w -def make_robot_state(translation, joints, status=RobotStatus.AVAILABLE) -> FakeRobotState: - translation = np.asarray(translation, dtype=np.float32) - joints = np.asarray(joints, dtype=np.float32) - return FakeRobotState(translation, joints, status) - - def emit_ready_payload(frame_emitter, robot_emitter, grip_emitter, robot_state): frame_adapter = pimm.shared_memory.NumpySMAdapter((2, 2, 3), np.dtype(np.uint8)) frame_adapter.array[:] = np.zeros((2, 2, 3), dtype=np.uint8)