Align pi0.7 minimal structure with paper

tests/test_context.py

@@ -6,8 +6,8 @@ from pi07mini.context import (
     PolicyMetadata,
     build_ambiguous_history_pair,
     drop_prompt_components,
-    endpoint_subgoal_heatmap,
     format_policy_prompt,
+    future_observation_proxy,
     make_history_proxy_frames,
     prompt_components,
 )
@@ -18,19 +18,13 @@ def _history() -> torch.Tensor:
     return torch.stack((x, 0.04 * x.square()), -1)
 
 
-def test_subgoal_heatmap_depends_on_endpoint_only() -> None:
-    endpoint = torch.tensor([20.0, 4.0])
-    first = endpoint_subgoal_heatmap(endpoint, size=24)
-    # Two future paths with radically different interiors produce exactly the
-    # same condition because only their shared endpoint is passed to the API.
-    path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
-    path_b = path_a.clone()
-    path_b[1:-1, 1] += 100.0
-    second = endpoint_subgoal_heatmap(path_b[-1], size=24)
-    torch.testing.assert_close(first, second, atol=0.0, rtol=0.0)
-    assert first.shape == (1, 24, 24)
-    assert float(first.max()) == 1.0
-    assert torch.isfinite(first).all()
+def test_future_subgoal_has_image_contract_not_coordinate_heatmap() -> None:
+    image = torch.linspace(0.0, 1.0, 3 * 24 * 24).reshape(3, 24, 24)
+    future = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
+    subgoal = future_observation_proxy(image, _history(), future)
+    assert subgoal.shape == image.shape
+    assert torch.isfinite(subgoal).all()
+    assert not torch.equal(subgoal, image)
 
 
 def test_history_proxy_is_oldest_to_newest_and_future_free() -> None:

tests/test_data.py

@@ -76,7 +76,6 @@ def test_policy_variants_and_metadata() -> None:
         "quality": 5,
         "mistake": False,
         "control_mode": "ee",
-        "source_type": "demonstration",
     }
     assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4])
 
@@ -105,12 +104,13 @@ def test_dataset_contract_and_roundtrip(tmp_path: Path) -> None:
     assert item["history_states"].shape == (6, 2)
     assert item["current_state"].shape == (2,)
     assert item["actions"].shape == (16, 2)
-    assert item["subgoal_heatmap"].shape == (1, 32, 32)
+    assert item["future_observation"].shape == (3, 32, 32)
     assert item["subgoal_pixel_values"].shape == (3, 32, 32)
     assert item["ambiguous_history_pair"].shape == (2, 6, 2)
     assert 0 <= item["rtc_delay_steps"] <= 12
-    assert item["rtc_target_actions"].shape == (16, 2)
-    assert item["rtc_valid_mask"].shape == (16,)
+    assert item["rtc_action_prefix"].shape == (16, 2)
+    assert item["rtc_prefix_mask"].shape == (16,)
+    assert item["rtc_postfix_mask"].shape == (16,)
     assert item["reasoning"] == "Steer left."
     assert item["prompt_component_mask"].all()
     assert item["token_metadata"]["quality"] == 5
@@ -131,4 +131,4 @@ def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None:
     second = dataset[3]
     assert first["prompt_components"] == second["prompt_components"]
     assert first["rtc_delay_steps"] == second["rtc_delay_steps"]
-    torch.testing.assert_close(first["rtc_target_actions"], second["rtc_target_actions"])
+    torch.testing.assert_close(first["rtc_action_prefix"], second["rtc_action_prefix"])

tests/test_model.py

@@ -31,7 +31,7 @@ def tiny_model() -> Pi07Mini:
 
 def sample_sequence():
     tokenizer = MiniTokenizer()
-    metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration")
+    metadata = EpisodeMetadata(2000, 5, False, "ee")
     return tokenizer, tokenizer.build_policy_sequence(
         "drive", "steer_left", torch.zeros(2, 2), metadata
     )
