Initial pi0.6 minimal validation

2026-07-11 22:35:54 +08:00
commit 53ef020b46
27 changed files with 3459 additions and 0 deletions
@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import torch
+
+from pi06mini.experience import (
+    build_synthetic_experience,
+    discounted_returns,
+    n_step_advantages,
+    normalize_returns_to_unit_interval,
+    quantile_advantage_labels,
+)
+
+
+def _geometry() -> tuple[torch.Tensor, torch.Tensor]:
+    history = torch.stack((torch.linspace(-7.0, 0.0, 8), torch.zeros(8)), dim=-1)
+    x = torch.linspace(1.0, 16.0, 16)
+    future = torch.stack((x, 0.025 * x.square()), dim=-1)
+    return history, future
+
+
+def test_synthetic_buffer_has_all_experience_sources_and_failure_geometry() -> None:
+    history, future = _geometry()
+    episodes = build_synthetic_experience(history, future, episode_prefix="pai")
+
+    assert [episode.source_type for episode in episodes] == [
+        "demo_success",
+        "autonomous_understeer",
+        "autonomous_stop",
+        "intervention_correction",
+    ]
+    demo, understeer, stopped, correction = episodes
+    torch.testing.assert_close(demo.future_xy, future)
+    assert understeer.future_xy[-1, 1] < future[-1, 1]
+    assert torch.equal(stopped.future_xy[-1], stopped.future_xy[-2])
+    assert correction.intervention_mask.any()
+    assert not correction.intervention_mask[0]
+    torch.testing.assert_close(correction.future_xy[-1], future[-1])
+    assert torch.all(demo.rewards == 0)
+    assert float(understeer.rewards.min()) < 0
+    assert float(stopped.rewards.min()) < 0
+
+
+def test_returns_are_normalized_to_closed_negative_unit_interval() -> None:
+    rewards = torch.tensor([[0.0, 0.0, 0.0], [-0.2, -0.4, -1.0]])
+    returns = discounted_returns(rewards, gamma=1.0)
+    normalized = normalize_returns_to_unit_interval(returns)
+
+    assert normalized.shape == rewards.shape
+    assert float(normalized.min()) == -1.0
+    assert float(normalized.max()) == 0.0
+    assert torch.all(normalized <= 0.0)
+    assert torch.all(normalized >= -1.0)
+    assert torch.equal(
+        normalize_returns_to_unit_interval(torch.ones(3)), torch.zeros(3)
+    )
+
+
+def test_n_step_advantage_bootstraps_and_stops_at_terminal() -> None:
+    rewards = torch.tensor([1.0, 2.0, 3.0])
+    values = torch.tensor([0.5, 1.0, 1.5, 8.0])
+    dones = torch.tensor([False, True, False])
+    advantages = n_step_advantages(
+        rewards, values, n_step=3, gamma=0.5, dones=dones
+    )
+
+    # t=0 sees r0 + gamma*r1, then the terminal blocks the large bootstrap.
+    assert torch.isclose(advantages[0], torch.tensor(1.5))
+    assert torch.isclose(advantages[1], torch.tensor(1.0))
+    # t=2 reaches the explicit end and therefore has no bootstrap.
+    assert torch.isclose(advantages[2], torch.tensor(1.5))
+
+
+def test_quantile_labels_force_interventions_positive() -> None:
+    advantages = torch.tensor([-3.0, -2.0, -1.0, 2.0])
+    interventions = torch.tensor([True, False, False, False])
+    labels, threshold = quantile_advantage_labels(
+        advantages, quantile=0.75, intervention_mask=interventions
+    )
+
+    assert labels.tolist() == [True, False, False, True]
+    assert float(threshold) > -1.0
@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+
+from pi06mini.data import Pi06ExperienceDataset, 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_experience_dataset_contract_and_intervention_label(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 = Pi06ExperienceDataset(manifest, image_size=32)
+    assert len(dataset) == 4
+    assert {item.source_type for item in dataset.episodes} == {
+        "demo_success",
+        "autonomous_understeer",
+        "autonomous_stop",
+        "intervention_correction",
+    }
+    assert float(dataset.normalized_returns.min()) == -1.0
+    assert float(dataset.normalized_returns.max()) == 0.0
+
+    items = [dataset[index] for index in range(len(dataset))]
+    intervention = next(item for item in items if item["is_intervention"])
+    assert bool(intervention["positive_condition"])
+    assert torch.all(intervention["positive_steps"][intervention["intervention_mask"]])
+    assert intervention["image"].shape == (3, 32, 32)
+    assert intervention["state"].shape == (16,)
+    assert intervention["actions"].shape == (16, 2)
+    assert intervention["returns"].shape == (16,)
+    assert intervention["return_target"].shape == ()
+    assert intervention["value_target"].shape == ()
+    assert intervention["advantages"].shape == (16,)
+    assert intervention["advantage_label"] == "positive"
+    assert intervention["reasoning"] == "Steer left."
+    decoded = dataset.action_normalizer.decode(
+        intervention["actions"], intervention["history_xy"]
+    )
+    torch.testing.assert_close(decoded, intervention["future_xy"], atol=1e-5, rtol=1e-5)
@@ -0,0 +1,123 @@
+import torch
+
+from pi06mini.model import Pi06Config, Pi06Mini
+from pi06mini.tokenizer import MiniTokenizer
+
+
+def tiny_model() -> Pi06Mini:
+    tokenizer = MiniTokenizer()
+    return Pi06Mini(
+        Pi06Config(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            max_text_tokens=32,
+            max_action_tokens=4,
+            action_start_token_id=tokenizer.ACTION_START,
+        )
+    )
+
+
+def test_flow_reads_subtask_and_advantage_but_not_fast_targets() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
+    )
+    model = tiny_model()
+    input_ids = sequence.input_ids[None]
+    derived = model._derive_action_prefix_mask(input_ids)
+    torch.testing.assert_close(derived[0], sequence.action_prefix_mask)
+
+    text_offset = 1
+    prefix, flow = model._masks(
+        batch=1,
+        prefix_length=1 + input_ids.shape[1],
+        action_length=2,
+        text_offset=text_offset,
+        attention_mask=torch.ones_like(input_ids, dtype=torch.bool),
+        action_prefix_mask=derived,
+        generated_token_start=sequence.generated_token_start,
+        device=torch.device("cpu"),
+    )
+    generated_start = text_offset + sequence.generated_token_start
+    assert prefix[0, 0, :generated_start].all()
+    assert not prefix[0, 0, generated_start:].any()
+    # Flow sees image plus all text strictly before ACTION_START.
+    readable_prefix = text_offset + sequence.action_start_index
+    assert flow[0, :, :readable_prefix].all()
+    assert not flow[0, :, readable_prefix : 1 + input_ids.shape[1]].any()
+    assert flow[0, :, 1 + input_ids.shape[1] :].all()
+
+
+def test_flow_loss_is_knowledge_insulated_from_every_vlm_parameter() -> None:
+    torch.manual_seed(4)
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "steer_left", torch.zeros(2, 2), "positive"
+    )
+    model = tiny_model()
+    output = model(
+        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],
+        generated_token_start=sequence.generated_token_start,
+        target_velocity=torch.randn(1, 2, 2),
+    )
+    assert output.flow_loss is not None
+    output.flow_loss.backward()
+
+    vlm_gradients = [parameter.grad for parameter in model.vlm_parameters()]
+    assert vlm_gradients
+    assert all(gradient is None or torch.count_nonzero(gradient) == 0 for gradient in vlm_gradients)
+    action_gradients = [parameter.grad for parameter in model.action_parameters()]
+    assert any(gradient is not None and torch.count_nonzero(gradient) > 0 for gradient in action_gradients)
+
+
+def test_knowledge_insulation_stops_gradients_not_conditioning() -> None:
+    torch.manual_seed(8)
+    tokenizer = MiniTokenizer()
+    model = tiny_model().eval()
+    positive = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "positive"
+    )
+    negative = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "negative"
+    )
+    image = torch.randn(1, 3, 16, 16)
+    noise = torch.randn(1, 2, 2)
+    timestep = torch.full((1,), 0.5)
+    with torch.no_grad():
+        pos = model(
+            image, positive.input_ids[None], noise, timestep,
+            generated_token_start=positive.generated_token_start,
+        ).vector_field
+        neg = model(
+            image, negative.input_ids[None], noise, timestep,
+            generated_token_start=negative.generated_token_start,
+        ).vector_field
+    assert not torch.allclose(pos, neg)
+
+
+def test_joint_loss_still_trains_vlm_through_token_objective() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "stop", torch.zeros(2, 2), "negative"
+    )
+    model = tiny_model()
+    output = model(
+        torch.randn(1, 3, 16, 16),
+        sequence.input_ids[None],
+        torch.randn(1, 2, 2),
+        torch.rand(1),
+        generated_token_start=sequence.generated_token_start,
+        text_labels=sequence.labels[None],
+        target_velocity=torch.randn(1, 2, 2),
+    )
+    output.loss.backward()
+    assert model.token_embedding.weight.grad is not None
+    assert torch.count_nonzero(model.token_embedding.weight.grad) > 0
@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch
+
+from pi06mini.recap import (
+    apply_condition_dropout,
+    classifier_free_guidance,
+    condition_dropout_mask,
+)
+
+
+def test_default_condition_dropout_is_statistically_thirty_percent() -> None:
+    generator = torch.Generator().manual_seed(11)
+    mask = condition_dropout_mask(20_000, generator=generator)
+
+    assert mask.dtype == torch.bool
+    assert abs(float(mask.float().mean()) - 0.3) < 0.015
+
+
+def test_condition_dropout_replaces_whole_examples() -> None:
+    conditions = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+    dropped, mask = apply_condition_dropout(
+        conditions,
+        null_condition=-1.0,
+        mask=torch.tensor([False, True, False]),
+    )
+
+    assert mask.tolist() == [False, True, False]
+    torch.testing.assert_close(dropped[0], conditions[0])
+    torch.testing.assert_close(dropped[1], torch.tensor([-1.0, -1.0]))
+    torch.testing.assert_close(dropped[2], conditions[2])
+
+
+def test_cfg_interpolates_and_extrapolates_flow_predictions() -> None:
+    conditional = torch.tensor([[2.0, 4.0]])
+    unconditional = torch.tensor([[0.0, 2.0]])
+
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=0.0),
+        unconditional,
+    )
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=1.0),
+        conditional,
+    )
+    torch.testing.assert_close(
+        classifier_free_guidance(conditional, unconditional, guidance_scale=1.5),
+        torch.tensor([[3.0, 5.0]]),
+    )
@@ -0,0 +1,65 @@
+import torch
+
+from pi06mini.tokenizer import MiniTokenizer
+
+
+def test_advantage_condition_is_between_subtask_and_actions() -> None:
+    tokenizer = MiniTokenizer()
+    history = torch.tensor([[0.0, 0.0], [1.0, 0.2]])
+    actions = torch.tensor([[0.1, 0.0], [0.2, -0.1]])
+
+    positive = tokenizer.build_policy_sequence(
+        history, "steer_left", actions, "positive"
+    )
+    assert positive.advantage_index == positive.subtask_index + 1
+    assert positive.action_start_index == positive.advantage_index + 1
+    assert positive.input_ids[positive.advantage_index] == tokenizer.ADVANTAGE_POSITIVE
+    assert positive.input_ids[positive.action_start_index] == tokenizer.ACTION_START
+    assert positive.action_prefix_mask[: positive.action_start_index].all()
+    assert not positive.action_prefix_mask[positive.action_start_index :].any()
+
+    negative = tokenizer.build_policy_sequence(
+        history, "steer_left", actions, "negative"
+    )
+    assert negative.input_ids[negative.advantage_index] == tokenizer.ADVANTAGE_NEGATIVE
+
+
+def test_omitted_advantage_is_absent_not_a_third_condition_token() -> None:
+    tokenizer = MiniTokenizer()
+    sequence = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "keep_lane", torch.zeros(2, 2), "omitted"
+    )
+
+    assert sequence.advantage_index is None
+    assert sequence.action_start_index == sequence.subtask_index + 1
+    assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
+    assert tokenizer.ADVANTAGE_POSITIVE not in sequence.input_ids
+    assert tokenizer.ADVANTAGE_NEGATIVE not in sequence.input_ids
+
+
+def test_labels_predict_subtask_and_fast_but_not_external_advantage() -> None:
+    tokenizer = MiniTokenizer()
+    positive = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "stop", torch.zeros(2, 2), True
+    )
+    assert positive.labels[positive.subtask_index - 1] == positive.input_ids[positive.subtask_index]
+    assert positive.labels[positive.subtask_index] == -100
+    assert positive.labels[positive.advantage_index] == tokenizer.ACTION_START
+
+    omitted = tokenizer.build_policy_sequence(
+        torch.zeros(2, 2), "stop", torch.zeros(2, 2), None
+    )
+    assert omitted.labels[omitted.subtask_index] == 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
@@ -0,0 +1,73 @@
+import torch
+
+from pi06mini.model import Pi06Config, Pi06Mini
+from pi06mini.tokenizer import MiniTokenizer
+from pi06mini.value import DistributionalValueModel, ValueConfig
+
+
+def tiny_value_model() -> DistributionalValueModel:
+    tokenizer = MiniTokenizer()
+    return DistributionalValueModel(
+        ValueConfig(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            max_text_tokens=16,
+        )
+    )
+
+
+def test_distributional_value_has_201_bins_and_expected_value() -> None:
+    tokenizer = MiniTokenizer()
+    model = tiny_value_model()
+    state = tokenizer.build_value_input(torch.zeros(2, 2))[None].repeat(2, 1)
+    output = model(
+        torch.randn(2, 3, 16, 16),
+        state,
+        target_values=torch.tensor([-1.0, 0.0]),
+    )
+
+    assert output.logits.shape == (2, 201)
+    assert output.probabilities.shape == (2, 201)
+    torch.testing.assert_close(output.probabilities.sum(-1), torch.ones(2))
+    assert torch.all((-1.0 <= output.expectation) & (output.expectation <= 0.0))
+    assert output.target_bins.tolist() == [0, 200]
+    assert output.loss is not None and output.loss.ndim == 0
+
+
+def test_value_discretization_and_expectation_use_configured_support() -> None:
+    model = tiny_value_model()
+    bins = model.values_to_bins(torch.tensor([-2.0, -0.5, 1.0]))
+    assert bins.tolist() == [0, 100, 200]
+    torch.testing.assert_close(
+        model.bins_to_values(torch.tensor([0, 100, 200])),
+        torch.tensor([-1.0, -0.5, 0.0]),
+    )
+    logits = torch.full((3, 201), -100.0)
+    logits[0, 0] = logits[1, 100] = logits[2, 200] = 100.0
+    torch.testing.assert_close(
+        model.expectation_from_logits(logits),
+        torch.tensor([-1.0, -0.5, 0.0]),
+    )
+
+
+def test_value_model_and_policy_share_no_parameters() -> None:
+    tokenizer = MiniTokenizer()
+    policy = Pi06Mini(
+        Pi06Config(
+            vocab_size=tokenizer.vocab_size,
+            image_size=16,
+            patch_size=16,
+            d_model=32,
+            n_heads=4,
+            n_layers=1,
+            action_start_token_id=tokenizer.ACTION_START,
+        )
+    )
+    value = tiny_value_model()
+    policy_ids = {id(parameter) for parameter in policy.parameters()}
+    value_ids = {id(parameter) for parameter in value.parameters()}
+    assert policy_ids.isdisjoint(value_ids)