Initial pi0.7 minimal validation

2026-07-12 00:52:00 +08:00
commit b4c8db846b
27 changed files with 3724 additions and 0 deletions
@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import torch
+
+from pi07mini.context import (
+    PolicyMetadata,
+    build_ambiguous_history_pair,
+    drop_prompt_components,
+    endpoint_subgoal_heatmap,
+    format_policy_prompt,
+    make_history_proxy_frames,
+    prompt_components,
+)
+
+
+def _history() -> torch.Tensor:
+    x = torch.linspace(-8.0, 0.0, 8)
+    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_history_proxy_is_oldest_to_newest_and_future_free() -> None:
+    image = torch.linspace(0.0, 1.0, 3 * 20 * 20).reshape(3, 20, 20)
+    frames = make_history_proxy_frames(image, _history(), frames=6)
+    assert frames.shape == (6, 3, 20, 20)
+    torch.testing.assert_close(frames[-1], image * 0.98)
+    assert not torch.equal(frames[0], frames[-1])
+
+
+def test_ambiguous_histories_are_distinct_but_share_current_state() -> None:
+    pair = build_ambiguous_history_pair(_history(), frames=6)
+    assert pair.shape == (2, 6, 2)
+    torch.testing.assert_close(pair[0, -1], pair[1, -1])
+    assert not torch.allclose(pair[0, :-1], pair[1, :-1])
+
+
+def test_prompt_component_dropout_is_independent_and_deterministic() -> None:
+    metadata = PolicyMetadata("fast", "good", "none", "autonomous")
+    components = prompt_components(metadata, task="Drive safely.")
+    generator_a = torch.Generator().manual_seed(42)
+    generator_b = torch.Generator().manual_seed(42)
+    visible_a, mask_a = drop_prompt_components(
+        components, drop_probability=0.5, generator=generator_a
+    )
+    visible_b, mask_b = drop_prompt_components(
+        components, drop_probability=0.5, generator=generator_b
+    )
+    assert visible_a == visible_b
+    torch.testing.assert_close(mask_a, mask_b)
+    assert visible_a["task"] == "Drive safely."
+    assert mask_a.shape == (5,)
+    text = format_policy_prompt(visible_a)
+    assert "Drive safely." in text
+    for dropped in set(components) - set(visible_a):
+        assert f"{dropped}=" not in text
+
+
+def test_full_metadata_dropout_keeps_only_task() -> None:
+    components = prompt_components(PolicyMetadata("slow", "poor", "stop", "autonomous"))
+    visible, mask = drop_prompt_components(components, drop_probability=1.0)
+    assert set(visible) == {"task"}
+    assert mask.tolist() == [True, False, False, False, False]
@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi07mini.data import (
+    Pi07SteerableDataset,
+    build_policy_trajectories,
+    load_manifest,
+    materialize_subset,
+)
+
+
+def _make_source_project(root: Path) -> Path:
+    processed = root / "data" / "processed"
+    (processed / "images").mkdir(parents=True)
+    (processed / "trajectories").mkdir(parents=True)
+    pixels = np.zeros((40, 64, 3), dtype=np.uint8)
+    pixels[..., 0] = 80
+    pixels[..., 1] = np.arange(64, dtype=np.uint8)[None, :]
+    Image.fromarray(pixels).save(processed / "images" / "sample.jpg")
+    hx = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
+    history = np.stack((hx, 0.01 * hx, np.zeros(8)), axis=-1)
+    fx = np.linspace(1.0, 24.0, 16, dtype=np.float32)
+    future = np.stack((fx, 0.015 * fx**2, np.zeros(16)), axis=-1)
+    np.savez_compressed(
+        processed / "trajectories" / "sample.npz",
+        history_xyz=history,
+        future_xyz=future,
+    )
+    manifest = processed / "overfit_manifest.jsonl"
+    manifest.write_text(
+        json.dumps(
+            {
+                "clip_id": "clip-a",
+                "t0_us": 123,
+                "reasoning": "Steer left.",
+                "image_path": "data/processed/images/sample.jpg",
+                "trajectory_path": "data/processed/trajectories/sample.npz",
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    return manifest
+
+
+def test_policy_variants_and_metadata() -> None:
+    x = torch.linspace(-7.0, 0.0, 8)
+    history = torch.stack((x, torch.zeros_like(x)), -1)
+    fx = torch.linspace(1.0, 24.0, 16)
+    future = torch.stack((fx, 0.02 * fx.square()), -1)
+    policies = build_policy_trajectories(history, future)
+    assert [policy.kind for policy in policies] == [
+        "expert",
+        "fast",
+        "slow",
+        "understeer",
+        "stop",
+    ]
+    metadata = {policy.kind: policy.metadata.as_dict() for policy in policies}
+    assert metadata["expert"] == {
+        "speed": "normal",
+        "quality": "expert",
+        "mistake": "none",
+        "control": "human",
+    }
+    assert metadata["understeer"]["mistake"] == "understeer"
+    assert metadata["stop"]["quality"] == "poor"
+    assert policies[0].metadata.as_token_metadata() == {
+        "speed_steps": 1_000,
+        "quality": 5,
+        "mistake": False,
+        "control_mode": "ee",
+        "source_type": "demonstration",
+    }
+    assert torch.allclose(policies[-1].future_xy[-1], policies[-1].future_xy[4])
+
+
+def test_dataset_contract_and_roundtrip(tmp_path: Path) -> None:
+    source = _make_source_project(tmp_path / "source")
+    manifest = materialize_subset(source, tmp_path / "target")
+    assert load_manifest(manifest)[0].image_path.is_file()
+    dataset = Pi07SteerableDataset(
+        manifest,
+        image_size=32,
+        prompt_dropout=0.0,
+        max_rtc_delay_steps=12,
+    )
+    assert len(dataset) == 5
+    assert {policy.kind for policy in dataset.policies} == {
+        "expert",
+        "fast",
+        "slow",
+        "understeer",
+        "stop",
+    }
+    item = dataset[0]
+    assert item["pixel_values"].shape == (3, 32, 32)
+    assert item["history_pixel_values"].shape == (6, 3, 32, 32)
+    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["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["reasoning"] == "Steer left."
+    assert item["prompt_component_mask"].all()
+    assert item["token_metadata"]["quality"] == 5
+    assert item["control_mode"] == "ee"
+    decoded = dataset.action_normalizer.decode(item["actions"], item["history_xy"])
+    torch.testing.assert_close(decoded, item["future_xy"], atol=1e-5, rtol=1e-5)
+
+
+def test_dataset_is_deterministic_per_index(tmp_path: Path) -> None:
+    source = _make_source_project(tmp_path / "source")
+    dataset = Pi07SteerableDataset(
+        source,
+        image_size=24,
+        prompt_dropout=0.5,
+        seed=11,
+    )
+    first = dataset[3]
+    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"])
@@ -0,0 +1,63 @@
+import torch
+
+from pi07mini.memory import MEMConfig, MEMVideoEncoder
+
+
+def tiny_encoder() -> MEMVideoEncoder:
+    return MEMVideoEncoder(
+        MEMConfig(
+            image_size=16,
+            patch_size=8,
+            d_model=16,
+            n_heads=4,
+            n_layers=2,
+            max_history_frames=6,
+        )
+    ).eval()
+
+
+def test_mem_output_token_count_is_independent_of_history_length() -> None:
+    encoder = tiny_encoder()
+    current = torch.randn(2, 3, 16, 16)
+    with torch.no_grad():
+        no_history = encoder(current)
+        long_history = encoder(current, torch.randn(2, 6, 3, 16, 16))
+    assert no_history.shape == (2, 4, 16)
+    assert long_history.shape == no_history.shape
+
+
+def test_mem_is_sensitive_to_temporal_order() -> None:
+    torch.manual_seed(7)
+    encoder = tiny_encoder()
+    current = torch.randn(1, 3, 16, 16)
+    history = torch.stack(
+        (
+            torch.zeros(3, 16, 16),
+            torch.ones(3, 16, 16),
+            torch.full((3, 16, 16), 2.0),
+        )
+    )[None]
+    with torch.no_grad():
+        chronological = encoder(current, history)
+        reversed_history = encoder(current, history.flip(1))
+    assert not torch.allclose(chronological, reversed_history)
+
+
+def test_masked_history_frame_cannot_change_memory() -> None:
+    torch.manual_seed(9)
+    encoder = tiny_encoder()
+    current = torch.randn(1, 3, 16, 16)
+    history = torch.randn(1, 2, 3, 16, 16)
+    changed = history.clone()
+    changed[:, 0] = 1000.0
+    mask = torch.tensor([[False, True]])
+    with torch.no_grad():
+        first = encoder(current, history, mask)
+        second = encoder(current, changed, mask)
+    torch.testing.assert_close(first, second)
+
+
+def test_current_frame_temporal_position_is_exactly_zero() -> None:
+    encoding = MEMVideoEncoder.temporal_position(torch.tensor([3, 2, 1, 0]), 15)
+    assert torch.count_nonzero(encoding[-1]) == 0
+
@@ -0,0 +1,180 @@
+import torch
+
+from pi07mini.model import (
+    HighLevelPolicy,
+    Pi07Config,
+    Pi07Mini,
+    PrefixLayout,
+    TinyGoalGenerator,
+)
+from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer
+
+
+def tiny_model() -> Pi07Mini:
+    tokenizer = MiniTokenizer()
+    return Pi07Mini(
+        Pi07Config(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            memory_layers=1,
+            max_history_frames=3,
+            max_text_tokens=64,
+            max_action_tokens=4,
+            action_start_token_id=tokenizer.ACTION_START,
+        )
+    )
+
+
+def sample_sequence():
+    tokenizer = MiniTokenizer()
+    metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration")
+    return tokenizer, tokenizer.build_policy_sequence(
+        "drive", "steer_left", torch.zeros(2, 2), metadata
+    )
+
+
+def common_inputs(sequence):
+    return dict(
+        pixel_values=torch.randn(1, 3, 16, 16),
+        history_pixel_values=torch.randn(1, 2, 3, 16, 16),
+        history_states=torch.randn(1, 2, 2),
+        current_state=torch.randn(1, 2),
+        subgoal_pixel_values=torch.randn(1, 3, 16, 16),
+        input_ids=sequence.input_ids[None],
+        noisy_actions=torch.randn(1, 2, 2),
+        timesteps=torch.rand(1),
+        action_prefix_mask=sequence.action_prefix_mask[None],
+    )
+
+
+def test_block_causal_attention_matches_pi07_modal_order() -> None:
+    layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3)
+    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()
+
+
+def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None:
+    _, sequence = sample_sequence()
+    model = tiny_model()
+    inputs = common_inputs(sequence)
+    prefix, _, layout = model.encode_prefix(
+        inputs["pixel_values"],
+        inputs["history_pixel_values"],
+        inputs["history_states"],
+        inputs["current_state"],
+        inputs["input_ids"],
+        inputs["subgoal_pixel_values"],
+    )
+    assert layout.observation_tokens == 1 + 2 + 1  # MEM + history states + current
+    assert layout.goal_tokens == 1
+    assert prefix.shape[1] == layout.total_tokens
+
+
+def test_flow_cannot_read_teacher_forced_fast_tokens() -> None:
+    torch.manual_seed(12)
+    tokenizer, sequence = sample_sequence()
+    model = tiny_model().eval()
+    changed_ids = sequence.input_ids.clone()
+    changed_ids[sequence.action_start_index + 1 : -1] = tokenizer.ACTION_TOKEN_START
+    inputs = common_inputs(sequence)
+    with torch.no_grad():
+        first = model(**inputs).vector_field
+        inputs["input_ids"] = changed_ids[None]
+        second = model(**inputs).vector_field
+    torch.testing.assert_close(first, second)
+
+
+def test_flow_loss_is_knowledge_insulated_from_all_vlm_parameters() -> None:
+    torch.manual_seed(13)
+    _, sequence = sample_sequence()
+    model = tiny_model()
+    output = model(
+        **common_inputs(sequence), target_velocity=torch.randn(1, 2, 2)
+    )
+    assert output.flow_loss is not None
+    output.flow_loss.backward()
+    assert all(
+        parameter.grad is None or torch.count_nonzero(parameter.grad) == 0
+        for parameter in model.vlm_parameters()
+    )
+    assert any(
+        parameter.grad is not None and torch.count_nonzero(parameter.grad) > 0
+        for parameter in model.action_parameters()
+    )
+
+
+def test_fast_cross_entropy_still_trains_vlm() -> None:
+    _, sequence = sample_sequence()
+    model = tiny_model()
+    output = model(
+        **common_inputs(sequence), text_labels=sequence.labels[None]
+    )
+    assert output.text_loss is not None
+    output.text_loss.backward()
+    assert model.token_embedding.weight.grad is not None
+    assert torch.count_nonzero(model.token_embedding.weight.grad) > 0
+
+
+def test_visual_subgoal_conditions_flow_expert() -> None:
+    torch.manual_seed(14)
+    _, sequence = sample_sequence()
+    model = tiny_model().eval()
+    inputs = common_inputs(sequence)
+    with torch.no_grad():
+        first = model(**inputs).vector_field
+        inputs["subgoal_pixel_values"] = inputs["subgoal_pixel_values"] + 3.0
+        second = model(**inputs).vector_field
+    assert not torch.allclose(first, second)
+
+
+def test_metadata_cfg_uses_paper_extrapolation_formula() -> None:
+    conditioned = torch.tensor([1.0, 2.0])
+    unconditional = torch.tensor([0.5, 1.5])
+    torch.testing.assert_close(
+        Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 0.0), conditioned
+    )
+    torch.testing.assert_close(
+        Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 1.0),
+        torch.tensor([1.5, 2.5]),
+    )
+
+
+def test_high_level_policy_predicts_next_subtask_token() -> None:
+    model = HighLevelPolicy(tiny_model().config)
+    output = model(
+        pixel_values=torch.randn(2, 3, 16, 16),
+        history_pixel_values=torch.randn(2, 2, 3, 16, 16),
+        history_states=torch.randn(2, 2, 2),
+        current_state=torch.randn(2, 2),
+        input_ids=torch.tensor([[256, 257, 264], [256, 257, 264]]),
+        target_token=torch.tensor([271, 272]),
+    )
+    assert output.logits.shape == (2, tiny_model().config.vocab_size)
+    assert output.loss is not None
+
+
+def test_tiny_goal_generator_predicts_image_flow() -> None:
+    tokenizer = MiniTokenizer()
+    generator = TinyGoalGenerator(tokenizer.vocab_size)
+    current = torch.randn(2, 3, 16, 16)
+    noisy = torch.randn_like(current)
+    output = generator(
+        current,
+        torch.tensor([[256, 257, 264], [256, 257, 264]]),
+        noisy,
+        torch.rand(2),
+        target_velocity=torch.randn_like(current),
+    )
+    assert output.vector_field.shape == current.shape
+    assert output.loss is not None
@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from pi07mini.rtc import align_actions_for_rtc, 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:
+    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)
+
+
+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_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
+    generator = torch.Generator().manual_seed(9)
+    values = {
+        sample_rtc_delay(generator=generator, action_horizon=16) for _ in range(300)
+    }
+    assert min(values) == 0
+    assert max(values) == 12
+
+
+@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)
@@ -0,0 +1,65 @@
+import torch
+
+from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer
+
+
+def metadata() -> EpisodeMetadata:
+    return EpisodeMetadata(
+        speed_steps=2100,
+        quality=5,
+        mistake=False,
+        control_mode="ee",
+        source_type="autonomous_rollout",
+    )
+
+
+def test_prompt_contains_rich_metadata_and_control_tokens() -> None:
+    tokenizer = MiniTokenizer()
+    prompt, metadata_indices = tokenizer.encode_prompt(
+        "drive on the road", "steer_left", metadata()
+    )
+    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
+
+
+def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.zeros(3, 2)
+    conditioned, unconditional = tokenizer.build_metadata_cfg_pair(
+        "drive", "keep_lane", actions, metadata()
+    )
+    assert tokenizer.QUALITY_MARKER in conditioned.input_ids
+    assert tokenizer.QUALITY_MARKER not in unconditional.input_ids
+    assert tokenizer.CONTROL_EE in conditioned.input_ids
+    assert tokenizer.CONTROL_EE in unconditional.input_ids
+    assert conditioned.metadata_indices
+    assert not unconditional.metadata_indices
+
+
+def test_flow_prefix_ends_before_fast_stream() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        "drive", "keep_lane", torch.zeros(4, 2), metadata()
+    )
+    assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
+    assert sequence.action_prefix_mask[: sequence.action_start_index].all()
+    assert not sequence.action_prefix_mask[sequence.action_start_index :].any()
+    assert sequence.labels[sequence.action_start_index - 1] == tokenizer.ACTION_START
+
+
+def test_fast_lite_dct_round_trip_has_bounded_error() -> None:
+    tokenizer = MiniTokenizer()
+    actions = torch.tensor(
+        [[0.0, 0.0], [0.2, -0.1], [0.5, -0.2], [0.8, -0.1]],
+        dtype=torch.float32,
+    )
+    exact = tokenizer.idct_actions(tokenizer.dct_actions(actions))
+    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
+