@@ -52,16 +52,17 @@ def common_inputs(sequence):
 
 
 def test_block_causal_attention_matches_pi07_modal_order() -> None:
-    layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3)
+    layout = PrefixLayout(
+        observation_tokens=2, prompt_tokens=2, goal_tokens=2, fast_tokens=2
+    )
     valid = torch.ones(1, layout.total_tokens, dtype=torch.bool)
     edges = Pi07Mini.build_prefix_attention_mask(layout, valid)[0]
     assert edges[0, :2].all()
-    assert not edges[0, 2:].any()  # observation cannot read goal or text
-    assert edges[2, :4].all()
-    assert not edges[2, 4:].any()  # goal reads observation + goal, not text
-    assert edges[4, :5].all()
-    assert not edges[4, 5:].any()  # first text token is causal
-    assert edges[6].all()
+    assert not edges[0, 2:].any()
+    assert edges[2, :3].all() and not edges[2, 3:].any()  # causal prompt
+    assert edges[4, :6].all() and not edges[4, 6:].any()  # goal reads obs + text + goal
+    assert edges[6, :7].all() and not edges[6, 7]  # FAST is causal after goal
+    assert edges[7].all()
 
 
 def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None:
@@ -78,9 +79,19 @@ def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None:
     )
     assert layout.observation_tokens == 1 + 2 + 1  # MEM + history states + current
     assert layout.goal_tokens == 1
+    assert layout.prompt_tokens == sequence.action_start_index
+    assert layout.fast_tokens == sequence.input_ids.numel() - sequence.action_start_index
     assert prefix.shape[1] == layout.total_tokens
 
 
+def test_action_expert_accepts_per_token_rtc_timesteps() -> None:
+    _, sequence = sample_sequence()
+    inputs = common_inputs(sequence)
+    inputs["timesteps"] = torch.tensor([[1.0, 0.25]])
+    output = tiny_model()(**inputs)
+    assert output.vector_field.shape == (1, 2, 2)
+
+
 def test_flow_cannot_read_teacher_forced_fast_tokens() -> None:
     torch.manual_seed(12)
     tokenizer, sequence = sample_sequence()

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)

tests/test_tokenizer.py

@@ -9,7 +9,6 @@ def metadata() -> EpisodeMetadata:
         quality=5,
         mistake=False,
         control_mode="ee",
-        source_type="autonomous_rollout",
     )
 
 
@@ -21,10 +20,18 @@ def test_prompt_contains_rich_metadata_and_control_tokens() -> None:
     assert tokenizer.SPEED_MARKER in prompt
     assert tokenizer.QUALITY_MARKER in prompt
     assert tokenizer.MISTAKE_MARKER in prompt
-    assert tokenizer.SOURCE_MARKER in prompt
     assert tokenizer.CONTROL_MARKER in prompt
     assert tokenizer.CONTROL_EE in prompt
-    assert len(metadata_indices) == 8
+    assert len(metadata_indices) == 6
+
+
+def test_high_level_prompt_contains_past_subtasks_and_next_query() -> None:
+    tokenizer = MiniTokenizer()
+    prompt = tokenizer.encode_high_level_prompt(
+        "drive", ("keep_lane", "slow_down"), metadata()
+    )
+    assert prompt[-1] == tokenizer.NEXT_SUBTASK_MARKER
+    assert torch.count_nonzero(prompt == tokenizer.SUBTASK_MARKER) == 2
 
 
 def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None:
@@ -62,4 +69,3 @@ def test_fast_lite_dct_round_trip_has_bounded_error() -> None:
     torch.testing.assert_close(exact, actions, atol=1e-5, rtol=1e-5)
     decoded = tokenizer.decode_actions(tokenizer.encode_actions(actions))
     assert torch.sqrt(torch.mean((decoded - actions) ** 2)) < 0.3
-