Align pi0.7 minimal structure with paper

tests/test_rtc.py

@@ -3,36 +3,39 @@ from __future__ import annotations
 import pytest
 import torch
 
-from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay
+from pi07mini.flow import FlowMatching
+from pi07mini.rtc import build_rtc_conditioning, sample_rtc_delay
 
 
 def _actions() -> torch.Tensor:
     return torch.arange(32, dtype=torch.float32).reshape(16, 2)
 
 
-def test_zero_delay_preserves_entire_chunk() -> None:
+def test_zero_delay_has_empty_prefix_and_full_postfix() -> None:
+    rtc = build_rtc_conditioning(_actions(), 0)
+    assert not rtc.prefix_mask.any()
+    assert rtc.postfix_mask.all()
+    assert torch.count_nonzero(rtc.action_prefix) == 0
+
+
+def test_delay_keeps_ground_truth_prefix_without_shifting_target() -> None:
     actions = _actions()
-    aligned = align_actions_for_rtc(actions, 0)
-    torch.testing.assert_close(aligned.delayed_target_actions, actions)
-    assert aligned.valid_mask.all()
-    assert aligned.executed_prefix_actions.shape == (0, 2)
+    rtc = build_rtc_conditioning(actions, 12)
+    torch.testing.assert_close(rtc.action_prefix[:12], actions[:12])
+    assert torch.count_nonzero(rtc.action_prefix[12:]) == 0
+    assert rtc.prefix_mask.tolist() == [True] * 12 + [False] * 4
+    assert rtc.postfix_mask.tolist() == [False] * 12 + [True] * 4
 
 
-def test_delay_shifts_target_and_masks_padded_tail() -> None:
-    actions = _actions()
-    aligned = align_actions_for_rtc(actions, 12)
-    torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:])
-    assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0
-    assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12
-    torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12])
-
-
-def test_rtc_alignment_preserves_batch_dimensions() -> None:
-    actions = torch.stack((_actions(), _actions() + 100.0))
-    aligned = align_actions_for_rtc(actions, 3)
-    assert aligned.delayed_target_actions.shape == (2, 16, 2)
-    assert aligned.valid_mask.shape == (2, 16)
-    torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:])
+def test_flow_training_sets_prefix_time_to_one_and_masks_its_loss() -> None:
+    actions = _actions()[None]
+    noisy, times, _, loss_mask = FlowMatching().training_batch(
+        actions, delays=torch.tensor([3]), generator=torch.Generator().manual_seed(4)
+    )
+    torch.testing.assert_close(noisy[:, :3], actions[:, :3])
+    torch.testing.assert_close(times[:, :3], torch.ones(1, 3))
+    assert not loss_mask[:, :3].any()
+    assert loss_mask[:, 3:].all()
 
 
 def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
@@ -47,4 +50,4 @@ def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
 @pytest.mark.parametrize("delay", [-1, 13, 16])
 def test_invalid_delay_is_rejected(delay: int) -> None:
     with pytest.raises(ValueError):
-        align_actions_for_rtc(_actions(), delay)
+        build_rtc_conditioning(_actions(), delay)