fix(rollout): rtc reanchor to non normalized state

Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>

src/lerobot/processor/relative_action_processor.py

@@ -142,6 +142,10 @@ class RelativeActionsProcessorStep(ProcessorStep):
         new_transition[TransitionKey.ACTION] = to_relative_actions(action, state, mask)
         return new_transition
 
+    def get_cached_state(self) -> torch.Tensor | None:
+        """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
+        return self._last_state
+
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,
@@ -182,7 +186,8 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
                 "but relative_step is None. Ensure relative_step is set when constructing the postprocessor."
             )
 
-        if self.relative_step._last_state is None:
+        cached_state = self.relative_step.get_cached_state()
+        if cached_state is None:
             raise RuntimeError(
                 "AbsoluteActionsProcessorStep requires state from RelativeActionsProcessorStep "
                 "but no state has been cached. Ensure the preprocessor runs before the postprocessor."
@@ -194,9 +199,7 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
             return new_transition
 
         mask = self.relative_step._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_absolute_actions(
-            action, self.relative_step._last_state, mask
-        )
+        new_transition[TransitionKey.ACTION] = to_absolute_actions(action, cached_state, mask)
         return new_transition
 
     def get_config(self) -> dict[str, Any]:

src/lerobot/rollout/inference/rtc.py

@@ -43,7 +43,6 @@ from lerobot.processor import (
     create_transition,
     to_relative_actions,
 )
-from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.feature_utils import build_dataset_frame
 
 from ..robot_wrapper import ThreadSafeRobot
@@ -318,13 +317,15 @@ class RTCInferenceEngine(InferenceEngine):
                         preprocessed = self._preprocessor(obs_batch)
 
                         if prev_actions is not None and self._relative_step is not None:
-                            state_tensor = preprocessed.get(OBS_STATE)
-                            if state_tensor is not None:
+                            # Rebase against the raw cached state so the leftover tail stays in
+                            # the training-time coordinate frame.
+                            raw_state = self._relative_step.get_cached_state()
+                            if raw_state is not None:
                                 prev_abs = queue.get_processed_left_over()
                                 if prev_abs is not None and prev_abs.numel() > 0:
                                     prev_actions = _reanchor_relative_rtc_prefix(
                                         prev_actions_absolute=prev_abs,
-                                        current_state=state_tensor,
+                                        current_state=raw_state,
                                         relative_step=self._relative_step,
                                         normalizer_step=self._normalizer_step,
                                         policy_device=policy_device,

tests/policies/rtc/test_rtc_relative_actions.py

@@ -22,13 +22,14 @@ from lerobot.configs.types import (
     PolicyFeature,
     RTCAttentionSchedule,
 )
-from lerobot.processor import TransitionKey, batch_to_transition
+from lerobot.processor import TransitionKey, batch_to_transition, create_transition
 from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
 from lerobot.processor.relative_action_processor import (
     AbsoluteActionsProcessorStep,
     RelativeActionsProcessorStep,
     to_relative_actions,
 )
+from lerobot.rollout.inference.rtc import _reanchor_relative_rtc_prefix
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 
@@ -218,7 +219,9 @@ class TestFullPipelineRelativeRTC:
 
     def test_roundtrip_with_identity_normalization(self):
         """Actions → relative → normalize → [model] → unnormalize → absolute should recover originals.
-        Using mean=0, std=1 normalization (identity)."""
+
+        Using mean=0, std=1 normalization (identity).
+        """
         relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
 
         state = torch.randn(1, ACTION_DIM)
@@ -400,6 +403,106 @@ class TestStateRebasingApproximation:
         assert error_excluded < 1e-6, f"Excluded joint should have zero error, got {error_excluded}"
 
 
+class TestRTCReanchoringWithStateNormalizer:
+    """RTC re-anchoring under non-identity OBS_STATE normalization."""
+
+    @staticmethod
+    def _build_normalizer_with_state_stats():
+        """Build a relative-action preprocessor with non-trivial OBS_STATE stats."""
+        features = {
+            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)),
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(ACTION_DIM,)),
+        }
+        norm_map = {
+            FeatureType.ACTION: NormalizationMode.MEAN_STD,
+            FeatureType.STATE: NormalizationMode.MEAN_STD,
+        }
+        stats = {
+            ACTION: {
+                "mean": torch.zeros(ACTION_DIM).numpy(),
+                "std": (0.5 * torch.ones(ACTION_DIM)).numpy(),
+            },
+            OBS_STATE: {
+                "mean": (5.0 * torch.ones(ACTION_DIM)).numpy(),
+                "std": (2.0 * torch.ones(ACTION_DIM)).numpy(),
+            },
+        }
+        relative_step = RelativeActionsProcessorStep(enabled=True)
+        normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+        return relative_step, normalizer
+
+    def test_reanchor_with_raw_state_matches_normalize_of_absolute_minus_state(self):
+        """Reanchoring with the raw cached state yields ``normalize(prev_actions_absolute - raw_state)``."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+        relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
+
+        prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
+
+        result = _reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=relative_step.get_cached_state(),
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+
+        expected_relative = to_relative_actions(prev_actions_absolute, raw_state, [True] * ACTION_DIM)
+        expected = normalizer(create_transition(action=expected_relative))[TransitionKey.ACTION]
+        torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5)
+
+    def test_reanchor_with_normalized_state_produces_wrong_result(self):
+        """Reanchoring with raw vs. normalized state produces meaningfully different outputs."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+        relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
+
+        normalized_obs = normalizer(batch_to_transition({OBS_STATE: raw_state.clone()}))
+        normalized_state = normalized_obs[TransitionKey.OBSERVATION][OBS_STATE]
+        assert not torch.allclose(normalized_state, raw_state)
+
+        prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
+
+        result_raw = _reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=raw_state,
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+        result_normalized = _reanchor_relative_rtc_prefix(
+            prev_actions_absolute=prev_actions_absolute,
+            current_state=normalized_state,
+            relative_step=relative_step,
+            normalizer_step=normalizer,
+            policy_device="cpu",
+        )
+
+        max_abs_diff = (result_raw - result_normalized).abs().max()
+        assert max_abs_diff > 0.5, (
+            f"Raw and normalized state produced near-identical outputs (max diff {max_abs_diff:.4f}); "
+            "OBS_STATE stats are too close to identity to be sensitive."
+        )
+
+    def test_engine_pipeline_cached_state_is_raw_after_full_preprocess(self):
+        """``get_cached_state()`` returns raw OBS_STATE after the full preprocessor pipeline runs."""
+        relative_step, normalizer = self._build_normalizer_with_state_stats()
+
+        raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
+
+        transition = batch_to_transition({OBS_STATE: raw_state.clone()})
+        transition = relative_step(transition)
+        preprocessed = normalizer(transition)
+
+        cached = relative_step.get_cached_state()
+        torch.testing.assert_close(cached, raw_state, atol=1e-6, rtol=1e-6)
+
+        post_normalize_state = preprocessed[TransitionKey.OBSERVATION][OBS_STATE]
+        assert not torch.allclose(cached, post_normalize_state, atol=1e-3)
+
+
 def _detect_relative_actions(preprocessor) -> bool:
     """Mirror of the helper in lerobot-rollout for testing without importing it."""
     return any(isinstance(step, RelativeActionsProcessorStep) and step.enabled for step in preprocessor.steps)