From 61c9d034a48a51a44d474d0262cac87c39cb7308 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Fri, 17 Jul 2026 16:28:31 +0000 Subject: [PATCH] sync engine: address review feedback on relative-action anchoring - stop() restores the original predict_action_chunk so the policy object is not left permanently patched after the engine goes away - snapshot the chunk anchor with clone() (guarded against None on the first tick) so it survives any future in-place mutation of the cached state - document the engine->step action_names lazy-fill and the torch.compile ordering the probe relies on - add tests for reset() re-anchoring across episodes and stop() probe restore Co-Authored-By: Claude Opus 4.8 (1M context) --- src/lerobot/rollout/inference/sync.py | 25 ++++++++++++++++--- tests/test_rollout.py | 36 +++++++++++++++++++++++++++ 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index 8d0cdd653..ee7032fb7 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/src/lerobot/rollout/inference/sync.py @@ -83,7 +83,11 @@ class SyncInferenceEngine(InferenceEngine): # Set by the probe for the current tick / ever, respectively. self._chunk_predicted = False self._ever_predicted_chunk = False + self._original_predict_action_chunk = None # set while the probe is installed if self._relative_step is not None: + # ``action_names`` is optional on the step; fill it lazily from the + # policy/dataset so the relative<->absolute mask is built correctly. This is + # a deliberate engine->step side effect (the step is configured by its consumer). if self._relative_step.action_names is None: cfg_names = getattr(policy.config, "action_feature_names", None) self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys) @@ -112,6 +116,11 @@ class SyncInferenceEngine(InferenceEngine): def stop(self) -> None: """No background resources to stop.""" + # Undo the probe so the policy object isn't left permanently patched + # (it may outlive this engine or be reused by another). + if self._original_predict_action_chunk is not None: + self._policy.predict_action_chunk = self._original_predict_action_chunk + self._original_predict_action_chunk = None logger.info("SyncInferenceEngine stopped") def reset(self) -> None: @@ -129,8 +138,13 @@ class SyncInferenceEngine(InferenceEngine): def _install_chunk_probe(self) -> None: """Wrap the policy's public ``predict_action_chunk`` so we learn which ticks predict a fresh chunk (when the anchor must advance) without introspecting any - private action queue. Chunking policies call it from ``select_action``.""" - inner = self._policy.predict_action_chunk + private action queue. Chunking policies call it from ``select_action``. + + Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d + one, since ``build_rollout_context`` compiles before building the engine); undone + in ``stop()``.""" + self._original_predict_action_chunk = self._policy.predict_action_chunk + inner = self._original_predict_action_chunk def probe(*args, **kwargs): self._chunk_predicted = True @@ -179,7 +193,12 @@ class SyncInferenceEngine(InferenceEngine): ) # Snapshot the chunk anchor before the preprocessor overwrites it with this # tick's state; restore it below if this tick only served a cached action. - anchor_before = self._relative_step.get_cached_state() if self._relative_step is not None else None + # ``clone`` so the snapshot survives even if the cached tensor is ever mutated + # in place (today it is only rebound, but the copy is cheap for a state vector). + anchor_before = None + if self._relative_step is not None: + cached = self._relative_step.get_cached_state() + anchor_before = cached.clone() if cached is not None else None self._chunk_predicted = False with torch.inference_mode(), autocast_ctx: observation = prepare_observation_for_inference( diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 27956e3d6..e9aa2c543 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -486,6 +486,29 @@ def test_sync_relative_holds_anchor_across_chunk(): torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next])) +def test_sync_relative_reset_reanchors_new_episode(): + """After ``reset()`` the first tick of the next episode anchors to the new state.""" + n = 3 + chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)]) + pre, post, relative_step = _relative_pre_post() + policy = _fake_relative_policy(chunk_rel, n_action_steps=n) + engine = _build_sync_engine(policy, pre, post) + + # Episode 1: one tick anchors to s0 and leaves cached actions in the queue. + engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0])) + assert policy._predict_state["predict_calls"] == 1 + + engine.reset() # clears the queue and the per-episode chunk flags + + # Episode 2: a fresh state must produce a fresh chunk anchored to that state, + # not carry over the previous episode's anchor. + s_new = [7.0, 8.0, 9.0, 10.0] + out = engine.get_action(_obs_frame(s_new)) + assert policy._predict_state["predict_calls"] == 2 + torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0]) + torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new])) + + def test_sync_relative_non_chunking_policy_refreshes_every_tick(): """A policy that never calls ``predict_action_chunk`` must not freeze the anchor.""" n = 3 @@ -509,3 +532,16 @@ def test_sync_engine_no_relative_step_is_none(): policy.config.use_amp = False engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock()) assert engine._relative_step is None + + +def test_sync_relative_stop_restores_policy_method(): + """``stop()`` un-patches the probe so the policy object isn't permanently modified.""" + n = 3 + chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)]) + pre, post, _ = _relative_pre_post() + policy = _fake_relative_policy(chunk_rel, n_action_steps=n) + original = policy.predict_action_chunk + engine = _build_sync_engine(policy, pre, post) + assert policy.predict_action_chunk is not original # probe installed + engine.stop() + assert policy.predict_action_chunk is original # restored