seeker
118 lines
4.1 KiB
Python
118 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from pi0mini.data import ActionNormalizer, Pi0OverfitDataset, load_manifest, materialize_subset
|
|
|
|
|
|
def _make_source_project(root: Path) -> Path:
|
|
processed = root / "data" / "processed"
|
|
image_dir = processed / "images"
|
|
trajectory_dir = processed / "trajectories"
|
|
image_dir.mkdir(parents=True)
|
|
trajectory_dir.mkdir(parents=True)
|
|
|
|
pixels = np.zeros((48, 80, 3), dtype=np.uint8)
|
|
pixels[..., 0] = 64
|
|
pixels[..., 1] = np.arange(80, dtype=np.uint8)[None, :]
|
|
Image.fromarray(pixels).save(image_dir / "sample.jpg")
|
|
|
|
history_x = np.linspace(-7.0, 0.0, 8, dtype=np.float32)
|
|
history = np.stack((history_x, 0.02 * history_x, np.zeros(8)), axis=-1)
|
|
future_x = np.linspace(1.0, 32.0, 16, dtype=np.float32)
|
|
future = np.stack((future_x, 0.002 * future_x**2, np.zeros(16)), axis=-1)
|
|
np.savez_compressed(
|
|
trajectory_dir / "sample.npz", history_xyz=history, future_xyz=future
|
|
)
|
|
|
|
manifest = processed / "overfit_manifest.jsonl"
|
|
manifest.write_text(
|
|
json.dumps(
|
|
{
|
|
"clip_id": "clip-a",
|
|
"t0_us": 1234,
|
|
"image_path": "data/processed/images/sample.jpg",
|
|
"trajectory_path": "data/processed/trajectories/sample.npz",
|
|
}
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
return manifest
|
|
|
|
|
|
def test_action_normalizer_is_reversible() -> None:
|
|
history = torch.tensor([[-1.0, 0.0], [0.0, 0.1]])
|
|
future = torch.tensor([[1.0, 0.2], [2.5, 0.4], [4.0, 0.7]])
|
|
normalizer = ActionNormalizer.fit([history], [future])
|
|
actions = normalizer.encode(history, future)
|
|
|
|
assert actions.shape == (3, 2)
|
|
assert torch.allclose(normalizer.decode(actions, history), future, atol=1e-6)
|
|
restored = ActionNormalizer.from_dict(normalizer.as_dict())
|
|
assert torch.allclose(restored.decode(actions, history), future, atol=1e-6)
|
|
|
|
|
|
def test_materialized_dataset_contract_and_determinism(tmp_path: Path) -> None:
|
|
source_manifest = _make_source_project(tmp_path / "source")
|
|
copied_manifest = materialize_subset(source_manifest, tmp_path / "target")
|
|
records = load_manifest(copied_manifest)
|
|
|
|
assert len(records) == 1
|
|
assert records[0].image_path.is_file()
|
|
assert records[0].trajectory_path.is_file()
|
|
|
|
dataset = Pi0OverfitDataset(copied_manifest, seed=19)
|
|
repeated = Pi0OverfitDataset(
|
|
copied_manifest,
|
|
seed=19,
|
|
action_statistics=dataset.action_normalizer.as_dict(),
|
|
)
|
|
assert len(dataset) == 4
|
|
assert dataset.state_dim == 16
|
|
|
|
original = dataset[0]
|
|
augmented = dataset[1]
|
|
repeated_augmented = repeated[1]
|
|
assert set(original) == {
|
|
"image",
|
|
"state",
|
|
"actions",
|
|
"history_xy",
|
|
"future_xy",
|
|
"sample_id",
|
|
}
|
|
assert original["image"].shape == (3, 64, 64)
|
|
assert original["state"].shape == (16,)
|
|
assert original["actions"].shape == (16, 2)
|
|
assert original["history_xy"].shape == (8, 2)
|
|
assert original["future_xy"].shape == (16, 2)
|
|
assert original["sample_id"].endswith(":aug0")
|
|
assert augmented["sample_id"].endswith(":aug1")
|
|
assert not torch.equal(original["future_xy"], augmented["future_xy"])
|
|
assert torch.equal(augmented["image"], repeated_augmented["image"])
|
|
assert torch.equal(augmented["state"], repeated_augmented["state"])
|
|
assert torch.equal(augmented["actions"], repeated_augmented["actions"])
|
|
|
|
decoded = dataset.action_normalizer.decode(
|
|
original["actions"], original["history_xy"]
|
|
)
|
|
assert torch.allclose(decoded, original["future_xy"], atol=1e-5)
|
|
|
|
|
|
def test_statistics_can_be_saved(tmp_path: Path) -> None:
|
|
source_manifest = _make_source_project(tmp_path / "source")
|
|
manifest = materialize_subset(source_manifest, tmp_path / "target")
|
|
dataset = Pi0OverfitDataset(manifest)
|
|
path = dataset.save_action_statistics(tmp_path / "stats.json")
|
|
values = json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
assert len(values["mean"]) == 2
|
|
assert len(values["scale"]) == 2
|
|
assert min(values["scale"]) >= 0.05
|