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) <noreply@anthropic.com>

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``."""
+        private action queue. Chunking policies call it from ``select_action``.
-        inner = self._policy.predict_action_chunk
+
+        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(

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