diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index d9f97f2c6..e1e65acb1 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/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]: diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py index ae8719b77..4ec47aca2 100644 --- a/src/lerobot/rollout/inference/rtc.py +++ b/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, diff --git a/tests/policies/rtc/test_rtc_relative_actions.py b/tests/policies/rtc/test_rtc_relative_actions.py index 14c115764..510f79c5d 100644 --- a/tests/policies/rtc/test_rtc_relative_actions.py +++ b/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)