Align pi0.7 minimal structure with paper

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()