adding relative actions to sync engine

src/lerobot/processor/relative_action_processor.py

@@ -112,6 +112,12 @@ 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:
@@ -136,8 +142,9 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep
-        if state is not None:
+        # Always cache state for the paired AbsoluteActionsProcessorStep, unless a
+        # chunked inference engine has frozen the anchor for the current chunk.
+        if state is not None and not self._hold_state:
             self._last_state = state
 
         if not self.enabled:
@@ -156,6 +163,15 @@ 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:
+        """Freeze (``True``) or resume (``False``) refreshing of the cached anchor state.
+
+        While held, ``__call__`` keeps the previously cached ``_last_state`` so the
+        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 {
             "enabled": self.enabled,

src/lerobot/rollout/context.py

@@ -43,7 +43,6 @@ from lerobot.processor import (
     make_default_processors,
     rename_stats,
 )
-from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
@@ -52,7 +51,6 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
     InferenceEngine,
     RTCInferenceConfig,
-    SyncInferenceConfig,
     create_inference_engine,
 )
 from .robot_wrapper import ThreadSafeRobot
@@ -399,15 +397,6 @@ def build_rollout_context(
         },
     )
 
-    if isinstance(cfg.inference, SyncInferenceConfig) and any(
-        isinstance(step, RelativeActionsProcessorStep) and step.enabled
-        for step in getattr(preprocessor, "steps", ())
-    ):
-        raise NotImplementedError(
-            "SyncInferenceEngine does not support policies with relative actions for now."
-            "Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
-        )
-
     # --- 7. Inference strategy (needs policy + pre/post + hardware) --
     logger.info(
         "Creating inference engine (type=%s)...",

src/lerobot/rollout/inference/sync.py

@@ -24,26 +24,32 @@ import torch
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
-from lerobot.processor import PolicyProcessorPipeline
+from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
 
 from .base import InferenceEngine
 
 logger = logging.getLogger(__name__)
 
 
-# TODO(Steven): support relative-action policies.  The per-tick flow refreshes
-# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
-# actions popped on later ticks get reanchored to the *current* robot state and
-# absolute targets drift through the chunk.  Relative-action policies are
-# rejected at context-build time today; RTC postprocesses the whole chunk and
-# is unaffected.
+# Relative-action support (drift-free anchoring)
+# ----------------------------------------------
+# Relative-action policies predict a *chunk* of offsets anchored to the robot
+# state at chunk-prediction time.  ``select_action`` serves that chunk one action
+# per tick from an internal ``_action_queue``, recomputing only when the queue is
+# empty.  The per-tick flow here runs the full pre/post pipeline every call, and
+# ``RelativeActionsProcessorStep`` would otherwise refresh its cached anchor state
+# on every tick — so actions popped from the queue on later ticks would be
+# re-anchored to the *current* (already-moved) state and absolute targets would
+# drift through the chunk.
 #
-# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
-# local FIFO of postprocessed actions.  Eliminates drift by construction and
-# saves per-tick pre/post work, but bypasses ``select_action`` — needs
-# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
-# ``select_action``), and Diffusion-family (obs-history queues populated as a
-# side effect of ``select_action``).
+# 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):
@@ -73,6 +79,24 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
+
+        # Relative-action policies need the chunk anchor held while cached actions
+        # are popped (see module docstring).  Introspect the preprocessor for an
+        # enabled RelativeActionsProcessorStep, mirroring the RTC engine.
+        self._relative_step = next(
+            (
+                s
+                for s in getattr(preprocessor, "steps", ())
+                if isinstance(s, RelativeActionsProcessorStep) and s.enabled
+            ),
+            None,
+        )
+        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")
+
         logger.info(
             "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
             self._device,
@@ -93,6 +117,21 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
+        # ``policy.reset()`` empties ``_action_queue`` so the next ``get_action``
+        # recomputes and refreshes the anchor; clear any leftover hold defensively.
+        if self._relative_step is not None:
+            self._relative_step.set_hold(False)
+
+    def _policy_will_recompute(self) -> bool:
+        """True if the next ``select_action`` will predict a fresh chunk (queue empty/absent).
+
+        Relative-action policies expose an ``_action_queue`` deque that is refilled
+        only when empty.  When it is non-empty the upcoming ``select_action`` will
+        pop a cached action, so the anchor state must be held.  Policies without the
+        attribute recompute every tick, so we always refresh the anchor.
+        """
+        queue = getattr(self._policy, "_action_queue", None)
+        return queue is None or len(queue) == 0
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -107,13 +146,25 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
-        with torch.inference_mode(), autocast_ctx:
-            observation = prepare_observation_for_inference(
-                observation, self._device, self._task, self._robot_type
-            )
-            observation = self._preprocessor(observation)
-            action = self._policy.select_action(observation)
-            action = self._postprocessor(action)
+        # For relative-action policies, hold the cached anchor on ticks that pop a
+        # cached action so the whole chunk stays anchored to the state captured when
+        # it was predicted.  Decided before the pipeline runs (the queue is drained
+        # inside ``select_action``); always released in ``finally`` so a hold never
+        # leaks across ticks or on exception.
+        hold_anchor = self._relative_step is not None and not self._policy_will_recompute()
+        if self._relative_step is not None:
+            self._relative_step.set_hold(hold_anchor)
+        try:
+            with torch.inference_mode(), autocast_ctx:
+                observation = prepare_observation_for_inference(
+                    observation, self._device, self._task, self._robot_type
+                )
+                observation = self._preprocessor(observation)
+                action = self._policy.select_action(observation)
+                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

@@ -346,3 +346,41 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
 
     result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
     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."""
