diff --git a/src/lerobot/policies/rtc/action_queue.py b/src/lerobot/policies/rtc/action_queue.py index 4c6c58297..a8c6d1f6f 100644 --- a/src/lerobot/policies/rtc/action_queue.py +++ b/src/lerobot/policies/rtc/action_queue.py @@ -236,13 +236,8 @@ class ActionQueue: if action_index_before_inference is not None: indexes_diff = max(0, self.last_index - action_index_before_inference) if indexes_diff != real_delay: - # The latency estimate (`real_delay`) and the number of actions the robot - # actually consumed during inference (`indexes_diff`) disagree. This happens - # when the queue starved (robot idle) or on the first chunk (nothing consumed - # yet). Discarding `real_delay` here would drop actions the arm never executed - # and splice the queue `real_delay` steps ahead of the physical pose — a hard - # jump/slam, worst on slow policies where `real_delay` is large. Never discard - # more than was actually consumed. + # take the min of both to avoid discarding actions that were not + # actually consumed during inference, which would cause a jump in the queue resolved = min(real_delay, indexes_diff) logger.warning( "Indexes diff != real delay (indexes_diff=%d, real_delay=%d); " diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py index c16744c05..1421082b5 100644 --- a/src/lerobot/rollout/inference/rtc.py +++ b/src/lerobot/rollout/inference/rtc.py @@ -40,7 +40,6 @@ from lerobot.processor import ( PolicyProcessorPipeline, RelativeActionsProcessorStep, ) -from lerobot.utils.constants import ACTION from lerobot.utils.feature_utils import build_dataset_frame from ..robot_wrapper import ThreadSafeRobot @@ -64,17 +63,7 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0 def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor: - """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference. - - Padding repeats the last real action ("hold") rather than filling with zeros. The RTC - guidance pulls the new chunk toward this prefix at the padded indices (they fall inside - the weighted region when the real leftover is shorter than ``target_steps``). A zero in - the model's normalized action space decodes to the dataset *mean* action — a nonzero - offset that yanks the spliced action toward a mean/neutral pose for one step, producing - an intermittent seam (e.g. 95 -> 103 -> 95). Holding the last real action keeps the - padded targets continuous with the prefix, so no fake target enters the guided region. - The fixed output length is preserved so ``torch.compile`` policies keep stable shapes. - """ + """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference.""" if prev_actions.ndim != 2: raise ValueError(f"Expected 2D [T, A] tensor, got shape={tuple(prev_actions.shape)}") steps, _ = prev_actions.shape @@ -84,6 +73,7 @@ def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int return prev_actions[:target_steps] if steps == 0: raise ValueError("Cannot pad an empty prefix: no last action to hold.") + # repeat the last action to fill the remaining steps to avoid a sudden jump hold = prev_actions[-1:].expand(target_steps - steps, -1) return torch.cat([prev_actions, hold], dim=0) @@ -124,31 +114,8 @@ class RTCInferenceEngine(InferenceEngine): self._postprocessor = postprocessor self._robot = robot_wrapper self._rtc_config = rtc_config - # Build observations with the SAME feature spec sync uses (post - # `robot_observation_processor`), not the raw-hardware spec. `build_dataset_frame` - # orders `observation.state` by this spec's `names`; using the raw-hardware order - # here (as before) desynced the state vector from sync whenever the observation - # processor reorders/renames state keys, corrupting both normalization and the - # relative-action anchor. The `prefix="observation"` filter ignores the action - # entries in the combined dict. self._obs_features = dataset_features - # The model emits actions in `dataset_features[ACTION]` order (the order it was - # trained on); the robot expects them in `ordered_action_keys` order. Sync remaps - # by NAME via `make_robot_action` + reindex (sync.py) before returning; RTC must do - # the SAME, otherwise the engine-agnostic strategy (`send_next_action`) maps the raw - # model-order tensor onto `ordered_action_keys` positionally and mis-assigns joints - # whenever the two orders differ — a per-joint permutation that drives the arm wrong. self._ordered_action_keys = ordered_action_keys - state_ft = dataset_features.get("observation.state") - if state_ft is not None: - logger.info("RTC observation.state layout: %s", state_ft.get("names")) - action_ft = dataset_features.get(ACTION) - if action_ft is not None: - logger.info( - "RTC action layout: model/dataset=%s -> robot=%s", - action_ft.get("names"), - self._ordered_action_keys, - ) self._task = task self._fps = fps self._device = device or "cpu" @@ -267,17 +234,14 @@ class RTCInferenceEngine(InferenceEngine): # ------------------------------------------------------------------ def get_action(self, obs_frame: dict | None) -> torch.Tensor | None: - """Pop the next action from the RTC queue (ignores ``obs_frame``). - - The queued action is in the model's ``dataset_features[ACTION]`` order; remap it - by NAME into ``ordered_action_keys`` order before returning, so the engine-agnostic - strategy maps values onto the correct joints. Mirrors ``SyncInferenceEngine.get_action``. - """ + """Pop the next action from the RTC queue (ignores ``obs_frame``).""" if self._action_queue is None: return None action = self._action_queue.get() if action is None: return None + + # properly reorder the action dict to match the robot's expected order action_dict = make_robot_action(action, self._obs_features) return torch.tensor([action_dict[k] for k in self._ordered_action_keys]) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 5a09b783e..b1fa2a212 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -252,93 +252,10 @@ def test_create_inference_engine_sync(): # --------------------------------------------------------------------------- -# Observation feature-spec consistency (sync vs RTC) -# -# ``build_dataset_frame`` orders the ``observation.state`` vector by the ``names`` -# list of the feature spec it is handed. Sync uses ``dataset_features`` (joint -# layout AFTER ``robot_observation_processor``); RTC must use the SAME spec, not the -# raw-hardware layout, or its state vector desyncs from sync — corrupting both the -# normalizer's per-joint stats and the relative-action anchor, which sends the arm -# to systematically wrong absolute targets. +# RTC action-key ordering and prefix padding # --------------------------------------------------------------------------- -def test_build_dataset_frame_state_layout_depends_on_feature_spec(): - """A reordered feature spec yields a reordered ``observation.state`` vector.""" - import numpy as np - - from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features - - # Raw hardware joint order. - obs_hw = {"shoulder.pos": float, "elbow.pos": float, "wrist.pos": float, "gripper.pos": float} - hw_features = hw_to_dataset_features(obs_hw, "observation") - - # What a ``robot_observation_processor`` that reorders state keys would produce - # (e.g. gripper moved to the front). Same keys, different order. - dataset_features = { - "observation.state": { - "dtype": "float32", - "shape": (4,), - "names": ["gripper.pos", "shoulder.pos", "elbow.pos", "wrist.pos"], - } - } - - values = {"shoulder.pos": 10.0, "elbow.pos": 20.0, "wrist.pos": 30.0, "gripper.pos": 40.0} - - hw_state = build_dataset_frame(hw_features, values, prefix="observation")["observation.state"] - ds_state = build_dataset_frame(dataset_features, values, prefix="observation")["observation.state"] - - np.testing.assert_array_equal(hw_state, [10.0, 20.0, 30.0, 40.0]) - np.testing.assert_array_equal(ds_state, [40.0, 10.0, 20.0, 30.0]) - # The two layouts genuinely differ: this is the sync/RTC divergence the fix removes. - assert not np.array_equal(hw_state, ds_state) - - -def test_create_inference_engine_rtc_uses_dataset_features(): - """The RTC engine must build observations from ``dataset_features`` (sync's spec), - not the raw-hardware ``hw_features`` — so its ``observation.state`` matches sync.""" - from lerobot.rollout import RTCInferenceConfig, RTCInferenceEngine, create_inference_engine - - dataset_features = { - "observation.state": { - "dtype": "float32", - "shape": (4,), - "names": ["gripper.pos", "shoulder.pos", "elbow.pos", "wrist.pos"], - } - } - # Deliberately different order to prove the engine ignores it. - hw_features = { - "observation.state": { - "dtype": "float32", - "shape": (4,), - "names": ["shoulder.pos", "elbow.pos", "wrist.pos", "gripper.pos"], - } - } - - engine = create_inference_engine( - RTCInferenceConfig(), - policy=MagicMock(), - # No relative/normalizer steps => __init__ introspection stays trivial. - preprocessor=MagicMock(steps=[]), - postprocessor=MagicMock(steps=[]), - robot_wrapper=MagicMock(robot_type="mock"), - hw_features=hw_features, - dataset_features=dataset_features, - ordered_action_keys=["k"], - task="test", - fps=30.0, - device="cpu", - ) - assert isinstance(engine, RTCInferenceEngine) - assert engine._obs_features is dataset_features - assert engine._obs_features["observation.state"]["names"] == [ - "gripper.pos", - "shoulder.pos", - "elbow.pos", - "wrist.pos", - ] - - def test_rtc_get_action_remaps_model_order_to_ordered_action_keys(): """RTC must remap the model-order action vector to ``ordered_action_keys`` by NAME before returning — matching sync. Otherwise the strategy maps model outputs onto the