diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index 5b1039291..884edd7b0 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/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, 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..5c7e5b479 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/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() diff --git a/tests/policies/test_relative_actions.py b/tests/policies/test_relative_actions.py index 15ef0a31b..1867183c4 100644 --- a/tests/policies/test_relative_actions.py +++ b/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"} diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 85a29ff4c..e9aa2c543 100644 --- a/tests/test_rollout.py +++ b/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