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>

src/lerobot/processor/relative_action_processor.py

@@ -112,12 +112,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
     exclude_joints: list[str] = field(default_factory=list)
     action_names: list[str] | None = None
     _last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
-    # When True, ``__call__`` stops refreshing ``_last_state``. Chunked inference
-    # engines set this so every action popped from a policy's action queue is
-    # re-anchored to the state captured when the chunk was predicted, instead of
-    # drifting against the current per-tick state. Episode-scoped runtime state,
-    # not part of ``get_config``.
-    _hold_state: bool = field(default=False, init=False, repr=False)
 
     def _build_mask(self, action_dim: int) -> list[bool]:
         if not self.exclude_joints or self.action_names is None:
@@ -142,9 +136,8 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep, unless a
+        # Always cache state for the paired AbsoluteActionsProcessorStep.
-        # chunked inference engine has frozen the anchor for the current chunk.
+        if state is not None:
-        if state is not None and not self._hold_state:
             self._last_state = state
 
         if not self.enabled:
@@ -163,14 +156,10 @@ class RelativeActionsProcessorStep(ProcessorStep):
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
-    def set_hold(self, flag: bool) -> None:
+    def set_cached_state(self, state: torch.Tensor | None) -> None:
-        """Freeze (``True``) or resume (``False``) refreshing of the cached anchor state.
+        """Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
-
+        per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
-        While held, ``__call__`` keeps the previously cached ``_last_state`` so the
+        self._last_state = state
-        paired :class:`AbsoluteActionsProcessorStep` re-anchors every action of a
-        predicted chunk to the same state, avoiding intra-chunk drift.
-        """
-        self._hold_state = flag
 
     def get_config(self) -> dict[str, Any]:
         return {

src/lerobot/rollout/inference/sync.py

@@ -31,25 +31,14 @@ from .base import InferenceEngine
 logger = logging.getLogger(__name__)
 
 
-# Relative-action support (drift-free anchoring)
+# Relative-action support: a predicted chunk of offsets is anchored to the robot
-# ----------------------------------------------
+# state at prediction time, but the sync engine reruns the pre/post pipeline every
-# Relative-action policies predict a *chunk* of offsets anchored to the robot
+# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
-# state at chunk-prediction time.  ``select_action`` serves that chunk one action
+# current (moved) state and drift through the chunk. We pin the anchor per chunk:
-# per tick from an internal ``_action_queue``, recomputing only when the queue is
+# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
-# empty.  The per-tick flow here runs the full pre/post pipeline every call, and
+# predict a fresh chunk; on the others the engine restores the anchor the relative
-# ``RelativeActionsProcessorStep`` would otherwise refresh its cached anchor state
+# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
-# on every tick — so actions popped from the queue on later ticks would be
+# (e.g. LingBot-VA keyframe feedback) are preserved.
-# re-anchored to the *current* (already-moved) state and absolute targets would
-# drift through the chunk.
-#
-# Fix: detect chunk boundaries by inspecting the policy's ``_action_queue`` length
-# *before* running the pipeline, and freeze the relative step's cached anchor
-# (``set_hold``) on ticks that pop a cached action.  The whole chunk is then
-# anchored to a single state, exactly like RTC.  ``select_action`` stays on the
-# hot path, so policy-specific side effects (e.g. LingBot-VA's per-tick keyframe
-# feedback) are preserved.  Policies without an ``_action_queue`` (e.g. ACT
-# temporal ensembling, which recomputes every tick) fall back to refreshing the
-# anchor every tick, which is the correct behaviour there.
 
 
 class SyncInferenceEngine(InferenceEngine):
@@ -81,9 +70,8 @@ class SyncInferenceEngine(InferenceEngine):
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
 
-        # Relative-action policies need the chunk anchor held while cached actions
+        # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
-        # are popped (see module docstring).  Introspect the preprocessor for an
+        # (see module comment), mirroring the RTC engine.
-        # enabled RelativeActionsProcessorStep, mirroring the RTC engine.
         self._relative_step = next(
             (
                 s
@@ -92,11 +80,15 @@ class SyncInferenceEngine(InferenceEngine):
             ),
             None,
         )
+        # Set by the probe for the current tick / ever, respectively.
+        self._chunk_predicted = False
+        self._ever_predicted_chunk = False
         if self._relative_step is not None:
             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)
-            logger.info("Relative actions enabled: chunk anchor will be held per chunk")
+            self._install_chunk_probe()
+            logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
 
         # Intermediate-prediction visualization (e.g. a world model's imagined video). When on,
         # ``get_action`` requests predictions and keeps the current chunk's frame stacks; a playhead
@@ -128,23 +120,24 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
-        # ``policy.reset()`` empties ``_action_queue`` so the next ``get_action``
+        # New episode: the next tick predicts a fresh chunk and re-anchors.
-        # recomputes and refreshes the anchor; clear any leftover hold defensively.
+        self._chunk_predicted = False
-        if self._relative_step is not None:
+        self._ever_predicted_chunk = False
-            self._relative_step.set_hold(False)
         self._pred_stacks = {}
         self._pred_cursor = 0
 
-    def _policy_will_recompute(self) -> bool:
+    def _install_chunk_probe(self) -> None:
-        """True if the next ``select_action`` will predict a fresh chunk (queue empty/absent).
+        """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
 
-        Relative-action policies expose an ``_action_queue`` deque that is refilled
+        def probe(*args, **kwargs):
-        only when empty.  When it is non-empty the upcoming ``select_action`` will
+            self._chunk_predicted = True
-        pop a cached action, so the anchor state must be held.  Policies without the
+            self._ever_predicted_chunk = True
-        attribute recompute every tick, so we always refresh the anchor.
+            return inner(*args, **kwargs)
-        """
+
-        queue = getattr(self._policy, "_action_queue", None)
+        self._policy.predict_action_chunk = probe
-        return queue is None or len(queue) == 0
 
     def get_intermediate_predictions(self) -> dict | None:
         """Serve one imagined frame per key for this tick, advancing the playhead.
@@ -184,34 +177,30 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
-        # For relative-action policies, hold the cached anchor on ticks that pop a
+        # Snapshot the chunk anchor before the preprocessor overwrites it with this
-        # cached action so the whole chunk stays anchored to the state captured when
+        # tick's state; restore it below if this tick only served a cached action.
-        # it was predicted.  Decided before the pipeline runs (the queue is drained
+        anchor_before = self._relative_step.get_cached_state() if self._relative_step is not None else None
-        # inside ``select_action``); always released in ``finally`` so a hold never
+        self._chunk_predicted = False
-        # leaks across ticks or on exception.
+        with torch.inference_mode(), autocast_ctx:
-        hold_anchor = self._relative_step is not None and not self._policy_will_recompute()
+            observation = prepare_observation_for_inference(
-        if self._relative_step is not None:
+                observation, self._device, self._task, self._robot_type
-            self._relative_step.set_hold(hold_anchor)
+            )
-        try:
+            observation = self._preprocessor(observation)
-            with torch.inference_mode(), autocast_ctx:
+            if self._visualize_predictions:
-                observation = prepare_observation_for_inference(
+                action, predictions = unpack_action_output(
-                    observation, self._device, self._task, self._robot_type
+                    self._policy.select_action(observation, return_intermediate_predictions=True)
                 )
-                observation = self._preprocessor(observation)
+                if predictions:
-                if self._visualize_predictions:
+                    # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead.
-                    action, predictions = unpack_action_output(
+                    self._pred_stacks = predictions
-                        self._policy.select_action(observation, return_intermediate_predictions=True)
+                    self._pred_cursor = 0
-                    )
+            else:
-                    if predictions:
+                action = self._policy.select_action(observation)
-                        # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead.
+            # Hold the anchor only for a chunking policy serving a cached action this
-                        self._pred_stacks = predictions
+            # tick; policies that never chunk or that recomputed keep refreshing.
-                        self._pred_cursor = 0
+            if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
-                else:
+                self._relative_step.set_cached_state(anchor_before)
-                    action = self._policy.select_action(observation)
+            action = self._postprocessor(action)
-                action = self._postprocessor(action)
-        finally:
-            if self._relative_step is not None:
-                self._relative_step.set_hold(False)
         action_tensor = action.squeeze(0).cpu()
 
         # Reorder to match dataset action ordering so the caller can treat

tests/policies/test_relative_actions.py

@@ -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 test_cached_anchor_not_in_config():
-
+    """The cached anchor is ephemeral runtime state and must not leak into the config."""
-
-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."""
     step = RelativeActionsProcessorStep(enabled=True)
