Merge remote-tracking branch 'origin/main' into user/khalil-meftah/2026-02-16-rl-stack-refactor

tests/policies/groot/test_groot_lerobot.py

@@ -31,7 +31,7 @@ from lerobot.policies.groot.processor_groot import make_groot_pre_post_processor
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 from lerobot.utils.device_utils import auto_select_torch_device
-from tests.utils import require_cuda  # noqa: E402
+from tests.utils import require_cuda
 
 pytest.importorskip("transformers")
 

tests/policies/hilserl/test_modeling_classifier.py

@@ -14,13 +14,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 from lerobot.policies.sac.reward_model.modeling_classifier import ClassifierOutput
 from lerobot.utils.constants import OBS_IMAGE, REWARD
-from tests.utils import require_package
+from tests.utils import skip_if_package_missing
 
 
 def test_classifier_output():
@@ -36,7 +37,10 @@ def test_classifier_output():
     )
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
+@pytest.mark.skip(
+    reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
+)
 def test_binary_classifier_with_default_params():
     from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
@@ -77,7 +81,10 @@ def test_binary_classifier_with_default_params():
     assert not torch.isnan(output.hidden_states).any(), "Tensor contains NaN values"
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
+@pytest.mark.skip(
+    reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
+)
 def test_multiclass_classifier():
     from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
@@ -116,7 +123,10 @@ def test_multiclass_classifier():
     assert not torch.isnan(output.hidden_states).any(), "Tensor contains NaN values"
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
+@pytest.mark.skip(
+    reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
+)
 def test_default_device():
     from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
@@ -128,7 +138,10 @@ def test_default_device():
         assert p.device == torch.device("cpu")
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
+@pytest.mark.skip(
+    reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
+)
 def test_explicit_device_setup():
     from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 

tests/policies/multi_task_dit/test_multi_task_dit.py

@@ -0,0 +1,621 @@
+#!/usr/bin/env python
+
+# Copyright 2025 Bryson Jones and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# ruff: noqa: E402
+
+"""Test script for Multi-Task DiT policy.
+
+To run tests locally:
+    python -m pytest tests/policies/multi_task_dit/test_multi_task_dit.py -v
+"""
+
+import os
+
+import pytest
+import torch
+from torch import Tensor
+
+pytest.importorskip("transformers")
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
+    reason="This test requires local transformers installation and is not meant for CI",
+)
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
+from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
+from lerobot.policies.multi_task_dit.processor_multi_task_dit import (
+    make_multi_task_dit_pre_post_processors,
+)
+from lerobot.utils.constants import (
+    ACTION,
+    OBS_IMAGES,
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    OBS_STATE,
+)
+from lerobot.utils.random_utils import seeded_context, set_seed
+
+
+@pytest.fixture(autouse=True)
+def set_random_seed():
+    seed = 17
+    set_seed(seed)
+
+
+def create_train_batch(
+    batch_size: int = 2,
+    n_obs_steps: int = 2,
+    horizon: int = 16,
+    state_dim: int = 10,
+    action_dim: int = 10,
+    height: int = 224,
+    width: int = 224,
+) -> dict[str, Tensor]:
+    """Create a training batch with visual input and text."""
+    return {
+        "observation.state": torch.randn(batch_size, n_obs_steps, state_dim),
+        f"{OBS_IMAGES}.laptop": torch.rand(batch_size, n_obs_steps, 3, height, width),
+        ACTION: torch.randn(batch_size, horizon, action_dim),
+        "task": ["pick up the cube"] * batch_size,
+    }
+
+
+def create_observation_batch(
+    batch_size: int = 2, state_dim: int = 10, height: int = 224, width: int = 224
+) -> dict:
+    """Create observation batch for inference for a single timestep."""
+    return {
+        "observation.state": torch.randn(batch_size, state_dim),
+        f"{OBS_IMAGES}.laptop": torch.rand(batch_size, 3, height, width),
+        "task": ["pick up the red cube"] * batch_size,
+    }
+
+
+def create_config(
+    state_dim: int = 10,
+    action_dim: int = 10,
+    n_obs_steps: int = 2,
+    horizon: int = 16,
+    n_action_steps: int = 8,
+    with_visual: bool = True,
+    height: int = 224,
+    width: int = 224,
+) -> MultiTaskDiTConfig:
+    """Create a MultiTaskDiT config for testing.
+
+    Args:
+        state_dim: Dimension of state observations
+        action_dim: Dimension of actions
+        n_obs_steps: Number of observation steps
+        horizon: Action prediction horizon
+        n_action_steps: Number of action steps to execute
+        with_visual: Whether to include visual input (default: True)
+        height: Image height (only used if with_visual=True)
+        width: Image width (only used if with_visual=True)
+    """
+    input_features = {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))}
+
+    if with_visual:
+        input_features[f"{OBS_IMAGES}.laptop"] = PolicyFeature(
+            type=FeatureType.VISUAL, shape=(3, height, width)
+        )
+
+    config = MultiTaskDiTConfig(
+        input_features=input_features,
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))},
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+        # Use smaller model for faster tests
+        hidden_dim=128,
+        num_layers=2,
+        num_heads=4,
+    )
+
+    config.validate_features()
+    return config
+
+
+@pytest.mark.parametrize("batch_size,state_dim,action_dim", [(2, 10, 10), (1, 6, 6)])
+def test_multi_task_dit_policy_forward(batch_size: int, state_dim: int, action_dim: int):
+    """Test forward pass (training mode)."""
+    n_obs_steps = 2
+    horizon = 16
+    n_action_steps = 8
+
+    config = create_config(
+        state_dim=state_dim,
+        action_dim=action_dim,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+    )
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    policy.train()
+
+    # Use preprocessor to handle tokenization
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, _ = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+
+    batch = create_train_batch(
+        batch_size=batch_size,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        state_dim=state_dim,
+        action_dim=action_dim,
+    )
+
+    # Process batch through preprocessor to tokenize task text
+    processed_batch = preprocessor(batch)
+
+    # Test forward pass
+    loss, _ = policy.forward(processed_batch)
+    assert loss is not None
+    assert loss.item() is not None
+    assert loss.shape == ()
+
+    # Test backward pass
+    loss.backward()
+
+
+def test_multi_task_dit_pre_post_processors():
+    """Test pre and post processors for Multi-Task DiT policy."""
+    state_dim = 10
+    action_dim = 8
+    n_obs_steps = 2
+    horizon = 16
+
+    config = create_config(
+        state_dim=state_dim,
+        action_dim=action_dim,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=8,
+    )
+    config.device = "cpu"
+
+    # Set normalization mode to match the stats we're providing
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.MEAN_STD,  # Use MEAN_STD since we provide mean/std stats
+        "ACTION": NormalizationMode.MIN_MAX,
+    }
+
+    # Create dataset stats for normalization
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.zeros(state_dim),
+            "std": torch.ones(state_dim),
+        },
+        "action": {
+            "min": torch.full((action_dim,), -1.0),
+            "max": torch.ones(action_dim),
+        },
+    }
+
+    # Create processors
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(
+        config=config, dataset_stats=dataset_stats
+    )
+
+    # Test preprocessor with sample data
+    batch = {
+        "observation.state": torch.randn(state_dim),
+        f"{OBS_IMAGES}.laptop": torch.rand(3, 224, 224),
+        ACTION: torch.randn(action_dim),
+        "task": "pick up the cube",
+    }
+
+    processed_batch = preprocessor(batch)
+
+    # Check that data is batched
+    assert processed_batch["observation.state"].shape == (1, state_dim)
+    assert processed_batch[f"{OBS_IMAGES}.laptop"].shape == (1, 3, 224, 224)
+    assert processed_batch[ACTION].shape == (1, action_dim)
+    # Check that task text was tokenized
+    assert OBS_LANGUAGE_TOKENS in processed_batch
+    assert OBS_LANGUAGE_ATTENTION_MASK in processed_batch
+    assert processed_batch[OBS_LANGUAGE_TOKENS].shape[0] == 1  # batch dimension
+    assert processed_batch[OBS_LANGUAGE_ATTENTION_MASK].shape[0] == 1  # batch dimension
+
+    # Check that data is on correct device
+    assert processed_batch["observation.state"].device.type == "cpu"
+    assert processed_batch[f"{OBS_IMAGES}.laptop"].device.type == "cpu"
+    assert processed_batch[ACTION].device.type == "cpu"
+
+    # Test postprocessor with sample action (PolicyAction is just a torch.Tensor)
+    action = torch.randn(1, action_dim)
+    processed_action = postprocessor(action)
+
+    # Check that action is unnormalized and on CPU
+    assert processed_action.shape == (1, action_dim)
+    assert processed_action.device.type == "cpu"
+
+
+def test_multi_task_dit_pre_post_processors_normalization():
+    """Test that normalization and unnormalization work correctly with simple sanity check numbers."""
+    state_dim = 3
+    action_dim = 2
+
+    config = create_config(
+        state_dim=state_dim,
+        action_dim=action_dim,
+        n_obs_steps=2,
+        horizon=16,
+        n_action_steps=8,
+    )
+    config.device = "cpu"
+
+    # Set normalization mode to match the stats we're providing
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.MEAN_STD,  # Use MEAN_STD since we provide mean/std stats
+        "ACTION": NormalizationMode.MIN_MAX,
+    }
+
+    # Use simple stats that will actually transform the values
+    dataset_stats = {
+        "observation.state": {
+            "mean": torch.full((state_dim,), 5.0),
+            "std": torch.full((state_dim,), 2.0),
+        },
+        "action": {
+            "min": torch.zeros(action_dim),
+            "max": torch.full((action_dim,), 2.0),
+        },
+    }
+
+    # Create processors
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(
+        config=config, dataset_stats=dataset_stats
+    )
+
+    # Use simple input values
+    input_state = torch.tensor([7.0, 5.0, 3.0])  # Will normalize to [1.0, 0.0, -1.0]
+    input_action = torch.tensor([1.0, 2.0])  # Will normalize to [0.0, 1.0]
+
+    batch = {
+        "observation.state": input_state,
+        f"{OBS_IMAGES}.laptop": torch.rand(3, 224, 224),
+        ACTION: input_action,
+        "task": "test task",
+    }
+
+    # Process through preprocessor
+    processed_batch = preprocessor(batch)
+
+    # State normalization: (x - mean) / std
+    expected_normalized_state = torch.tensor([1.0, 0.0, -1.0])
+    assert torch.allclose(processed_batch["observation.state"][0], expected_normalized_state, atol=1e-5)
+
+    # Action normalization: (x - min) / (max - min) * 2 - 1
+    expected_normalized_action = torch.tensor([0.0, 1.0])
+    assert torch.allclose(processed_batch[ACTION][0], expected_normalized_action, atol=1e-5)
+
+    # Test unnormalization: should recover original values
+    normalized_action_tensor = processed_batch[ACTION][0:1]  # Keep batch dimension
+    unnormalized_action = postprocessor(normalized_action_tensor)
+
+    # Should recover original action values
+    assert torch.allclose(unnormalized_action[0], input_action, atol=1e-4)
+
+
+@pytest.mark.parametrize("batch_size,state_dim,action_dim", [(2, 10, 10), (1, 6, 6)])
+def test_multi_task_dit_policy_select_action(batch_size: int, state_dim: int, action_dim: int):
+    """Test select_action (inference mode)."""
+    n_obs_steps = 2
+    horizon = 16
+    n_action_steps = 8
+
+    config = create_config(
+        state_dim=state_dim,
+        action_dim=action_dim,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+    )
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    policy.eval()
+    policy.reset()  # Reset queues before inference
+
+    # Create processors - use IDENTITY normalization when no stats provided
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+
+    with torch.no_grad():
+        observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
+        # Process observation through preprocessor
+        processed_obs = preprocessor(observation_batch)
+        selected_action = policy.select_action(processed_obs)
+        # Process action through postprocessor (PolicyAction is just a torch.Tensor)
+        processed_action = postprocessor(selected_action)
+        assert processed_action.shape == (batch_size, action_dim)
+
+
+def test_multi_task_dit_policy_diffusion_objective():
+    """Test policy with diffusion objective."""
+    batch_size = 2
+    state_dim = 10
+    action_dim = 10
+    n_obs_steps = 2
+    horizon = 16
+    n_action_steps = 8
+
+    input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
+        f"{OBS_IMAGES}.laptop": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+
+    config = MultiTaskDiTConfig(
+        input_features=input_features,
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))},
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+        # Use diffusion objective
+        objective="diffusion",
+        noise_scheduler_type="DDPM",
+        num_train_timesteps=100,
+        num_inference_steps=10,
+        # Smaller model for tests
+        hidden_dim=128,
+        num_layers=2,
+        num_heads=4,
+    )
+    config.validate_features()
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    policy.train()
+
+    # Use preprocessor to handle tokenization
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, _ = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+
+    batch = create_train_batch(
+        batch_size=batch_size,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        state_dim=state_dim,
+        action_dim=action_dim,
+    )
+
+    # Process batch through preprocessor to tokenize task text
+    processed_batch = preprocessor(batch)
+
+    # Test forward pass
+    loss, _ = policy.forward(processed_batch)
+    assert loss is not None
+    assert loss.item() is not None
+
+    # Test inference
+    policy.eval()
+    # Use IDENTITY normalization when no stats provided
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+    with torch.no_grad():
+        observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
+        # Process observation through preprocessor
+        processed_obs = preprocessor(observation_batch)
+        selected_action = policy.select_action(processed_obs)
+        # Process action through postprocessor (PolicyAction is just a torch.Tensor)
+        processed_action = postprocessor(selected_action)
+        assert processed_action.shape == (batch_size, action_dim)
+
+
+def test_multi_task_dit_policy_flow_matching_objective():
+    """Test policy with flow matching objective."""
+    batch_size = 2
+    state_dim = 10
+    action_dim = 10
+    n_obs_steps = 2
+    horizon = 16
+    n_action_steps = 8
+
+    input_features = {
+        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
+        f"{OBS_IMAGES}.laptop": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
+    }
+
+    config = MultiTaskDiTConfig(
+        input_features=input_features,
+        output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))},
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+        # Use flow matching objective
+        objective="flow_matching",
+        sigma_min=0.0,
+        num_integration_steps=10,  # Fewer steps for faster tests
+        integration_method="euler",
+        # Smaller model for tests
+        hidden_dim=128,
+        num_layers=2,
+        num_heads=4,
+    )
+    config.validate_features()
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    policy.train()
+
+    # Use preprocessor to handle tokenization
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, _ = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+
+    batch = create_train_batch(
+        batch_size=batch_size,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        state_dim=state_dim,
+        action_dim=action_dim,
+    )
+
+    # Process batch through preprocessor to tokenize task text
+    processed_batch = preprocessor(batch)
+
+    # Test forward pass
+    loss, _ = policy.forward(processed_batch)
+    assert loss is not None
+    assert loss.item() is not None
+
+    # Test inference
+    policy.eval()
+    # Use IDENTITY normalization when no stats provided
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+    with torch.no_grad():
+        observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
+        # Process observation through preprocessor
+        processed_obs = preprocessor(observation_batch)
+        selected_action = policy.select_action(processed_obs)
+        # Process action through postprocessor (PolicyAction is just a torch.Tensor)
+        processed_action = postprocessor(selected_action)
+        assert processed_action.shape == (batch_size, action_dim)
+
+
+def test_multi_task_dit_policy_save_and_load(tmp_path):
+    """Test that the policy can be saved and loaded correctly."""
+    root = tmp_path / "test_multi_task_dit_save_and_load"
+
+    state_dim = 10
+    action_dim = 10
+    batch_size = 2
+    n_obs_steps = 2
+    horizon = 16
+    n_action_steps = 8
+
+    config = create_config(
+        state_dim=state_dim,
+        action_dim=action_dim,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        n_action_steps=n_action_steps,
+    )
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    policy.eval()
+
+    policy.save_pretrained(root)
+    loaded_policy = MultiTaskDiTPolicy.from_pretrained(root, config=config)
+    loaded_policy.to(config.device)
+    loaded_policy.eval()
+
+    batch = create_train_batch(
+        batch_size=batch_size,
+        n_obs_steps=n_obs_steps,
+        horizon=horizon,
+        state_dim=state_dim,
+        action_dim=action_dim,
+    )
+
+    # Use preprocessor to handle tokenization
+    config.normalization_mapping = {
+        "VISUAL": NormalizationMode.IDENTITY,
+        "STATE": NormalizationMode.IDENTITY,
+        "ACTION": NormalizationMode.IDENTITY,
+    }
+    preprocessor, postprocessor = make_multi_task_dit_pre_post_processors(config=config, dataset_stats=None)
+
+    with torch.no_grad():
+        with seeded_context(12):
+            # Process batch through preprocessor
+            processed_batch = preprocessor(batch)
+            # Collect policy values before saving
+            loss, _ = policy.forward(processed_batch)
+
+            observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
+            # Process observation through preprocessor
+            processed_obs = preprocessor(observation_batch)
+            actions = policy.select_action(processed_obs)
+
+        with seeded_context(12):
+            # Process batch through preprocessor
+            processed_batch = preprocessor(batch)
+            # Collect policy values after loading
+            loaded_loss, _ = loaded_policy.forward(processed_batch)
+
+            loaded_observation_batch = create_observation_batch(batch_size=batch_size, state_dim=state_dim)
+            processed_obs = preprocessor(loaded_observation_batch)
+            loaded_actions = loaded_policy.select_action(processed_obs)
+
+        # Compare state dicts
+        assert policy.state_dict().keys() == loaded_policy.state_dict().keys()
+        for k in policy.state_dict():
+            assert torch.allclose(policy.state_dict()[k], loaded_policy.state_dict()[k], atol=1e-6)
+
+        # Compare values before and after saving and loading
+        assert torch.allclose(loss, loaded_loss)
+        assert torch.allclose(actions, loaded_actions)
+
+
+def test_multi_task_dit_policy_get_optim_params():
+    """Test that the policy returns correct optimizer parameter groups."""
+    config = create_config(
+        state_dim=10,
+        action_dim=10,
+        n_obs_steps=2,
+        horizon=16,
+        n_action_steps=8,
+    )
+
+    policy = MultiTaskDiTPolicy(config=config)
+    policy.to(config.device)
+    param_groups = policy.get_optim_params()
+
+    # Should have 2 parameter groups: non-vision and vision encoder
+    assert len(param_groups) == 2
+
+    # First group is non-vision params (no lr specified, will use default)
+    assert "params" in param_groups[0]
+    assert len(param_groups[0]["params"]) > 0
+
+    # Second group is vision encoder params with different lr
+    assert "params" in param_groups[1]
+    assert "lr" in param_groups[1]
+    expected_lr = config.optimizer_lr * config.vision_encoder_lr_multiplier
+    assert param_groups[1]["lr"] == expected_lr

