seeker
101 lines
3.8 KiB
Python
101 lines
3.8 KiB
Python
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["value_histories"].shape == (17, 8, 2)
|
|
assert intervention["value_targets"].shape == (17,)
|
|
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)
|
|
|
|
|
|
def test_learned_value_predictions_drive_advantage_relabeling(tmp_path: Path) -> None:
|
|
manifest = _make_source_project(tmp_path / "source")
|
|
dataset = Pi06ExperienceDataset(manifest, image_size=32)
|
|
before = dataset.advantages.clone()
|
|
predicted = torch.zeros(len(dataset.episodes), dataset.action_horizon + 1)
|
|
predicted[:, :-1] = torch.linspace(-0.9, -0.1, dataset.action_horizon)
|
|
|
|
dataset.relabel_from_value_predictions(predicted)
|
|
|
|
assert not torch.allclose(dataset.advantages, before)
|
|
intervention_index = next(
|
|
index
|
|
for index, episode in enumerate(dataset.episodes)
|
|
if episode.source_type == "intervention_correction"
|
|
)
|
|
assert bool(dataset.positive_episodes[intervention_index])
|