admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 18:26:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

sync engine: robust per-chunk anchor via predict_action_chunk probe

Browse Source 
Replace the relative-action anchor mechanism from "adding relative actions
to sync engine" with one that detects fresh chunks through the public
predict_action_chunk contract instead of introspecting the policy's private
_action_queue. Drops set_hold/_hold_state in favour of a set_cached_state
snapshot/restore, keeps select_action on the hot path (so the prediction-
visualization path is unaffected), and refreshes the anchor every tick for
policies that never chunk.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Maxime Ellerbach
2026-07-17 15:13:55 +00:00
parent a185f9dde2
commit fd6aed87b7
4 changed files with 92 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/processor/relative_action_processor.py
+6 -17
View File
@@ -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 {
src/lerobot/rollout/inference/sync.py
+52 -63
View File
@@ -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
tests/policies/test_relative_actions.py
+3 -34
View File
@@ -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"}
tests/test_rollout.py
+31 -31
View File
@@ -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():