cleaning up to keep changes minimal

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
+                # take the min of both to avoid discarding actions that were not
-                # actually consumed during inference (`indexes_diff`) disagree. This happens
+                # actually consumed during inference, which would cause a jump in the queue
-                # 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.
                 resolved = min(real_delay, indexes_diff)
                 logger.warning(
                     "Indexes diff != real delay (indexes_diff=%d, real_delay=%d); "

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.
+    """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.
-    """
     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``).
+        """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``.
-        """
         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])
 

tests/test_rollout.py

@@ -252,93 +252,10 @@ def test_create_inference_engine_sync():
 
 
 # ---------------------------------------------------------------------------
-# Observation feature-spec consistency (sync vs RTC)
+# RTC action-key ordering and prefix padding
-#
-# ``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.
 # ---------------------------------------------------------------------------
 
 
-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