diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index b6e1e2826..fd8d2f7ef 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/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 - # chunked inference engine has frozen the anchor for the current chunk. - if state is not None and not self._hold_state: + # Always cache state for the paired AbsoluteActionsProcessorStep. + if state is not None: 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: - """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 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 { diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index aa30c34bb..8d0cdd653 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/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 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. -# -# 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. +# 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): @@ -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 - # are popped (see module docstring). Introspect the preprocessor for an - # enabled RelativeActionsProcessorStep, mirroring the RTC engine. + # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk + # (see module comment), 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`` - # recomputes and refreshes the anchor; clear any leftover hold defensively. - if self._relative_step is not None: - self._relative_step.set_hold(False) + # New episode: the next tick predicts a fresh chunk and re-anchors. + self._chunk_predicted = False + self._ever_predicted_chunk = False self._pred_stacks = {} self._pred_cursor = 0 - def _policy_will_recompute(self) -> bool: - """True if the next ``select_action`` will predict a fresh chunk (queue empty/absent). + 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``.""" + inner = self._policy.predict_action_chunk - 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 probe(*args, **kwargs): + self._chunk_predicted = True + self._ever_predicted_chunk = True + return inner(*args, **kwargs) + + self._policy.predict_action_chunk = probe 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 - # 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 + # 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. + anchor_before = self._relative_step.get_cached_state() if self._relative_step 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) + if self._visualize_predictions: + action, predictions = unpack_action_output( + self._policy.select_action(observation, return_intermediate_predictions=True) ) - observation = self._preprocessor(observation) - if self._visualize_predictions: - action, predictions = unpack_action_output( - self._policy.select_action(observation, return_intermediate_predictions=True) - ) - if predictions: - # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead. - self._pred_stacks = predictions - self._pred_cursor = 0 - else: - action = self._policy.select_action(observation) - action = self._postprocessor(action) - finally: - if self._relative_step is not None: - self._relative_step.set_hold(False) + if predictions: + # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead. + self._pred_stacks = predictions + self._pred_cursor = 0 + else: + 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() # 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 4dc97b9b5..1867183c4 100644 --- a/tests/policies/test_relative_actions.py +++ b/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 _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.""" +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) - - 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() + 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 450c2b53f..27956e3d6 100644 --- a/tests/test_rollout.py +++ b/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. - - 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. - """ + """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, @@ -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 - # it caches (or holds) the anchor, then pass the batch through unchanged. + # 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 @@ -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): - """Fake chunk policy: refills an ``_action_queue`` with ``chunk_rel`` when empty.""" +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) - 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 predict_action_chunk(_batch=None, **_kwargs): + state["predict_calls"] += 1 + return chunk_rel.unsqueeze(0) # [B=1, n, dim] 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) + 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 = 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(): - """A policy without ``_action_queue`` refreshes the anchor every tick.""" +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, 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():