diff --git a/positronic/policy/base.py b/positronic/policy/base.py index b3d93fac9..a76cd0bec 100644 --- a/positronic/policy/base.py +++ b/positronic/policy/base.py @@ -73,6 +73,19 @@ def __call__(self, obs: Mapping[str, Any], time_ns: int) -> list[dict[str, Any]] ``time_ns`` is the caller's clock reading in nanoseconds. A session reads no clock of its own. """ + def reads_observation(self, obs: Mapping[str, Any], time_ns: int) -> bool: + """Whether the call this query precedes would read the observation it is given. + + Answering False permits the caller to hand over whatever it already has instead of building + what this session asked for, so a session whose ``__call__`` reads the observation must answer + True — including one that only records it. The call happens either way. + + It takes what ``__call__`` takes, so a session deciding on a field decides on the same value in + both: a query answered against the caller's clock and a call answered against the observation's + own stamp disagree wherever the two readings straddle a boundary. + """ + return True + @property def meta(self) -> dict[str, Any]: """What this session reports about its model and its episode.""" diff --git a/positronic/policy/codec.py b/positronic/policy/codec.py index 78c0f6bb2..b837c7619 100644 --- a/positronic/policy/codec.py +++ b/positronic/policy/codec.py @@ -114,8 +114,13 @@ def __init__(self, inner: Session, codec: 'Codec'): super().__init__(inner) self._codec = codec + def reads_observation(self, obs, time_ns): + # Forwarded because this session encodes rather than reads: told no, it hands its own input on. + return self._inner.reads_observation(obs, time_ns) + def __call__(self, obs, time_ns): - encoded = self._codec.encode(obs) + reads = self._inner.reads_observation(obs, time_ns) + encoded = self._codec.encode(obs) if reads else obs action = self._inner(encoded, time_ns) if action is None: return None diff --git a/positronic/policy/layers.py b/positronic/policy/layers.py index 530039e06..96b43bee9 100644 --- a/positronic/policy/layers.py +++ b/positronic/policy/layers.py @@ -80,6 +80,12 @@ def __init__(self, inner: Session): super().__init__(inner) self._trajectory_end: float | None = None + def reads_observation(self, obs, time_ns): + # The same guard ``__call__`` makes, against the same instant, so the two cannot disagree. + if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: + return False + return self._inner.reads_observation(obs, time_ns) + def __call__(self, obs, time_ns): if self._trajectory_end is not None and _obs_time(obs) < self._trajectory_end: return None @@ -180,7 +186,11 @@ def __init__(self, inner: Session, keys: tuple[str, ...], offsets_sec: tuple[flo def __call__(self, obs, time_ns): now = _obs_time(obs) + # Unconditional, including on the ticks that skip sampling below: a gap here is a hole in + # the window every later sample is taken from. self._buffer.append(now, {k: obs[k] for k in self._keys}) + if not self._inner.reads_observation(obs, time_ns): + return self._inner(obs, time_ns) return self._inner({**obs, **self._buffer.sample(now)}, time_ns) def cancel(self): diff --git a/positronic/policy/remote.py b/positronic/policy/remote.py index 875e5620b..1d5f9aa1a 100644 --- a/positronic/policy/remote.py +++ b/positronic/policy/remote.py @@ -93,6 +93,9 @@ def __call__(self, obs: cabc.Mapping[str, Any], time_ns: int) -> list[dict[str, return None return [result] if isinstance(result, dict) else result + def reads_observation(self, obs: cabc.Mapping[str, Any], time_ns: int) -> bool: + return self._answer is None + def cancel(self): # The cancel says the world the chunk applies to has gone. The session still reads the round trip # for its failure, and drops the chunk that comes with it. diff --git a/positronic/policy/tests/test_layers.py b/positronic/policy/tests/test_layers.py index 28a5d1007..a95d581f0 100644 --- a/positronic/policy/tests/test_layers.py +++ b/positronic/policy/tests/test_layers.py @@ -12,7 +12,7 @@ from positronic.geom import Rotation, Transform3D from positronic.policy import spec from positronic.policy.action import AbsoluteJointsAction, AbsolutePositionAction, IKJointsAction, JointDeltaAction -from positronic.policy.base import Layer, Policy, Session +from positronic.policy.base import DelegatingSession, Layer, Policy, Session from positronic.policy.codec import ( ActionHorizon, ActionTimestamp, @@ -326,6 +326,87 @@ def test_every_arm_channel_is_stamped(self): np.testing.assert_array_equal(decoded[keys.TARGET_JOINTS], np.zeros(7)) +class _DecliningSession(Session): + """A session that reads no observation — what a remote one answers with a round trip in flight.""" + + def __init__(self): + self.seen: dict[str, Any] = {} + + def __call__(self, obs, time_ns): + self.seen = dict(obs) + return None + + def reads_observation(self, obs, time_ns): + return False + + +class _CapturingInner(Session): + def __init__(self): + self.seen: dict[str, Any] = {} + + def __call__(self, obs, time_ns): + self.seen = dict(obs) + return None + + +class TestReadsObservation: + """The gate that lets a layer skip building an observation nothing below will read.""" + + def test_a_session_reads_its_observation_by_default(self): + assert _ConstSession(None).reads_observation(_obs(0), 0) is True + + def test_a_schedule_reads_nothing_while_its_chunk_plays(self): + session = ChunkedSchedule().make_session(_ConstSession([{keys.ACTION_TIMESTAMP: 2.0}])) + assert session.reads_observation(_obs(0), 0) is True # nothing emitted yet + session(_obs(0), 0) # emits, ending 2 s from now + assert session.reads_observation(_obs(1.0), int(1.0e9)) is False # mid-chunk + assert session.reads_observation(_obs(3.0), int(3.0e9)) is True # played out + + def test_the_gate_agrees_with_the_call(self): + """Both read `_trajectory_end`, in two places, so they are pinned against each other: a call + that answers None is one the gate would have declined.""" + session = ChunkedSchedule().make_session(_ConstSession([{keys.ACTION_TIMESTAMP: 2.0}])) + session(_obs(0), 0) + for t_ns in (int(0.5e9), int(1.9e9), int(2.0e9), int(2.1e9), int(5.0e9)): + declined = not session.reads_observation(_obs(t_ns / 1e9), t_ns) + answered_none = session(_obs(t_ns / 1e9), t_ns) is None + assert declined is answered_none, f'disagreed at {t_ns} ns' + + def test_the_gate_reads_the_observation_instant_the_call_reads(self): + """An observation ahead of the caller's clock: a gate reading the clock would decline, while + the call goes through to the inner session — which then gets whatever the layer above kept.""" + session = ChunkedSchedule().make_session(_ConstSession([{keys.ACTION_TIMESTAMP: 2.0}])) + session(_obs(0), 0) + assert session.reads_observation(_obs(2.5), int(1.0e9)) is True + assert session(_obs(2.5), int(1.0e9)) is not None + + def test_a_delegating_session_answers_for_itself(self): + """A session that wraps another still reads its own observation — a recording tap is the live + case. Forwarding the inner answer would let a layer above hand it something else to log.""" + + class _Tap(DelegatingSession): + def __call__(self, obs, time_ns): + self.logged = dict(obs) + return self._inner(obs, time_ns) + + assert _Tap(_DecliningSession()).reads_observation(_obs(0), 0) is True + + def test_a_stack_is_not_sampled_when_nothing_below_reads_it(self): + """The window is the expensive part, and most ticks do not use it. Recording still happens on + every tick — the history it builds is what the next real call is sampled from.""" + inner = _DecliningSession() + session = TemporalStack(keys=('cam',), offsets_sec=(-0.2, 0.0)).make_session(inner) + session({**_obs(0), 'cam': np.zeros((4, 4, 3), np.uint8)}, 0) + assert inner.seen['cam'].shape == (4, 4, 3), 'passed through unstacked' + + def test_a_stack_is_sampled_when_something_below_reads_it(self): + inner = _CapturingInner() + session = TemporalStack(keys=('cam',), offsets_sec=(-0.2, 0.0)).make_session(inner) + session({**_obs(0), 'cam': np.zeros((4, 4, 3), np.uint8)}, 0) + # A stacked entry carries the window's length on a new leading axis. + assert inner.seen['cam'].shape == (2, 4, 4, 3) + + class TestTemporalStack: OFFSETS = (-0.2, -0.1, 0.0)