+    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()
+    assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}

tests/test_rollout.py

@@ -348,3 +348,164 @@ def test_rollout_context_fields():
 
     field_names = {f.name for f in dataclasses.fields(RolloutContext)}
     assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
+
+
+# ---------------------------------------------------------------------------
+# Sync engine: relative-action anchoring (drift-free chunk execution)
+# ---------------------------------------------------------------------------
+
+_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
+_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
+
+
+def _relative_pre_post(exclude_joints=None):
+    """Build fake pre/post processors wrapping real relative/absolute steps.
+
+    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,
+        TransitionKey,
+        create_transition,
+    )
+    from lerobot.utils.constants import OBS_STATE
+
+    relative_step = RelativeActionsProcessorStep(
+        enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
+    )
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    class _Pre:
+        steps = [relative_step]
+
+        def __call__(self, observation):
+            # observation carries a batched OBS_STATE tensor; run the relative step so
+            # 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
+
+        def reset(self):
+            pass
+
+    class _Post:
+        def __call__(self, action):
+            transition = create_transition(action=action)
+            return absolute_step(transition)[TransitionKey.ACTION]
+
+        def reset(self):
+            pass
+
+    return _Pre(), _Post(), relative_step
+
+
+def _fake_relative_policy(chunk_rel, n_action_steps, with_queue=True):
+    """Fake chunk policy: refills an ``_action_queue`` with ``chunk_rel`` when empty."""
+    from collections import deque
+
+    policy = MagicMock()
+    policy.config.use_amp = False
+    policy.config.action_feature_names = list(_REL_ACTION_NAMES)
+    state = {"predict_calls": 0}
+
+    if with_queue:
+        policy._action_queue = deque(maxlen=n_action_steps)
+    else:
+        # Ensure the attribute is truly absent so getattr(...) falls back.
+        del policy._action_queue
+
+    def select_action(_observation):
+        if with_queue:
+            if len(policy._action_queue) == 0:
+                state["predict_calls"] += 1
+                policy._action_queue.extend(chunk_rel[i].unsqueeze(0) for i in range(n_action_steps))
+            return policy._action_queue.popleft()
+        # No queue: recompute every tick (like temporal ensembling).
+        state["predict_calls"] += 1
+        return chunk_rel[0].unsqueeze(0)
+
+    policy.select_action.side_effect = select_action
+    policy.reset.side_effect = lambda: policy._action_queue.clear() if with_queue else None
+    policy._predict_state = state
+    return policy
+
+
+def _build_sync_engine(policy, pre, post):
+    from lerobot.rollout import SyncInferenceEngine
+
+    return SyncInferenceEngine(
+        policy=policy,
+        preprocessor=pre,
+        postprocessor=post,
+        dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
+        ordered_action_keys=list(_REL_ACTION_NAMES),
+        task="test",
+        device="cpu",
+        robot_type="mock",
+    )
+
+
+def _obs_frame(state_values):
+    import numpy as np
+
+    return {"observation.state": np.asarray(state_values, dtype=np.float32)}
+
+
+def test_sync_relative_holds_anchor_across_chunk():
+    """Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
+    n = 4
+    # A distinct relative offset per chunk step so a wrong anchor would be visible.
+    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i 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)
+
+    assert engine._relative_step is relative_step  # introspection wired the step
+
+    s0 = [1.0, 2.0, 3.0, 4.0]
+    outputs = []
+    for tick in range(n):
+        # Feed a *different* state each tick; a drifting anchor would use it.
+        state = [v + tick for v in s0]
+        outputs.append(engine.get_action(_obs_frame(state)))
+
+    # Exactly one chunk was predicted across the n ticks.
+    assert policy._predict_state["predict_calls"] == 1
+    for tick in range(n):
+        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.
+    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
+
+
+def test_sync_relative_fallback_without_action_queue():
+    """A policy without ``_action_queue`` refreshes the anchor every tick."""
+    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)
+    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.
+        torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
+
+
+def test_sync_engine_no_relative_step_is_none():
+    """Without an enabled relative step, the engine takes the plain select_action path."""
+    policy = MagicMock()
+    policy.config.use_amp = False
+    engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
+    assert engine._relative_step is None