sync engine: robust per-chunk anchor via predict_action_chunk probe

Replace the relative-action anchor mechanism from "adding relative actions
to sync engine" with one that detects fresh chunks through the public
predict_action_chunk contract instead of introspecting the policy's private
_action_queue. Drops set_hold/_hold_state in favour of a set_cached_state
snapshot/restore, keeps select_action on the hot path (so the prediction-
visualization path is unaffected), and refreshes the anchor every tick for
policies that never chunk.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Maxime Ellerbach
2026-07-17 15:13:55 +00:00
parent a185f9dde2
commit fd6aed87b7
4 changed files with 92 additions and 145 deletions
+3 -34
View File
@@ -348,39 +348,8 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
torch.testing.assert_close(result_state, original_state)
# Anchor-hold semantics (used by the synchronous rollout engine to avoid intra-chunk drift)
def _state_transition(state):
return batch_to_transition({OBS_STATE: state})
def test_set_hold_freezes_cached_anchor_state():
"""While held, the cached anchor is not overwritten, but the observation passes through."""
def test_cached_anchor_not_in_config():
"""The cached anchor is ephemeral runtime state and must not leak into the config."""
step = RelativeActionsProcessorStep(enabled=True)
state_a = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
step(_state_transition(state_a))
torch.testing.assert_close(step.get_cached_state(), state_a)
# Freeze: a new state must NOT replace the cached anchor.
step.set_hold(True)
state_b = torch.tensor([[9.0, 9.0, 9.0, 9.0]])
out = step(_state_transition(state_b))
torch.testing.assert_close(step.get_cached_state(), state_a)
# The transition's observation state is still the current one (hold only affects caching).
torch.testing.assert_close(out[TransitionKey.OBSERVATION][OBS_STATE], state_b)
# Release: caching resumes.
step.set_hold(False)
state_c = torch.tensor([[5.0, 6.0, 7.0, 8.0]])
step(_state_transition(state_c))
torch.testing.assert_close(step.get_cached_state(), state_c)
def test_hold_state_not_in_config():
"""Hold is ephemeral runtime state and must not leak into the serialized config."""
step = RelativeActionsProcessorStep(enabled=True)
step.set_hold(True)
assert "_hold_state" not in step.get_config()
step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}