From eb1c18d172a4d15bdb1d142c82261632406877ef Mon Sep 17 00:00:00 2001
From: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
Date: Wed, 22 Jul 2026 14:14:30 +0000
Subject: [PATCH] rtc action ordering

---
 src/lerobot/rollout/inference/factory.py |  1 +
 src/lerobot/rollout/inference/rtc.py     | 31 +++++++++++++++++++++---
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/lerobot/rollout/inference/factory.py b/src/lerobot/rollout/inference/factory.py
index 8ca46c916..cc6d13ee0 100644
--- a/src/lerobot/rollout/inference/factory.py
+++ b/src/lerobot/rollout/inference/factory.py
@@ -117,6 +117,7 @@ def create_inference_engine(
             robot_wrapper=robot_wrapper,
             rtc_config=config.rtc,
             dataset_features=dataset_features,
+            ordered_action_keys=ordered_action_keys,
             task=task,
             fps=fps,
             device=device,
diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py
index 33a3de21d..bfb620150 100644
--- a/src/lerobot/rollout/inference/rtc.py
+++ b/src/lerobot/rollout/inference/rtc.py
@@ -34,12 +34,13 @@ import torch
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.utils import prepare_observation_for_inference
+from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
 from lerobot.processor import (
     NormalizerProcessorStep,
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
 )
+from lerobot.utils.constants import ACTION
 from lerobot.utils.feature_utils import build_dataset_frame
 
 from ..robot_wrapper import ThreadSafeRobot
@@ -98,6 +99,7 @@ class RTCInferenceEngine(InferenceEngine):
         robot_wrapper: ThreadSafeRobot,
         rtc_config: RTCConfig,
         dataset_features: dict,
+        ordered_action_keys: list[str],
         task: str,
         fps: float,
         device: str | None,
@@ -119,9 +121,23 @@ class RTCInferenceEngine(InferenceEngine):
         # 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"
@@ -240,10 +256,19 @@ 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
-        return self._action_queue.get()
+        action = self._action_queue.get()
+        if action is None:
+            return None
+        action_dict = make_robot_action(action, self._obs_features)
+        return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
 
     def notify_observation(self, obs: dict) -> None:
         """Publish the latest observation for the RTC thread to consume."""