-
+    step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
-    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()
     assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}

tests/test_rollout.py

@@ -359,13 +359,8 @@ _REL_ACTION_DIM = len(_REL_ACTION_NAMES)
 
 
 def _relative_pre_post(exclude_joints=None):
-    """Build fake pre/post processors wrapping real relative/absolute steps.
+    """Pre/post processors wrapping the real relative (caches anchor) and absolute
-
+    (relative + cached state) steps, mirroring what the sync engine feeds them."""
-    The preprocessor runs the ``RelativeActionsProcessorStep`` (caching/holding the
-    anchor state) and passes the observation through; the postprocessor runs the
-    paired ``AbsoluteActionsProcessorStep`` (relative + cached state) and returns the
-    absolute action tensor.  Shapes mirror what the sync engine feeds them.
-    """
     from lerobot.processor import (
         AbsoluteActionsProcessorStep,
         RelativeActionsProcessorStep,
@@ -383,8 +378,7 @@ def _relative_pre_post(exclude_joints=None):
         steps = [relative_step]
 
         def __call__(self, observation):
-            # observation carries a batched OBS_STATE tensor; run the relative step so
+            # Run the relative step so it caches the anchor, then pass the batch through.
-            # it caches (or holds) the anchor, then pass the batch through unchanged.
             transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
             relative_step(transition)
             return observation
@@ -403,33 +397,37 @@ def _relative_pre_post(exclude_joints=None):
     return _Pre(), _Post(), relative_step
 
 
-def _fake_relative_policy(chunk_rel, n_action_steps, with_queue=True):
+def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
-    """Fake chunk policy: refills an ``_action_queue`` with ``chunk_rel`` when empty."""
+    """Fake relative-action policy for the sync engine.
+
+    ``chunking=True`` buffers a chunk and serves it one action per tick, calling the
+    public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
+    returns an action directly and never calls it. The engine's anchor probe keys off
+    that public call, so the fake routes through it rather than any private queue.
+    """
     from collections import deque
 
     policy = MagicMock()
     policy.config.use_amp = False
     policy.config.action_feature_names = list(_REL_ACTION_NAMES)
     state = {"predict_calls": 0}
+    queue = deque(maxlen=n_action_steps)
 
-    if with_queue:
+    def predict_action_chunk(_batch=None, **_kwargs):
-        policy._action_queue = deque(maxlen=n_action_steps)
+        state["predict_calls"] += 1
-    else:
+        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]
-        # Ensure the attribute is truly absent so getattr(...) falls back.
-        del policy._action_queue
 
     def select_action(_observation):