tests/policies/rtc/test_action_interpolator.py

@@ -0,0 +1,559 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ActionInterpolator and its interaction with ActionQueue (RTC)."""
+
+import pytest
+import torch
+
+from lerobot.policies.rtc.action_interpolator import ActionInterpolator
+from lerobot.policies.rtc.action_queue import ActionQueue
+from lerobot.policies.rtc.configuration_rtc import RTCConfig
+
+# ====================== Fixtures ======================
+
+
+@pytest.fixture
+def interp2():
+    """Create an ActionInterpolator with multiplier=2."""
+    return ActionInterpolator(multiplier=2)
+
+
+@pytest.fixture
+def interp3():
+    """Create an ActionInterpolator with multiplier=3."""
+    return ActionInterpolator(multiplier=3)
+
+
+# ====================== Initialization Tests ======================
+
+
+def test_interpolator_multiplier_1_no_interpolation():
+    """Test multiplier=1 creates a disabled interpolator."""
+    interp = ActionInterpolator(multiplier=1)
+    assert interp.multiplier == 1
+    assert not interp.enabled
+
+
+def test_interpolator_multiplier_2_enabled():
+    """Test multiplier=2 creates an enabled interpolator."""
+    interp = ActionInterpolator(multiplier=2)
+    assert interp.multiplier == 2
+    assert interp.enabled
+
+
+def test_interpolator_multiplier_0_raises():
+    """Test multiplier=0 raises ValueError."""
+    with pytest.raises(ValueError, match="multiplier must be >= 1"):
+        ActionInterpolator(multiplier=0)
+
+
+def test_interpolator_negative_multiplier_raises():
+    """Test negative multiplier raises ValueError."""
+    with pytest.raises(ValueError, match="multiplier must be >= 1"):
+        ActionInterpolator(multiplier=-1)
+
+
+def test_interpolator_default_multiplier_is_1():
+    """Test default multiplier is 1 (disabled)."""
+    interp = ActionInterpolator()
+    assert interp.multiplier == 1
+    assert not interp.enabled
+
+
+# ====================== needs_new_action Tests ======================
+
+
+def test_needs_new_action_true_initially(interp2):
+    """Test needs_new_action() returns True before any action is added."""
+    assert interp2.needs_new_action()
+
+
+def test_needs_new_action_false_after_add(interp2):
+    """Test needs_new_action() returns False right after add()."""
+    interp2.add(torch.tensor([1.0, 2.0]))
+    assert not interp2.needs_new_action()
+
+
+def test_needs_new_action_true_after_buffer_exhausted(interp2):
+    """Test needs_new_action() returns True after consuming all buffered actions."""
+    interp2.add(torch.tensor([1.0, 2.0]))
+    interp2.get()
+    assert interp2.needs_new_action()
+
+
+def test_needs_new_action_true_after_all_interpolated_consumed(interp2):
+    """Test needs_new_action() tracks interpolated sub-steps correctly."""
+    interp2.add(torch.tensor([0.0, 0.0]))
+    interp2.get()
+    assert interp2.needs_new_action()
+
+    interp2.add(torch.tensor([2.0, 4.0]))
+    interp2.get()
+    assert not interp2.needs_new_action()
+    interp2.get()
+    assert interp2.needs_new_action()
+
+
+# ====================== Passthrough Tests (multiplier=1) ======================
+
+
+def test_passthrough_single_action_returned_as_is():
+    """Test multiplier=1 returns the action unchanged."""
+    interp = ActionInterpolator(multiplier=1)
+    action = torch.tensor([3.0, 5.0])
+    interp.add(action)
+
+    result = interp.get()
+    assert result is not None
+    torch.testing.assert_close(result, action)
+
+
+def test_passthrough_none_after_single_get():
+    """Test multiplier=1 returns None after consuming the single action."""
+    interp = ActionInterpolator(multiplier=1)
+    interp.add(torch.tensor([1.0]))
+    interp.get()
+    assert interp.get() is None
+
+
+def test_passthrough_sequential_actions():
+    """Test multiplier=1 passes through consecutive actions one at a time."""
+    interp = ActionInterpolator(multiplier=1)
+    for val in [1.0, 2.0, 3.0]:
+        action = torch.tensor([val])
+        interp.add(action)
+        result = interp.get()
+        torch.testing.assert_close(result, action)
+        assert interp.get() is None
+
+
+# ====================== Interpolation Tests (multiplier=2) ======================
+
+
+def test_interpolation_2x_first_action_no_interpolation(interp2):
+    """Test first action has no previous, so buffer is just [action]."""
+    interp2.add(torch.tensor([0.0, 0.0]))
+    result = interp2.get()
+    torch.testing.assert_close(result, torch.tensor([0.0, 0.0]))
+    assert interp2.get() is None
+
+
+def test_interpolation_2x_second_action_produces_two_steps(interp2):
+    """Test second action produces 2 interpolated sub-steps."""
+    interp2.add(torch.tensor([0.0, 0.0]))
+    interp2.get()
+
+    interp2.add(torch.tensor([2.0, 4.0]))
+    step1 = interp2.get()
+    step2 = interp2.get()
+
+    torch.testing.assert_close(step1, torch.tensor([1.0, 2.0]))
+    torch.testing.assert_close(step2, torch.tensor([2.0, 4.0]))
+    assert interp2.get() is None
+
+
+def test_interpolation_2x_three_consecutive_actions(interp2):
+    """Test interpolation across three consecutive actions."""
+    a0 = torch.tensor([0.0])
+    a1 = torch.tensor([4.0])
+    a2 = torch.tensor([10.0])
+
+    interp2.add(a0)
+    torch.testing.assert_close(interp2.get(), a0)
+
+    interp2.add(a1)
+    torch.testing.assert_close(interp2.get(), torch.tensor([2.0]))
+    torch.testing.assert_close(interp2.get(), torch.tensor([4.0]))
+
+    interp2.add(a2)
+    torch.testing.assert_close(interp2.get(), torch.tensor([7.0]))
+    torch.testing.assert_close(interp2.get(), torch.tensor([10.0]))
+
+
+# ====================== Interpolation Tests (multiplier=3) ======================
+
+
+def test_interpolation_3x_produces_three_steps(interp3):
+    """Test multiplier=3 produces 3 interpolated sub-steps."""
+    interp3.add(torch.tensor([0.0, 0.0]))
+    interp3.get()
+
+    interp3.add(torch.tensor([3.0, 6.0]))
+    s1 = interp3.get()
+    s2 = interp3.get()
+    s3 = interp3.get()
+
+    torch.testing.assert_close(s1, torch.tensor([1.0, 2.0]))
+    torch.testing.assert_close(s2, torch.tensor([2.0, 4.0]))
+    torch.testing.assert_close(s3, torch.tensor([3.0, 6.0]))
+    assert interp3.get() is None
+
+
+def test_interpolation_3x_last_step_equals_target(interp3):
+    """Test last interpolated step equals the target action exactly."""
+    interp3.add(torch.tensor([10.0]))
+    interp3.get()
+
+    target = torch.tensor([100.0])
+    interp3.add(target)
+    interp3.get()
+    interp3.get()
+    last = interp3.get()
+    torch.testing.assert_close(last, target)
+
+
+# ====================== Reset Tests ======================
+
+
+def test_reset_clears_buffer(interp2):
+    """Test reset() clears the action buffer."""
+    interp2.add(torch.tensor([1.0]))
+    interp2.reset()
+    assert interp2.needs_new_action()
+    assert interp2.get() is None
+
+
+def test_reset_clears_prev(interp2):
+    """Test after reset, next add produces single-element buffer (no prev)."""
+    interp2.add(torch.tensor([0.0]))
+    interp2.get()
+    interp2.add(torch.tensor([10.0]))
+    interp2.get()
+    interp2.get()
+
+    interp2.reset()
+    interp2.add(torch.tensor([5.0]))
+    result = interp2.get()
+    torch.testing.assert_close(result, torch.tensor([5.0]))
+    assert interp2.get() is None
+
+
+def test_reset_episode_boundary(interp2):
+    """Test reset between two simulated episodes."""
+    interp2.add(torch.tensor([0.0]))
+    interp2.get()
+    interp2.add(torch.tensor([10.0]))
+    interp2.get()
+    interp2.get()
+
+    interp2.reset()
+
+    interp2.add(torch.tensor([100.0]))
+    result = interp2.get()
+    torch.testing.assert_close(result, torch.tensor([100.0]))
+    assert interp2.get() is None
+
+
+# ====================== get_control_interval Tests ======================
+
+
+def test_control_interval_30fps_multiplier_1():
+    """Test control interval at 30fps with no interpolation."""
+    interp = ActionInterpolator(multiplier=1)
+    assert interp.get_control_interval(30.0) == pytest.approx(1.0 / 30.0)
+
+
+def test_control_interval_30fps_multiplier_2(interp2):
+    """Test control interval at 30fps with 2x interpolation."""
+    assert interp2.get_control_interval(30.0) == pytest.approx(1.0 / 60.0)
+
+
+def test_control_interval_30fps_multiplier_3(interp3):
+    """Test control interval at 30fps with 3x interpolation."""
+    assert interp3.get_control_interval(30.0) == pytest.approx(1.0 / 90.0)
+
+
+def test_control_interval_60fps_multiplier_2(interp2):
+    """Test control interval at 60fps with 2x interpolation."""
+    assert interp2.get_control_interval(60.0) == pytest.approx(1.0 / 120.0)
+
+
+# ====================== get() on Empty Tests ======================
+
+
+def test_get_returns_none_before_any_add():
+    """Test get() returns None when no action has been added."""
+    interp = ActionInterpolator(multiplier=2)
+    assert interp.get() is None
+
+
+def test_get_returns_none_after_reset(interp2):
+    """Test get() returns None after reset."""
+    interp2.add(torch.tensor([1.0]))
+    interp2.reset()
+    assert interp2.get() is None
+
+
+# ====================== Multi-Dimensional Action Tests ======================
+
+
+def test_6dof_interpolation(interp2):
+    """Test interpolation works correctly with 6-dimensional actions."""
+    prev = torch.zeros(6)
+    target = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
+
+    interp2.add(prev)
+    interp2.get()
+
+    interp2.add(target)
+    mid = interp2.get()
+    end = interp2.get()
+
+    torch.testing.assert_close(mid, target / 2)
+    torch.testing.assert_close(end, target)
+
+
+# ====================== Simulated Control Loop Tests ======================
+
+
+def test_control_loop_produces_correct_action_count():
+    """Test N policy actions with multiplier M yields 1 + (N-1)*M robot commands."""
+    multiplier = 3
+    n_policy_actions = 5
+    interp = ActionInterpolator(multiplier=multiplier)
+
+    robot_commands = 0
+    for i in range(n_policy_actions):
+        action = torch.tensor([float(i)])
+        if interp.needs_new_action():
+            interp.add(action)
+        while True:
+            a = interp.get()
+            if a is None:
+                break
+            robot_commands += 1
+
+    expected = 1 + (n_policy_actions - 1) * multiplier
+    assert robot_commands == expected
+
+
+def test_control_loop_monotonic_increase():
+    """Test actions [0, 1, 2, 3] with multiplier=2 produce monotonically increasing values."""
+    interp = ActionInterpolator(multiplier=2)
+    all_values = []
+
+    for i in range(4):
+        interp.add(torch.tensor([float(i)]))
+        while True:
+            a = interp.get()
+            if a is None:
+                break
+            all_values.append(a.item())
+
+    for i in range(1, len(all_values)):
+        assert all_values[i] >= all_values[i - 1]
+
+
+# ====================== ActionQueue + ActionInterpolator Integration Tests ======================
+
+
+def _make_chunk(n_steps: int, action_dim: int = 2, offset: float = 0.0) -> torch.Tensor:
+    """Create a simple action chunk: each row is [offset + step_idx, offset + step_idx]."""
+    return torch.arange(n_steps, dtype=torch.float32).unsqueeze(1).expand(-1, action_dim) + offset
+
+
+def test_queue_interpolator_consumption_rate_matches_base_fps():
+    """Test queue.get() is called at base fps rate, not multiplied fps."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=3)
+
+    chunk = _make_chunk(10)
+    queue.merge(chunk, chunk.clone(), real_delay=0)
+
+    queue_gets = 0
+    control_ticks = 0
+
+    while True:
+        if interp.needs_new_action():
+            if queue.empty():
+                break
+            action = queue.get()
+            if action is None:
+                break
+            interp.add(action)
+            queue_gets += 1
+
+        result = interp.get()
+        if result is not None:
+            control_ticks += 1
+
+    assert queue_gets == 10
+    assert control_ticks == 1 + 9 * 3
+
+
+def test_queue_interpolator_leftover_decreases_only_on_queue_get():
+    """Test get_left_over() shrinks only on queue.get(), not on interpolator sub-steps."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=3)
+
+    chunk = _make_chunk(6)
+    queue.merge(chunk, chunk.clone(), real_delay=0)
+
+    assert interp.needs_new_action()
+    interp.add(queue.get())
+    leftover_after_first_get = queue.get_left_over()
+    assert leftover_after_first_get is not None
+    assert len(leftover_after_first_get) == 5
+
+    interp.get()
+    assert len(queue.get_left_over()) == 5
+
+    interp.add(queue.get())
+    assert len(queue.get_left_over()) == 4
+
+    for _ in range(3):
+        assert interp.get() is not None
+    assert len(queue.get_left_over()) == 4
+
+
+def test_queue_interpolator_processed_leftover_tracks_queue_index():
+    """Test get_processed_left_over() reflects queue's last_index, not interpolator state."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=2)
+
+    original = _make_chunk(8, offset=0.0)
+    processed = _make_chunk(8, offset=100.0)
+    queue.merge(original, processed, real_delay=0)
+
+    left = queue.get_processed_left_over()
+    assert len(left) == 8
+
+    for _ in range(3):
+        if interp.needs_new_action():
+            action = queue.get()
+            if action is not None:
+                interp.add(action)
+        interp.get()
+
+    proc_left = queue.get_processed_left_over()
+    orig_left = queue.get_left_over()
+    assert proc_left is not None and orig_left is not None
+    assert len(proc_left) == len(orig_left)
+    assert proc_left[0, 0].item() >= 100.0
+    assert orig_left[0, 0].item() < 100.0
+
+
+def test_queue_interpolator_merge_resets_queue_but_interpolator_keeps_prev():
+    """Test queue merge doesn't affect interpolator's prev, enabling smooth transitions."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=2)
+
+    chunk1 = torch.tensor([[0.0], [2.0], [4.0], [6.0], [8.0]])
+    queue.merge(chunk1, chunk1.clone(), real_delay=0)
+
+    consumed = []
+    for _ in range(5):
+        if interp.needs_new_action():
+            a = queue.get()
+            if a is not None:
+                interp.add(a)
+        r = interp.get()
+        if r is not None:
+            consumed.append(r.item())
+
+    assert interp.needs_new_action()
+    assert consumed[-1] == pytest.approx(4.0)
+
+    idx_before = queue.get_action_index()
+
+    chunk2 = torch.tensor([[10.0], [12.0], [14.0]])
+    queue.merge(chunk2, chunk2.clone(), real_delay=0, action_index_before_inference=idx_before)
+
+    first_action = queue.get()
+    assert first_action is not None
+    interp.add(first_action)
+    first_from_new = interp.get()
+    assert first_from_new is not None
+    assert first_from_new.item() == pytest.approx(7.0)
+
+
+def test_queue_interpolator_reset_does_not_affect_queue():
+    """Test interpolator reset leaves queue state untouched."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=2)
+
+    chunk = _make_chunk(5)
+    queue.merge(chunk, chunk.clone(), real_delay=0)
+
+    interp.add(queue.get())
+    interp.get()
+    interp.add(queue.get())
+    interp.get()
+    interp.get()
+
+    assert queue.qsize() == 3
+
+    interp.reset()
+
+    assert queue.qsize() == 3
+    assert len(queue.get_left_over()) == 3
+
+    interp.add(queue.get())
+    result = interp.get()
+    assert result is not None
+    assert queue.qsize() == 2
+
+
+def test_queue_interpolator_no_interpolation_1_to_1():
+    """Test multiplier=1 produces exactly 1 robot command per queue.get()."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=1)
+
+    chunk = _make_chunk(5)
+    queue.merge(chunk, chunk.clone(), real_delay=0)
+
+    robot_commands = 0
+    while not queue.empty():
+        if interp.needs_new_action():
+            action = queue.get()
+            if action is not None:
+                interp.add(action)
+        result = interp.get()
+        if result is not None:
+            robot_commands += 1
+
+    assert robot_commands == 5
+
+
+def test_queue_interpolator_delay_skips_stale_actions():
+    """Test merge with delay correctly skips stale actions for the interpolator."""
+    cfg = RTCConfig(enabled=True, execution_horizon=10)
+    queue = ActionQueue(cfg)
+    interp = ActionInterpolator(multiplier=2)
+
+    chunk1 = _make_chunk(10)
+    queue.merge(chunk1, chunk1.clone(), real_delay=0)
+
+    for _ in range(5):
+        if interp.needs_new_action():
+            a = queue.get()
+            if a is not None:
+                interp.add(a)
+        interp.get()
+
+    assert queue.get_action_index() == 3
+
+    chunk2 = _make_chunk(10, offset=100.0)
+    queue.merge(chunk2, chunk2.clone(), real_delay=3, action_index_before_inference=0)
+
+    first_action = queue.get()
+    assert first_action is not None
+    torch.testing.assert_close(first_action, torch.tensor([103.0, 103.0]))

tests/policies/rtc/test_action_queue.py

@@ -25,7 +25,7 @@ import torch
 from lerobot.policies.rtc.action_queue import ActionQueue
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
 
-# ====================== Fixtures ======================
+# Fixtures
 
 
 @pytest.fixture
@@ -63,7 +63,7 @@ def action_queue_rtc_disabled(rtc_config_disabled):
     return ActionQueue(rtc_config_disabled)
 
 
-# ====================== Initialization Tests ======================
+# Initialization tests
 
 
 def test_action_queue_initialization_rtc_enabled(rtc_config_enabled):
@@ -84,7 +84,7 @@ def test_action_queue_initialization_rtc_disabled(rtc_config_disabled):
     assert queue.cfg.enabled is False
 
 
-# ====================== get() Tests ======================
+# get() tests
 
 
 def test_get_returns_none_when_empty(action_queue_rtc_enabled):
@@ -136,7 +136,7 @@ def test_get_increments_last_index(action_queue_rtc_enabled, sample_actions):
     assert action_queue_rtc_enabled.last_index == 2
 
 
-# ====================== qsize() Tests ======================
+# qsize() tests
 
 
 def test_qsize_returns_zero_when_empty(action_queue_rtc_enabled):
@@ -167,7 +167,7 @@ def test_qsize_after_exhaustion(action_queue_rtc_enabled, sample_actions):
     assert action_queue_rtc_enabled.qsize() == 0
 
 
-# ====================== empty() Tests ======================
+# empty() tests
 
 
 def test_empty_returns_true_when_empty(action_queue_rtc_enabled):
@@ -202,7 +202,7 @@ def test_empty_after_full_consumption(action_queue_rtc_enabled, sample_actions):
     assert action_queue_rtc_enabled.empty() is True
 
 
-# ====================== get_action_index() Tests ======================
+# get_action_index() tests
 
 
 def test_get_action_index_initial_value(action_queue_rtc_enabled):
@@ -222,7 +222,7 @@ def test_get_action_index_after_consumption(action_queue_rtc_enabled, sample_act
     assert action_queue_rtc_enabled.get_action_index() == 3
 
 
-# ====================== get_left_over() Tests ======================
+# get_left_over() tests
 
 
 def test_get_left_over_returns_none_when_empty(action_queue_rtc_enabled):
@@ -269,7 +269,7 @@ def test_get_left_over_returns_empty_after_exhaustion(action_queue_rtc_enabled,
     assert leftover.shape == (0, 6)
 
 
-# ====================== merge() with RTC Enabled Tests ======================
+# merge() with RTC enabled tests
 
 
 def test_merge_replaces_queue_when_rtc_enabled(action_queue_rtc_enabled, sample_actions):
@@ -336,7 +336,7 @@ def test_merge_with_large_delay(action_queue_rtc_enabled, sample_actions):
     assert action_queue_rtc_enabled.qsize() == 0
 
 
-# ====================== merge() with RTC Disabled Tests ======================
+# merge() with RTC disabled tests
 
 
 def test_merge_appends_when_rtc_disabled(action_queue_rtc_disabled, sample_actions):
@@ -402,7 +402,7 @@ def test_merge_first_call_with_rtc_disabled(action_queue_rtc_disabled, sample_ac
     assert action_queue_rtc_disabled.last_index == 0
 
 
-# ====================== merge() with Different Action Shapes Tests ======================
+# merge() with different action shapes tests
 
 
 def test_merge_with_different_action_dims():
@@ -431,7 +431,7 @@ def test_merge_with_different_lengths():
     assert queue.qsize() == 35
 
 
-# ====================== merge() Delay Validation Tests ======================
+# merge() delay validation tests
 
 
 def test_merge_validates_delay_consistency(action_queue_rtc_enabled, sample_actions, caplog):
@@ -509,7 +509,7 @@ def test_merge_skips_validation_when_action_index_none(action_queue_rtc_enabled,
     assert "Indexes diff is not equal to real delay" not in caplog.text
 
 
-# ====================== Thread Safety Tests ======================
+# Thread safety tests
 
 
 def test_get_is_thread_safe(action_queue_rtc_enabled, sample_actions):
@@ -621,7 +621,7 @@ def test_concurrent_get_and_merge(action_queue_rtc_disabled, sample_actions):
     assert consumed_count[0] <= 200
 
 
-# ====================== get_left_over() Thread Safety Tests ======================
+# get_left_over() thread safety tests
 
 
 def test_get_left_over_is_thread_safe(action_queue_rtc_enabled, sample_actions):
@@ -670,7 +670,7 @@ def test_get_left_over_is_thread_safe(action_queue_rtc_enabled, sample_actions):
     assert len(leftovers) > 0
 
 
-# ====================== Edge Cases Tests ======================
+# Edge cases tests
 
 
 def test_queue_with_single_action(action_queue_rtc_enabled):
@@ -773,7 +773,7 @@ def test_qsize_with_none_queue(action_queue_rtc_enabled):
     assert action_queue_rtc_enabled.qsize() == 0
 
 
-# ====================== Integration Tests ======================
+# Integration tests
 
 
 def test_typical_rtc_workflow(action_queue_rtc_enabled, sample_actions):

tests/policies/rtc/test_rtc_relative_actions.py

@@ -0,0 +1,607 @@
+"""Tests for RTC + relative actions integration.
+
+Validates that Real-Time Chunking (RTC) works correctly when the policy uses
+relative actions. The key invariant: RTC guidance operates in model space
+(normalized relative actions), while the robot receives absolute actions after postprocessing.
+
+Flow under test:
+  Preprocessor: raw obs → relative step caches state → normalizer
+  Model: generates normalized relative actions (guided by RTC using leftover relative actions)
+  Postprocessor: unnormalize → absolute step (relative + cached state) → robot actions
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import torch
+
+from lerobot.configs.types import (
+    FeatureType,
+    NormalizationMode,
+    PolicyFeature,
+    RTCAttentionSchedule,
+)
+from lerobot.processor import TransitionKey, batch_to_transition
+from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
+from lerobot.processor.relative_action_processor import (
+    AbsoluteActionsProcessorStep,
+    RelativeActionsProcessorStep,
+    to_relative_actions,
+)
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+def _import_rtc_module(module_name: str, filename: str):
+    """Import an RTC module directly from its file path, bypassing lerobot.policies.__init__."""
+    rtc_dir = Path(__file__).resolve().parents[3] / "src" / "lerobot" / "policies" / "rtc"
+    spec = importlib.util.spec_from_file_location(module_name, rtc_dir / filename)
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = mod
+    spec.loader.exec_module(mod)
+    return mod
+
+
+_rtc_cfg_mod = _import_rtc_module("lerobot.policies.rtc.configuration_rtc", "configuration_rtc.py")
+RTCConfig = _rtc_cfg_mod.RTCConfig
+
+_action_queue_mod = _import_rtc_module("lerobot.policies.rtc.action_queue", "action_queue.py")
+ActionQueue = _action_queue_mod.ActionQueue
+
+_rtc_debug_mod = _import_rtc_module("lerobot.policies.rtc.debug_tracker", "debug_tracker.py")
+_rtc_mod = _import_rtc_module("lerobot.policies.rtc.modeling_rtc", "modeling_rtc.py")
+RTCProcessor = _rtc_mod.RTCProcessor
+
+ACTION_DIM = 6
+CHUNK_SIZE = 50
+EXECUTION_HORIZON = 10
+
+
+def _make_rtc_config(enabled=True):
+    return RTCConfig(
+        enabled=enabled,
+        execution_horizon=EXECUTION_HORIZON,
+        max_guidance_weight=10.0,
+        prefix_attention_schedule=RTCAttentionSchedule.EXP,
+    )
+
+
+def _make_relative_pipeline(action_dim=ACTION_DIM, norm_mode=NormalizationMode.MEAN_STD):
+    """Build paired relative/absolute processor steps and normalizer/unnormalizer."""
+    features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))}
+    norm_map = {FeatureType.ACTION: norm_mode}
+
+    stats = {
+        ACTION: {
+            "mean": torch.zeros(action_dim).numpy(),
+            "std": torch.ones(action_dim).numpy(),
+            "q01": (-2 * torch.ones(action_dim)).numpy(),
+            "q99": (2 * torch.ones(action_dim)).numpy(),
+            "min": (-3 * torch.ones(action_dim)).numpy(),
+            "max": (3 * torch.ones(action_dim)).numpy(),
+        }
+    }
+
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+    return relative_step, normalizer, unnormalizer, absolute_step
+
+
+class TestActionQueueRelativeActions:
+    """Verify ActionQueue stores model-space (relative) actions for RTC and absolute for robot."""
+
+    def test_left_over_returns_relative_actions(self):
+        """get_left_over() should return the original (relative-space) actions."""
+        cfg = _make_rtc_config()
+        queue = ActionQueue(cfg)
+
+        relative_actions = torch.randn(CHUNK_SIZE, ACTION_DIM)
+        absolute_actions = torch.randn(CHUNK_SIZE, ACTION_DIM)
+        queue.merge(relative_actions, absolute_actions, real_delay=0)
+
+        for _ in range(5):
+            queue.get()
+
+        leftover = queue.get_left_over()
+        torch.testing.assert_close(leftover, relative_actions[5:])
+
+    def test_robot_receives_absolute_actions(self):
+        """The robot (via get()) should receive postprocessed absolute actions."""
+        cfg = _make_rtc_config()
+        queue = ActionQueue(cfg)
+
+        relative_actions = torch.randn(CHUNK_SIZE, ACTION_DIM)
+        absolute_actions = torch.randn(CHUNK_SIZE, ACTION_DIM)
+        queue.merge(relative_actions, absolute_actions, real_delay=0)
+
+        first_action = queue.get()
+        torch.testing.assert_close(first_action, absolute_actions[0])
+
+
+class TestRTCDenoiseWithRelativeLeftovers:
+    """Verify RTC denoise_step correctly handles relative-space prev_chunk_left_over."""
+
+    def test_first_chunk_no_guidance(self):
+        """First chunk (no leftovers) should return v_t without guidance."""
+        rtc = RTCProcessor(_make_rtc_config())
+        x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        def mock_denoise(x):
+            return torch.ones_like(x)
+
+        result = rtc.denoise_step(
+            x_t=x_t,
+            prev_chunk_left_over=None,
+            inference_delay=0,
+            time=0.5,
+            original_denoise_step_partial=mock_denoise,
+        )
+        torch.testing.assert_close(result, torch.ones_like(x_t))
+
+    def test_relative_leftovers_shape_preserved(self):
+        """RTC output should have the same shape as input regardless of leftover shape."""
+        rtc = RTCProcessor(_make_rtc_config())
+        x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        shorter_leftover = torch.randn(1, 20, ACTION_DIM)
+
+        def mock_denoise(x):
+            return torch.zeros_like(x)
+
+        result = rtc.denoise_step(
+            x_t=x_t,
+            prev_chunk_left_over=shorter_leftover,
+            inference_delay=5,
+            time=0.5,
+            original_denoise_step_partial=mock_denoise,
+        )
+        assert result.shape == x_t.shape
+
+    def test_guidance_steers_toward_previous_relative_actions(self):
+        """RTC guidance should push x1_t toward prev_chunk_left_over in relative space."""
+        rtc = RTCProcessor(_make_rtc_config())
+        x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        prev_relatives = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        def mock_denoise(x):
+            return torch.zeros_like(x)
+
+        result_without_guidance = rtc.denoise_step(
+            x_t=x_t.clone(),
+            prev_chunk_left_over=None,
+            inference_delay=5,
+            time=0.5,
+            original_denoise_step_partial=mock_denoise,
+        )
+
+        result_with_guidance = rtc.denoise_step(
+            x_t=x_t.clone(),
+            prev_chunk_left_over=prev_relatives,
+            inference_delay=5,
+            time=0.5,
+            original_denoise_step_partial=mock_denoise,
+        )
+
+        assert not torch.allclose(result_with_guidance, result_without_guidance, atol=1e-6)
+
+
+class TestFullPipelineRelativeRTC:
+    """End-to-end test of the RTC + relative actions pipeline matching eval_with_real_robot.py flow."""
+
+    def test_preprocessor_caches_state_for_postprocessor(self):
+        """Preprocessor's relative step should cache state so postprocessor can convert back."""
+        relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
+
+        state = torch.randn(1, ACTION_DIM)
+        actions = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        batch = {ACTION: actions, OBS_STATE: state}
+        transition = batch_to_transition(batch)
+
+        relative_step(transition)
+        assert relative_step._last_state is not None
+        torch.testing.assert_close(relative_step._last_state, state)
+
+    def test_preprocessor_caches_state_without_actions(self):
+        """During inference, preprocessor receives only observations (no actions).
+        Relative step should still cache state for the postprocessor."""
+        relative_step, _, _, _ = _make_relative_pipeline()
+
+        state = torch.randn(1, ACTION_DIM)
+        batch = {OBS_STATE: state}
+        transition = batch_to_transition(batch)
+
+        relative_step(transition)
+        assert relative_step._last_state is not None
+        torch.testing.assert_close(relative_step._last_state, state)
+
+    def test_roundtrip_with_identity_normalization(self):
+        """Actions → relative → normalize → [model] → unnormalize → absolute should recover originals.
+        Using mean=0, std=1 normalization (identity)."""
+        relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
+
+        state = torch.randn(1, ACTION_DIM)
+        actions = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        batch = {ACTION: actions.clone(), OBS_STATE: state}
+        transition = batch_to_transition(batch)
+
+        t1 = relative_step(transition)
+        t2 = normalizer(t1)
+
+        model_output = t2[TransitionKey.ACTION].clone()
+
+        model_transition = {TransitionKey.ACTION: model_output}
+        t3 = unnormalizer(model_transition)
+        t4 = absolute_step(t3)
+
+        recovered = t4[TransitionKey.ACTION]
+        torch.testing.assert_close(recovered, actions, atol=1e-5, rtol=1e-5)
+
+    def test_eval_loop_simulation(self):
+        """Simulate the eval_with_real_robot.py loop with relative actions.
+
+        Iteration 1: No leftovers → model generates relative actions → store for RTC
+        Iteration 2: Use leftovers as RTC guidance → model generates new relative actions
+        Both iterations: postprocessor converts relative actions to absolute for robot
+        """
+        relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
+        rtc = RTCProcessor(_make_rtc_config())
+        queue = ActionQueue(_make_rtc_config())
+
+        def mock_model(prev_chunk_left_over, inference_delay, state):
+            """Simulate model generating relative actions with RTC."""
+            x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+            def denoise(x):
+                return -0.1 * x
+
+            guided_v = rtc.denoise_step(
+                x_t=x_t,
+                prev_chunk_left_over=prev_chunk_left_over,
+                inference_delay=inference_delay,
+                time=0.5,
+                original_denoise_step_partial=denoise,
+            )
+            return x_t - 0.5 * guided_v
+
+        # --- Iteration 1: first chunk, no leftovers ---
+        state_1 = torch.randn(1, ACTION_DIM)
+        obs_batch_1 = {OBS_STATE: state_1}
+        relative_step(batch_to_transition(obs_batch_1))
+
+        model_relatives_1 = mock_model(prev_chunk_left_over=None, inference_delay=0, state=state_1)
+        original_actions_1 = model_relatives_1.squeeze(0)
+
+        model_transition_1 = {TransitionKey.ACTION: model_relatives_1}
+        postprocessed_1 = absolute_step(unnormalizer(model_transition_1))[TransitionKey.ACTION].squeeze(0)
+
+        queue.merge(original_actions_1, postprocessed_1, real_delay=0)
+
+        # Consume some actions (simulate robot executing)
+        for _ in range(5):
+            action = queue.get()
+            assert action is not None
+
+        # --- Iteration 2: use leftovers for RTC ---
+        prev_actions = queue.get_left_over()
+        assert prev_actions is not None
+        assert prev_actions.shape[0] == CHUNK_SIZE - 5
+
+        state_2 = state_1 + 0.01 * torch.randn(1, ACTION_DIM)
+        obs_batch_2 = {OBS_STATE: state_2}
+        relative_step(batch_to_transition(obs_batch_2))
+
+        model_relatives_2 = mock_model(
+            prev_chunk_left_over=prev_actions.unsqueeze(0), inference_delay=3, state=state_2
+        )
+        original_actions_2 = model_relatives_2.squeeze(0)
+
+        model_transition_2 = {TransitionKey.ACTION: model_relatives_2}
+        postprocessed_2 = absolute_step(unnormalizer(model_transition_2))[TransitionKey.ACTION].squeeze(0)
+
+        queue.merge(original_actions_2, postprocessed_2, real_delay=3)
+
+        # Postprocessed actions should be in absolute space
+        action = queue.get()
+        assert action is not None
+        assert action.shape == (ACTION_DIM,)
+
+        # Verify leftovers are in relative space (original_queue stores relative actions)
+        leftover_relatives = queue.get_left_over()
+        assert leftover_relatives is not None
+        assert leftover_relatives.shape[1] == ACTION_DIM
+
+    def test_postprocessor_uses_correct_state_per_iteration(self):
+        """Each iteration's postprocessor should use the state from that iteration's preprocessor,
+        not a stale state from a previous iteration."""
+        relative_step, _, unnormalizer, absolute_step = _make_relative_pipeline()
+
+        state_1 = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        state_2 = torch.tensor([[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]])
+        relatives = torch.zeros(1, 5, ACTION_DIM)
+
+        # Iteration 1: cache state_1
+        relative_step(batch_to_transition({OBS_STATE: state_1}))
+        result_1 = absolute_step(unnormalizer({TransitionKey.ACTION: relatives.clone()}))[
+            TransitionKey.ACTION
+        ]
+        # relative=0 + state_1 should give state_1
+        for t in range(5):
+            torch.testing.assert_close(result_1[0, t], state_1[0], atol=1e-5, rtol=1e-5)
+
+        # Iteration 2: cache state_2
+        relative_step(batch_to_transition({OBS_STATE: state_2}))
+        result_2 = absolute_step(unnormalizer({TransitionKey.ACTION: relatives.clone()}))[
+            TransitionKey.ACTION
+        ]
+        for t in range(5):
+            torch.testing.assert_close(result_2[0, t], state_2[0], atol=1e-5, rtol=1e-5)
+
+
+class TestStateRebasingApproximation:
+    """Verify that the approximation from not rebasing leftover relative actions is small
+    when state changes between inference calls are small (real-time control regime)."""
+
+    def test_small_state_change_produces_small_error(self):
+        """With small state changes (typical in real-time control),
+        using stale relative actions for RTC guidance introduces negligible error."""
+        state_old = torch.randn(1, ACTION_DIM)
+        state_new = state_old + 0.01 * torch.randn(1, ACTION_DIM)
+
+        actions_absolute = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        mask = [True] * ACTION_DIM
+
+        relatives_old = to_relative_actions(actions_absolute, state_old, mask)
+        relatives_new = to_relative_actions(actions_absolute, state_new, mask)
+
+        error = (relatives_old - relatives_new).abs().mean()
+        state_change = (state_old - state_new).abs().mean()
+
+        # Error should be proportional to state change
+        assert error < 0.1, (
+            f"Relative-action error {error:.4f} should be small for small state change {state_change:.4f}"
+        )
+
+    def test_large_state_change_produces_proportional_error(self):
+        """With large state changes, stale relative actions diverge more (but RTC guidance decays)."""
+        state_old = torch.randn(1, ACTION_DIM)
+        state_new = state_old + 10.0 * torch.randn(1, ACTION_DIM)
+
+        actions_absolute = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        mask = [True] * ACTION_DIM
+
+        relatives_old = to_relative_actions(actions_absolute, state_old, mask)
+        relatives_new = to_relative_actions(actions_absolute, state_new, mask)
+
+        error = (relatives_old - relatives_new).abs().mean()
+        state_change = (state_old - state_new).abs().mean()
+
+        # Error should be roughly equal to state change
+        torch.testing.assert_close(
+            error.clone().detach(), state_change.clone().detach(), atol=1e-5, rtol=1e-5
+        )
+
+    def test_excluded_joints_not_affected_by_state_change(self):
+        """Joints excluded from relative conversion should not contribute rebasing error."""
+        state_old = torch.randn(1, ACTION_DIM)
+        state_new = state_old.clone()
+        state_new[0, -1] = state_old[0, -1] + 100.0
+
+        actions = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        mask = [True] * (ACTION_DIM - 1) + [False]
+
+        relatives_old = to_relative_actions(actions, state_old, mask)
+        relatives_new = to_relative_actions(actions, state_new, mask)
+
+        # Last dim (excluded) should have zero error
+        error_excluded = (relatives_old[..., -1] - relatives_new[..., -1]).abs().max()
+        assert error_excluded < 1e-6, f"Excluded joint should have zero error, got {error_excluded}"
+
+
+def _detect_relative_actions(preprocessor) -> bool:
+    """Mirror of the helper in eval_with_real_robot.py for testing without importing it."""
+    return any(isinstance(step, RelativeActionsProcessorStep) and step.enabled for step in preprocessor.steps)
+
+
+class TestDetectRelativeActions:
+    """Test the _detect_relative_actions helper logic used by eval_with_real_robot.py."""
+
+    def test_detects_enabled_relative_step(self):
+        class FakePipeline:
+            steps = [RelativeActionsProcessorStep(enabled=True)]
+
+        assert _detect_relative_actions(FakePipeline()) is True
+
+    def test_ignores_disabled_relative_step(self):
+        class FakePipeline:
+            steps = [RelativeActionsProcessorStep(enabled=False)]
+
+        assert _detect_relative_actions(FakePipeline()) is False
+
+    def test_returns_false_when_no_relative_step(self):
+        class FakePipeline:
+            steps = []
+
+        assert _detect_relative_actions(FakePipeline()) is False
+
+
+class TestNonRelativePolicy:
+    """Verify the same pipeline works when relative actions are disabled (standard absolute policy)."""
+
+    def test_disabled_relative_step_is_noop(self):
+        relative_step = RelativeActionsProcessorStep(enabled=False)
+        absolute_step = AbsoluteActionsProcessorStep(enabled=False, relative_step=relative_step)
+
+        state = torch.randn(1, ACTION_DIM)
+        actions = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        transition = batch_to_transition({ACTION: actions.clone(), OBS_STATE: state})
+        t1 = relative_step(transition)
+        torch.testing.assert_close(t1[TransitionKey.ACTION], actions)
+
+        t2 = absolute_step({TransitionKey.ACTION: actions.clone()})
+        torch.testing.assert_close(t2[TransitionKey.ACTION], actions)
+
+    def test_eval_loop_without_relative_actions(self):
+        """Full eval loop simulation with relative actions disabled: original and processed actions are identical."""
+        features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,))}
+        norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
+        stats = {
+            ACTION: {
+                "mean": torch.zeros(ACTION_DIM).numpy(),
+                "std": torch.ones(ACTION_DIM).numpy(),
+            }
+        }
+
+        relative_step = RelativeActionsProcessorStep(enabled=False)
+        unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+        absolute_step = AbsoluteActionsProcessorStep(enabled=False, relative_step=relative_step)
+
+        rtc = RTCProcessor(_make_rtc_config())
+        queue = ActionQueue(_make_rtc_config())
+
+        state = torch.randn(1, ACTION_DIM)
+        relative_step(batch_to_transition({OBS_STATE: state}))
+
+        model_output = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        post = absolute_step(unnormalizer({TransitionKey.ACTION: model_output.clone()}))[
+            TransitionKey.ACTION
+        ].squeeze(0)
+        original = model_output.squeeze(0)
+
+        # With identity norm and no relative-action transform, original and postprocessed should match
+        torch.testing.assert_close(original, post, atol=1e-5, rtol=1e-5)
+
+        queue.merge(original, post, real_delay=0)
+
+        for _ in range(5):
+            queue.get()
+
+        prev_actions = queue.get_left_over()
+        assert prev_actions is not None
+
+        # RTC guidance works the same way (absolute space)
+        x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        result = rtc.denoise_step(
+            x_t=x_t,
+            prev_chunk_left_over=prev_actions.unsqueeze(0),
+            inference_delay=3,
+            time=0.5,
+            original_denoise_step_partial=lambda x: torch.zeros_like(x),
+        )
+        assert result.shape == x_t.shape
+
+    def test_detect_relative_returns_false_when_disabled(self):
+        class FakePipeline:
+            steps = [RelativeActionsProcessorStep(enabled=False)]
+
+        assert not _detect_relative_actions(FakePipeline())
+
+    def test_detect_relative_returns_false_when_absent(self):
+        class FakePipeline:
+            steps = []
+
+        assert not _detect_relative_actions(FakePipeline())
+
+
+class TestMultiChunkConsistency:
+    """Test multiple RTC iterations with relative actions maintain consistency."""
+
+    def test_three_iteration_pipeline(self):
+        """Simulate three consecutive RTC iterations and verify queue state consistency."""
+        relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
+        queue = ActionQueue(_make_rtc_config())
+
+        states = [torch.randn(1, ACTION_DIM) + i * 0.01 for i in range(3)]
+
+        for i in range(3):
+            queue.get_left_over()
+
+            relative_step(batch_to_transition({OBS_STATE: states[i]}))
+
+            model_output = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+            post_transition = absolute_step(unnormalizer({TransitionKey.ACTION: model_output.clone()}))
+            postprocessed = post_transition[TransitionKey.ACTION].squeeze(0)
+            original = model_output.squeeze(0)
+
+            delay = min(i * 2, CHUNK_SIZE - 1)
+            queue.merge(original, postprocessed, real_delay=delay)
+
+            for _ in range(5):
+                action = queue.get()
+                assert action is not None
+                assert action.shape == (ACTION_DIM,)
+
+        # After 3 iterations, queue should still be in valid state
+        assert queue.qsize() > 0
+        leftover = queue.get_left_over()
+        assert leftover is not None
+        assert leftover.ndim == 2
+        assert leftover.shape[1] == ACTION_DIM
+
+    def test_leftover_and_processed_differ_when_relative_enabled(self):
+        """With relative actions enabled, original leftovers (relative) must differ from processed (absolute)."""
+        relative_step, _, unnormalizer, absolute_step = _make_relative_pipeline()
+        queue = ActionQueue(_make_rtc_config())
+
+        state = torch.tensor([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        relative_step(batch_to_transition({OBS_STATE: state}))
+
+        model_relatives = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        post = absolute_step(unnormalizer({TransitionKey.ACTION: model_relatives.clone()}))[
+            TransitionKey.ACTION
+        ].squeeze(0)
+        original = model_relatives.squeeze(0)
+
+        queue.merge(original, post, real_delay=0)
+
+        relative_leftover = queue.get_left_over()
+
+        # Leftovers (relative) must differ from the postprocessed absolute actions
+        assert not torch.allclose(relative_leftover, post, atol=1e-3)
+        state_expanded = state.squeeze(0).unsqueeze(0).expand_as(relative_leftover)
+        torch.testing.assert_close(post, relative_leftover + state_expanded, atol=1e-5, rtol=1e-5)
+
+    def test_rtc_guidance_uses_relative_space(self):
+        """Verify that RTC denoise_step receives relative-space leftovers, not absolute."""
+        relative_step, _, unnormalizer, absolute_step = _make_relative_pipeline()
+        rtc = RTCProcessor(_make_rtc_config())
+        queue = ActionQueue(_make_rtc_config())
+
+        state = torch.tensor([[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]])
+        relative_step(batch_to_transition({OBS_STATE: state}))
+
+        model_relatives = 0.1 * torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+        post = absolute_step(unnormalizer({TransitionKey.ACTION: model_relatives.clone()}))[
+            TransitionKey.ACTION
+        ].squeeze(0)
+        original = model_relatives.squeeze(0)
+
+        queue.merge(original, post, real_delay=0)
+
+        for _ in range(5):
+            queue.get()
+
+        prev_left_over = queue.get_left_over()
+
+        # prev_left_over should be small relative offsets (around 0.1 * randn), not large absolute values
+        assert prev_left_over.abs().mean() < 5.0, (
+            f"Leftover should be small relative offsets, got mean abs {prev_left_over.abs().mean():.2f}"
+        )
+
+        x_t = torch.randn(1, CHUNK_SIZE, ACTION_DIM)
+
+        def denoise(x):
+            return torch.zeros_like(x)
+
+        result = rtc.denoise_step(
+            x_t=x_t,
+            prev_chunk_left_over=prev_left_over.unsqueeze(0),
+            inference_delay=3,
+            time=0.5,
+            original_denoise_step_partial=denoise,
+        )
+
+        assert result.shape == x_t.shape

tests/policies/smolvla/test_smolvla_rtc.py

@@ -19,15 +19,15 @@
 import pytest
 import torch
 
-from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule  # noqa: E402
-from lerobot.policies.factory import make_pre_post_processors  # noqa: E402
-from lerobot.policies.rtc.configuration_rtc import RTCConfig  # noqa: E402
+from lerobot.configs.types import FeatureType, PolicyFeature, RTCAttentionSchedule
+from lerobot.policies.factory import make_pre_post_processors
+from lerobot.policies.rtc.configuration_rtc import RTCConfig
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig  # noqa: F401
-from lerobot.utils.random_utils import set_seed  # noqa: E402
-from tests.utils import require_cuda, require_package  # noqa: E402
+from lerobot.utils.random_utils import set_seed
+from tests.utils import require_cuda, skip_if_package_missing
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
 @require_cuda
 def test_smolvla_rtc_initialization():
     from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
@@ -65,7 +65,7 @@ def test_smolvla_rtc_initialization():
     print("✓ SmolVLA RTC initialization: Test passed")
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
 @require_cuda
 def test_smolvla_rtc_initialization_without_rtc_config():
     from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy  # noqa: F401
@@ -87,7 +87,7 @@ def test_smolvla_rtc_initialization_without_rtc_config():
     print("✓ SmolVLA RTC initialization without RTC config: Test passed")
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
 @require_cuda
 @pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
 def test_smolvla_rtc_inference_with_prev_chunk():
@@ -170,7 +170,7 @@ def test_smolvla_rtc_inference_with_prev_chunk():
     print("✓ SmolVLA RTC inference with prev_chunk: Test passed")
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
 @require_cuda
 @pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
 def test_smolvla_rtc_inference_without_prev_chunk():
@@ -244,7 +244,7 @@ def test_smolvla_rtc_inference_without_prev_chunk():
     print("✓ SmolVLA RTC inference without prev_chunk: Test passed")
 
 
-@require_package("transformers")
+@skip_if_package_missing("transformers")
 @require_cuda
 @pytest.mark.skipif(True, reason="Requires pretrained SmolVLA model weights")
 def test_smolvla_rtc_validation_rules():

tests/policies/test_policies.py

@@ -20,18 +20,18 @@ from pathlib import Path
 import einops
 import pytest
 import torch
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
 from packaging import version
 from safetensors.torch import load_file
 
-from lerobot import available_policies
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.configs.types import FeatureType, PolicyFeature
-from lerobot.datasets.factory import make_dataset
-from lerobot.datasets.feature_utils import dataset_to_policy_features
-from lerobot.datasets.utils import cycle
+from lerobot.datasets import make_dataset
 from lerobot.envs.factory import make_env, make_env_config
-from lerobot.envs.utils import preprocess_observation
+from lerobot.envs.utils import close_envs, preprocess_observation
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.act.modeling_act import ACTTemporalEnsembler
@@ -45,10 +45,23 @@ from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.policies.vqbet.modeling_vqbet import VQBeTHead
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.utils.feature_utils import dataset_to_policy_features
+from lerobot.utils.import_utils import is_package_available
 from lerobot.utils.random_utils import seeded_context
+from lerobot.utils.utils import cycle
 from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats
 from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
 
+# Policies that require optional heavy dependencies to instantiate
+_POLICY_REQUIRED_PACKAGES: dict[str, tuple[str, ...]] = {
+    "diffusion": ("diffusers",),
+}
+
+_ALL_POLICIES = ["act", "diffusion", "tdmpc", "vqbet"]
+AVAILABLE_POLICIES = [
+    p for p in _ALL_POLICIES if all(is_package_available(pkg) for pkg in _POLICY_REQUIRED_PACKAGES.get(p, ()))
+]
+
 
 @pytest.fixture
 def dummy_dataset_metadata(lerobot_dataset_metadata_factory, info_factory, tmp_path):
@@ -84,7 +97,7 @@ def dummy_dataset_metadata(lerobot_dataset_metadata_factory, info_factory, tmp_p
     return ds_meta
 
 
-@pytest.mark.parametrize("policy_name", available_policies)
+@pytest.mark.parametrize("policy_name", AVAILABLE_POLICIES)
 def test_get_policy_and_config_classes(policy_name: str):
     """Check that the correct policy and config classes are returned."""
     policy_cls = get_policy_class(policy_name)
@@ -224,6 +237,8 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
     # Test step through policy
     env.step(action)
 
+    close_envs(envs)
+
 
 # TODO(rcadene, aliberts): This test is quite end-to-end. Move this test in test_optimizer?
 def test_act_backbone_lr():
@@ -253,7 +268,7 @@ def test_act_backbone_lr():
     assert len(optimizer.param_groups[1]["params"]) == 20
 
 
-@pytest.mark.parametrize("policy_name", available_policies)
+@pytest.mark.parametrize("policy_name", AVAILABLE_POLICIES)
 def test_policy_defaults(dummy_dataset_metadata, policy_name: str):
     """Check that the policy can be instantiated with defaults."""
     policy_cls = get_policy_class(policy_name)
@@ -266,7 +281,7 @@ def test_policy_defaults(dummy_dataset_metadata, policy_name: str):
     policy_cls(policy_cfg)
 
 
-@pytest.mark.parametrize("policy_name", available_policies)
+@pytest.mark.parametrize("policy_name", AVAILABLE_POLICIES)
 def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name: str):
     policy_cls = get_policy_class(policy_name)
     policy_cfg = make_policy_config(policy_name)
@@ -341,7 +356,7 @@ def test_multikey_construction(multikey: bool):
         # to normalize the image at all. In our current codebase we dont normalize at all. But there is still a minor difference
         # that fails the test. However, by testing to normalize the image with 0.5 0.5 in the current codebase, the test pass.
         # Thus, we deactivate this test for now.
-        (
+        pytest.param(
             "lerobot/pusht",
             "diffusion",
             {
@@ -350,6 +365,7 @@ def test_multikey_construction(multikey: bool):
                 "down_dims": [128, 256, 512],
             },
             "",
+            marks=pytest.mark.skipif(not is_package_available("diffusers"), reason="diffusers not installed"),
         ),
         ("lerobot/aloha_sim_insertion_human", "act", {"n_action_steps": 10}, ""),
         (

tests/policies/test_relative_actions.py

@@ -0,0 +1,348 @@
+"""Tests for relative action transforms — full pipeline validation.
+
+Tests the complete flow matching OpenPI:
+  raw actions → RelativeActions → Normalize(relative_stats) → model → Unnormalize → AbsoluteActions
+
+Uses real dataset: lerobot-data-collection/dagger_final_1_21
+"""
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.datasets.compute_stats import get_feature_stats
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.processor import TransitionKey, batch_to_transition
+from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
+from lerobot.processor.relative_action_processor import (
+    AbsoluteActionsProcessorStep,
+    RelativeActionsProcessorStep,
+    to_absolute_actions,
+    to_relative_actions,
+)
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+CHUNK_SIZE = 10
+REPO_ID = "lerobot-data-collection/dagger_final_1_21"
+
+
+@pytest.fixture(scope="module")
+def dataset():
+    return LeRobotDataset(REPO_ID, episodes=[0])
+
+
+@pytest.fixture(scope="module")
+def action_dim(dataset):
+    return dataset.meta.features["action"]["shape"][0]
+
+
+def _build_action_chunks(dataset, chunk_size, max_chunks=50):
+    """Build action chunks from hf_dataset, like the training script does."""
+    hf = dataset.hf_dataset
+    total = len(hf)
+    all_ep = torch.tensor([int(hf[i]["episode_index"]) for i in range(total)])
+    chunks, states = [], []
+    for i in range(total - chunk_size + 1):
+        if all_ep[i] != all_ep[i + chunk_size - 1]:
+            continue
+        chunk_actions = torch.stack([hf[i + k]["action"] for k in range(chunk_size)]).float()
+        state = hf[i]["observation.state"].float()
+        chunks.append(chunk_actions)
+        states.append(state)
+        if len(chunks) >= max_chunks:
+            break
+    assert len(chunks) > 0, f"No valid chunks found. total={total}, ep_indices={all_ep.tolist()}"
+    return torch.stack(chunks), torch.stack(states)
+
+
+def _compute_relative_chunk_stats(action_chunks, states, mask):
+    all_chunks = []
+    for actions, state in zip(action_chunks, states, strict=True):
+        relative = to_relative_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        all_chunks.append(relative.numpy())
+    all_relative = np.concatenate(all_chunks, axis=0)
+    return get_feature_stats(all_relative, axis=0, keepdims=all_relative.ndim == 1)
+
+
+# Basic roundtrip tests
+
+
+def test_roundtrip_3d(action_dim):
+    actions = torch.randn(4, CHUNK_SIZE, action_dim)
+    state = torch.randn(4, action_dim)
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(to_relative_actions(actions, state, mask), state, mask)
+    torch.testing.assert_close(recovered, actions)
+
+
+def test_roundtrip_2d(action_dim):
+    actions = torch.randn(4, action_dim)
+    state = torch.randn(4, action_dim)
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(to_relative_actions(actions, state, mask), state, mask)
+    torch.testing.assert_close(recovered, actions)
+
+
+def test_no_mutation(action_dim):
+    actions = torch.randn(2, CHUNK_SIZE, action_dim)
+    original = actions.clone()
+    state = torch.randn(2, action_dim)
+    to_relative_actions(actions, state, [True] * action_dim)
+    torch.testing.assert_close(actions, original)
+
+
+def test_exclude_joints_supports_partial_name_matching():
+    names = [
+        "right_joint_1.pos",
+        "right_gripper.pos",
+        "left_joint_1.pos",
+        "left_gripper.pos",
+    ]
+    step = RelativeActionsProcessorStep(enabled=True, exclude_joints=["gripper"], action_names=names)
+    assert step._build_mask(len(names)) == [True, False, True, False]
+
+
+# Chunk-level relative stats test
+
+
+def test_chunk_stats_have_larger_std_than_frame_stats(dataset, action_dim):
+    """Chunk-level relative stats should have larger std than per-frame relative stats."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    chunk_stats = _compute_relative_chunk_stats(action_chunks, states, mask)
+
+    # Per-frame stats
+    hf = dataset.hf_dataset
+    n = min(500, len(hf))
+    frame_actions = torch.stack([hf[i]["action"] for i in range(n)]).float()
+    frame_states = torch.stack([hf[i]["observation.state"] for i in range(n)]).float()
+    frame_relatives = to_relative_actions(frame_actions, frame_states, mask).numpy()
+    frame_stats = get_feature_stats(frame_relatives, axis=0, keepdims=frame_relatives.ndim == 1)
+
+    assert chunk_stats["std"].mean() >= frame_stats["std"].mean(), (
+        f"Chunk std ({chunk_stats['std'].mean():.4f}) should be >= "
+        f"frame std ({frame_stats['std'].mean():.4f})"
+    )
+
+
+# Full pipeline roundtrip: relative → normalize → unnormalize → absolute
+
+
+def test_full_pipeline_roundtrip(dataset, action_dim):
+    """Test the complete OpenPI pipeline: relative → normalize → unnormalize → absolute."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    relative_stats = _compute_relative_chunk_stats(action_chunks, states, mask)
+    stats = {ACTION: dict(relative_stats.items())}
+
+    features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))}
+    norm_map = {FeatureType.ACTION: NormalizationMode.MEAN_STD}
+
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    original_actions = action_chunks[0].unsqueeze(0)
+    state = states[0].unsqueeze(0)
+
+    batch = {ACTION: original_actions, OBS_STATE: state}
+    transition = batch_to_transition(batch)
+
+    # Forward: relative → normalize
+    t1 = relative_step(transition)
+    t2 = normalizer(t1)
+
+    normalized_action = t2[TransitionKey.ACTION]
+    assert normalized_action.abs().mean() < 10, (
+        f"Normalized actions should be in reasonable range, got mean abs {normalized_action.abs().mean():.2f}"
+    )
+
+    # Reverse: unnormalize → absolute
+    t3 = unnormalizer(t2)
+    t4 = absolute_step(t3)
+
+    recovered_actions = t4[TransitionKey.ACTION]
+    torch.testing.assert_close(recovered_actions, original_actions, atol=1e-4, rtol=1e-4)
+
+
+def test_normalized_relative_values_are_reasonable(dataset, action_dim):
+    """With correct chunk stats, normalized relative actions should be in a reasonable range."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    relative_stats = _compute_relative_chunk_stats(action_chunks, states, mask)
+    mean = torch.tensor(relative_stats["mean"]).float()
+    std = torch.tensor(relative_stats["std"]).float()
+
+    all_normalized = []
+    for actions, state in zip(action_chunks, states, strict=True):
+        relative = to_relative_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        normalized = (relative - mean) / (std + 1e-6)
+        all_normalized.append(normalized)
+
+    all_normalized = torch.cat(all_normalized, dim=0)
+
+    pct_in_range = (all_normalized.abs() < 5).float().mean()
+    assert pct_in_range > 0.9, (
+        f"Only {pct_in_range * 100:.1f}% of normalized values in [-5, 5], expected >90%"
+    )
+
+    assert all_normalized.mean().abs() < 1.0, (
+        f"Mean of normalized relative actions is {all_normalized.mean():.2f}, expected near 0"
+    )
+
+
+def test_processor_step_roundtrip(dataset, action_dim):
+    """RelativeActionsProcessorStep applies relative offsets; to_absolute_actions recovers original."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(4)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(4)]),
+    }
+    original_actions = batch[ACTION].clone()
+    transition = batch_to_transition(batch)
+
+    step = RelativeActionsProcessorStep(enabled=True)
+    relative_transition = step(transition)
+    assert not torch.allclose(relative_transition[TransitionKey.ACTION], original_actions)
+
+    state = transition[TransitionKey.OBSERVATION][OBS_STATE]
+    mask = [True] * action_dim
+    recovered = to_absolute_actions(relative_transition[TransitionKey.ACTION], state, mask)
+    torch.testing.assert_close(recovered, original_actions)
+
+
+def test_processor_step_disabled_is_noop(dataset, action_dim):
+    """enabled=False should be a no-op."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(2)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(2)]),
+    }
+    original = batch[ACTION].clone()
+    transition = batch_to_transition(batch)
+    result = RelativeActionsProcessorStep(enabled=False)(transition)
+    torch.testing.assert_close(result[TransitionKey.ACTION], original)
+
+
+# Training batch shape validation
+
+
+def test_relative_with_action_chunks(dataset, action_dim):
+    """Verify relative actions work correctly with (B, chunk_size, action_dim) shaped actions."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+
+    # Simulate a training batch: actions=(B, chunk_size, action_dim), state=(B, state_dim)
+    batch_actions = action_chunks[:4]  # (4, chunk_size, action_dim)
+    batch_states = states[:4]  # (4, state_dim)
+
+    mask = [True] * action_dim
+    relative = to_relative_actions(batch_actions, batch_states, mask)
+
+    # First action in each chunk should be close to zero (action[t] - state[t] ≈ small)
+    first_relatives = relative[:, 0, :]  # (B, action_dim)
+    assert first_relatives.abs().mean() < relative.abs().mean(), (
+        f"First action in chunk should have smaller relative offset than average. "
+        f"First: {first_relatives.abs().mean():.4f}, Average: {relative.abs().mean():.4f}"
+    )
+
+    # Later actions should have larger relative offsets
+    last_relatives = relative[:, -1, :]  # (B, action_dim)
+    assert last_relatives.abs().mean() >= first_relatives.abs().mean(), (
+        f"Last action in chunk should have >= relative offset than first. "
+        f"Last: {last_relatives.abs().mean():.4f}, First: {first_relatives.abs().mean():.4f}"
+    )
+
+    # Roundtrip
+    recovered = to_absolute_actions(relative, batch_states, mask)
+    torch.testing.assert_close(recovered, batch_actions)
+
+
+def test_relative_stats_match_actual_data_distribution(dataset, action_dim):
+    """Verify computed stats match the actual relative-action distribution."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    # Compute stats like the training script does
+    relative_stats = _compute_relative_chunk_stats(action_chunks, states, mask)
+
+    # Also compute directly
+    all_relatives = []
+    for actions, state in zip(action_chunks, states, strict=True):
+        rel = to_relative_actions(actions.unsqueeze(0), state.unsqueeze(0), mask).squeeze(0)
+        all_relatives.append(rel)
+    all_relatives_tensor = torch.cat(all_relatives, dim=0)
+
+    # Compare mean
+    actual_mean = all_relatives_tensor.mean(dim=0).numpy()
+    np.testing.assert_allclose(relative_stats["mean"], actual_mean, atol=0.01)
+
+    # Compare std
+    actual_std = all_relatives_tensor.std(dim=0).numpy()
+    np.testing.assert_allclose(relative_stats["std"], actual_std, atol=0.1)
+
+    # Verify q01 < mean < q99
+    assert (relative_stats["q01"] < relative_stats["mean"]).all(), "q01 should be < mean"
+    assert (relative_stats["mean"] < relative_stats["q99"]).all(), "mean should be < q99"
+
+
+def test_quantile_normalization_roundtrip(dataset, action_dim):
+    """Full roundtrip with QUANTILES normalization (what OpenPI uses for pi05)."""
+    action_chunks, states = _build_action_chunks(dataset, CHUNK_SIZE)
+    mask = [True] * action_dim
+
+    relative_stats = _compute_relative_chunk_stats(action_chunks, states, mask)
+    stats = {ACTION: dict(relative_stats.items())}
+
+    features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,))}
+    norm_map = {FeatureType.ACTION: NormalizationMode.QUANTILES}
+
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    unnormalizer = UnnormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    original_actions = action_chunks[0].unsqueeze(0)
+    state = states[0].unsqueeze(0)
+
+    batch = {ACTION: original_actions, OBS_STATE: state}
+    transition = batch_to_transition(batch)
+
+    # Forward: relative → quantile normalize
+    t1 = relative_step(transition)
+    t2 = normalizer(t1)
+
+    normalized = t2[TransitionKey.ACTION]
+    # Most values should be in [-1, 1] with quantile normalization
+    pct_in_range = (normalized.abs() < 2).float().mean()
+    assert pct_in_range > 0.5, f"Only {pct_in_range * 100:.1f}% in [-2, 2] after quantile norm, expected >50%"
+
+    # Reverse: unnormalize → absolute
+    t3 = unnormalizer(t2)
+    t4 = absolute_step(t3)
+
+    recovered = t4[TransitionKey.ACTION]
+    torch.testing.assert_close(recovered, original_actions, atol=1e-3, rtol=1e-3)
+
+
+def test_state_not_modified_by_relative_processor(dataset, action_dim):
+    """State should never be modified by the relative-action processor."""
+    hf = dataset.hf_dataset
+    batch = {
+        ACTION: torch.stack([hf[i]["action"] for i in range(4)]),
+        OBS_STATE: torch.stack([hf[i]["observation.state"] for i in range(4)]),
+    }
+    original_state = batch[OBS_STATE].clone()
+    transition = batch_to_transition(batch)
+
+    step = RelativeActionsProcessorStep(enabled=True)
+    result = step(transition)
+
+    result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
+    torch.testing.assert_close(result_state, original_state)