adding relative actions to sync engine

src/lerobot/processor/relative_action_processor.py

@@ -126,7 +126,7 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep
+        # Always cache state for the paired AbsoluteActionsProcessorStep.
         if state is not None:
             self._last_state = state
 
@@ -146,6 +146,11 @@ class RelativeActionsProcessorStep(ProcessorStep):
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
+    def set_cached_state(self, state: torch.Tensor | None) -> None:
+        """Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
+        per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
+        self._last_state = state
+
     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,21 @@ 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.
-#
-# 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``).
+# 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
+# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
+# current (moved) state and drift through the chunk. We pin the anchor per chunk:
+# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
+# predict a fresh chunk; on the others the engine restores the anchor the relative
+# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
+# (e.g. LingBot-VA keyframe feedback) are preserved.
 
 
 class SyncInferenceEngine(InferenceEngine):
@@ -73,6 +68,31 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
+
+        # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
+        # (see module comment), mirroring the RTC engine.
+        self._relative_step = next(
+            (
+                s
+                for s in getattr(preprocessor, "steps", ())
+                if isinstance(s, RelativeActionsProcessorStep) and s.enabled
+            ),
+            None,
+        )
+        # 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)
+            self._install_chunk_probe()
+            logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
+
         logger.info(
             "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
             self._device,
@@ -85,6 +105,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:
@@ -93,6 +118,27 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
+        # New episode: the next tick predicts a fresh chunk and re-anchors.
+        self._chunk_predicted = False
+        self._ever_predicted_chunk = False
+
+    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``.
+
+        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
+            self._ever_predicted_chunk = True
+            return inner(*args, **kwargs)
+
+        self._policy.predict_action_chunk = probe
 
     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,12 +153,25 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
+        # 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.
+        # ``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(
                 observation, self._device, self._task, self._robot_type
             )
             observation = self._preprocessor(observation)
             action = self._policy.select_action(observation)
+            # Hold the anchor only for a chunking policy serving a cached action this
+            # tick; policies that never chunk or that recomputed keep refreshing.
+            if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
+                self._relative_step.set_cached_state(anchor_before)
             action = self._postprocessor(action)
         action_tensor = action.squeeze(0).cpu()
 

tests/policies/test_relative_actions.py

@@ -346,3 +346,10 @@ 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)
+
+
+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)
+    step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
+    assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}

tests/test_rollout.py

@@ -348,3 +348,200 @@ 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):
+    """Pre/post processors wrapping the real relative (caches anchor) and absolute
+    (relative + cached state) steps, mirroring 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):
+            # Run the relative step so it caches the anchor, then pass the batch through.
+            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, chunking=True):
+    """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)
+
+    def predict_action_chunk(_batch=None, **_kwargs):
+        state["predict_calls"] += 1
+        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]
+
+    def select_action(_observation):
+        if not chunking:
+            return chunk_rel[0].unsqueeze(0)
+        if len(queue) == 0:
+            actions = policy.predict_action_chunk(_observation)
+            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
+        return queue.popleft()
+
+    policy.predict_action_chunk.side_effect = predict_action_chunk
+    policy.select_action.side_effect = select_action
+    policy.reset.side_effect = queue.clear
+    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 -> 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])
+    # 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_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
+    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, 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 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():
+    """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
+
+
+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