-        if with_queue:
+        if not chunking:
-            if len(policy._action_queue) == 0:
+            return chunk_rel[0].unsqueeze(0)
-                state["predict_calls"] += 1
+        if len(queue) == 0:
-                policy._action_queue.extend(chunk_rel[i].unsqueeze(0) for i in range(n_action_steps))
+            actions = policy.predict_action_chunk(_observation)
-            return policy._action_queue.popleft()
+            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
-        # No queue: recompute every tick (like temporal ensembling).
+        return queue.popleft()
-        state["predict_calls"] += 1
-        return chunk_rel[0].unsqueeze(0)
 
+    policy.predict_action_chunk.side_effect = predict_action_chunk
     policy.select_action.side_effect = select_action
-    policy.reset.side_effect = lambda: policy._action_queue.clear() if with_queue else None
+    policy.reset.side_effect = queue.clear
     policy._predict_state = state
     return policy
 
@@ -479,28 +477,30 @@ def test_sync_relative_holds_anchor_across_chunk():
         expected = torch.tensor(s0) + chunk_rel[tick]
         torch.testing.assert_close(outputs[tick], expected)
 
-    # Next tick empties the queue -> recompute -> anchor refreshes to the new state.
+    # Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
     s_next = [10.0, 20.0, 30.0, 40.0]
     out = engine.get_action(_obs_frame(s_next))
     assert policy._predict_state["predict_calls"] == 2
     torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
-    assert relative_step._hold_state is False  # released after every call
+    # The anchor now reflects the fresh-chunk state, not the held one.
+    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
 
 
-def test_sync_relative_fallback_without_action_queue():
+def test_sync_relative_non_chunking_policy_refreshes_every_tick():
-    """A policy without ``_action_queue`` refreshes the anchor every tick."""
+    """A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
     n = 3
     chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
     pre, post, _ = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, with_queue=False)
+    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
     engine = _build_sync_engine(policy, pre, post)
 
     s0 = [1.0, 1.0, 1.0, 1.0]
     for tick in range(3):
         state = [v + tick for v in s0]
         out = engine.get_action(_obs_frame(state))
-        # Anchor tracks the current state every tick.
+        # Anchor must track the current state every tick (no chunk => no hold).
         torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
+    assert policy._predict_state["predict_calls"] == 0
 
 
 def test_sync_engine_no_relative_step_is_none():