feat(pi052): add training-time RTC

tests/policies/pi052/test_pi052_training_time_rtc.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+
+# Copyright 2026 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.
+
+"""Unit coverage for Pi052 training-time RTC conditioning."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _prepare_trained_rtc_prefix,
+    create_sinusoidal_pos_embedding,
+)
+from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
+from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
+    PI05Pytorch as PI052Pytorch,
+    _build_flow_matching_inputs,
+    _flow_loss_per_sample,
+)
+
+
+def test_training_rtc_uses_clean_prefix_and_per_token_time():
+    actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
+    noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
+    time = torch.tensor([0.25])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
+
+    assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
+    assert torch.equal(x_t[:, :2], actions[:, :2])
+    assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
+
+
+def test_training_rtc_loss_averages_over_postfix_only():
+    flow_loss = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    per_sample = _flow_loss_per_sample(flow_loss, prefix_mask)
+
+    assert per_sample.tolist() == [3.0]
+
+
+def test_per_token_time_embedding_preserves_action_axis():
+    time = torch.tensor([[0.0, 0.5, 1.0]])
+
+    embedding = create_sinusoidal_pos_embedding(time, 8, 4e-3, 4.0, time.device)
+
+    assert embedding.shape == (1, 3, 8)
+    assert not torch.equal(embedding[:, 0], embedding[:, 1])
+
+
+def test_action_expert_embeds_per_token_flow_times():
+    model = PI052Pytorch.__new__(PI052Pytorch)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace(chunk_size=3, min_period=4e-3, max_period=4.0)
+    model.gradient_checkpointing_enabled = False
+    model.action_in_proj = nn.Linear(2, 8)
+    model.time_mlp_in = nn.Linear(8, 8)
+    model.time_mlp_out = nn.Linear(8, 8)
+
+    suffix, _, _, adarms_cond = model.embed_suffix(
+        torch.randn(2, 3, 2),
+        torch.tensor([[0.0, 0.5, 0.5], [0.0, 0.0, 0.5]]),
+    )
+
+    assert suffix.shape == (2, 3, 8)
+    assert adarms_cond.shape == (2, 3, 8)
+
+
+def test_trained_rtc_prefix_is_padded_and_masked():
+    x_t = torch.randn(1, 5, 4)
+    previous = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+
+    prefix, mask = _prepare_trained_rtc_prefix(x_t, previous, inference_delay=2, training_max_delay=3)
+
+    assert prefix.shape == x_t.shape
+    assert mask.shape == x_t.shape
+    assert torch.equal(prefix[0, :2, :2], previous[:2])
+    assert torch.count_nonzero(prefix[0, :2, 2:]) == 0
+    assert mask[0, :2].all()
+    assert not mask[0, 2:].any()
+
+
+def test_trained_rtc_rejects_delay_outside_training_distribution():
+    with pytest.raises(ValueError, match="exceeds the checkpoint"):
+        _prepare_trained_rtc_prefix(
+            torch.randn(1, 5, 4),
+            torch.randn(4, 2),
+            inference_delay=4,
+            training_max_delay=3,
+        )
+
+
+@pytest.mark.parametrize("max_delay", [-1, 5])
+def test_pi052_config_rejects_invalid_training_rtc_delay(max_delay):
+    with pytest.raises(ValueError, match="rtc_training_max_delay"):
+        PI052Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)

tests/policies/rtc/test_configuration_rtc.py

@@ -16,6 +16,8 @@
 
 """Tests for RTC configuration module."""
 
+import pytest
+
 from lerobot.configs.types import RTCAttentionSchedule
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
 
@@ -27,6 +29,7 @@ def test_rtc_config_default_initialization():
     config = RTCConfig()
 
     assert config.enabled is True
+    assert config.mode == "guided"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
     assert config.max_guidance_weight == 10.0
     assert config.execution_horizon == 10
@@ -34,10 +37,16 @@ def test_rtc_config_default_initialization():
     assert config.debug_maxlen == 100
 
 
+def test_rtc_config_rejects_unknown_mode():
+    with pytest.raises(ValueError, match="mode must be"):
+        RTCConfig(mode="unknown")
+
+
 def test_rtc_config_custom_initialization():
     """Test RTCConfig initializes with custom values."""
     config = RTCConfig(
         enabled=True,
+        mode="trained",
         prefix_attention_schedule=RTCAttentionSchedule.EXP,
         max_guidance_weight=5.0,
         execution_horizon=20,
@@ -46,6 +55,7 @@ def test_rtc_config_custom_initialization():
     )
 
     assert config.enabled is True
+    assert config.mode == "trained"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP
     assert config.max_guidance_weight == 5.0
     assert config.execution_horizon == 20