From 01a029a05ee780b72121a6ad8c5ada7249619f1b Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Thu, 9 Jul 2026 14:39:06 +0000 Subject: [PATCH] adding relative actions to sync engine --- .../processor/relative_action_processor.py | 20 ++- src/lerobot/rollout/context.py | 11 -- src/lerobot/rollout/inference/sync.py | 91 +++++++--- tests/policies/test_relative_actions.py | 38 +++++ tests/test_rollout.py | 161 ++++++++++++++++++ 5 files changed, 288 insertions(+), 33 deletions(-) diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index 796c376f0..b6e1e2826 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/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, diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 20a7d715a..d3320ccf3 100644 --- a/src/lerobot/rollout/context.py +++ b/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)...", diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index 2bb05b6ab..d8ed591e3 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/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 diff --git a/tests/policies/test_relative_actions.py b/tests/policies/test_relative_actions.py index 15ef0a31b..4dc97b9b5 100644 --- a/tests/policies/test_relative_actions.py +++ b/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"} diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 85a29ff4c..450c2b53f 100644 --- a/tests/test_rollout.py +++ b/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