seeker
135 lines
4.5 KiB
Python
135 lines
4.5 KiB
Python
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"])
|