commit ffbcca8d2b6d3e2edb908c332ec95632e2fd028e Author: seeker Date: Sat Jul 11 20:04:18 2026 +0800 Initial pi0.5 minimal validation diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ff92181 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.ruff_cache/ +*.egg-info/ +.venv/ +data/ +outputs/**/*.pt +.DS_Store diff --git a/README.md b/README.md new file mode 100644 index 0000000..b40f44c --- /dev/null +++ b/README.md @@ -0,0 +1,141 @@ +# π0.5 Minimal Validation + +这是一个 **π0.5 风格的最小、可运行架构验证项目**。它从相邻 Alpamayo 验证项目准备的 PAI 子集中读取 1 个驾驶样本，生成 4 个确定性视觉增强变体做过拟合，目标是检查高层语义预测、blockwise attention、action expert 与 flow matching action chunk 能否组成完整训练和推理链路。 + +它不是 Physical Intelligence 官方代码，也不声称复现论文中的规模、数据或机器人能力。 + +## 验证内容 + +```text +图像 + 本体状态 + │ + ▼ +小型视觉/语言主干 ──→ 高层语义 subtask + │ │ + └──── semantic prefix┘ + │ + noisy action + time │ + └──────────▼ + action expert（双向 action attention） + │ + ▼ + flow velocity → action chunk +``` + +默认任务使用 `64×64` 前视相机图像、离散化历史 XY 状态和一句 subtask，预测未来 `16×2` 的二维轨迹。样本少且重复训练，因此成功标准是“能够记忆训练集”，而不是泛化。 + +## 与 π0.5 的对应关系 + +保留的关键思想： + +- 视觉/语言语义主干与独立 action expert； +- observation/text 与 action token 采用 blockwise attention 信息流； +- action token 内部采用双向 attention，一次预测整个 action chunk； +- flow matching 训练连续动作，推理时从高斯噪声做 Euler 积分； +- timestep 通过 adaptive RMSNorm 调制 action expert； +- 同时优化高层语义目标和低层连续动作目标。 +- VLM 自回归输出高层 subtask 以及 64-bin 离散动作 token； +- 推理时分别还原 VLM token 轨迹和 flow/diffusion 轨迹，与真值三方比较。 + +为了单机快速验证而简化的部分： + +- 小型自定义视觉编码器和 Transformer 取代 PaliGemma 2B； +- 隐藏宽度、层数、图像尺寸和 action horizon 大幅缩小； +- 使用固定小词表的 subtask 分类/生成，而非开放式语言理解； +- 使用二维驾驶轨迹，取代多机器人、多相机和高维关节控制； +- 默认使用 1 个 PAI 驾驶样本的 4 个视觉增强版本，没有机器人或网页联合训练数据； +- 使用正交 DCT-II、低频优先排列和 64-bin 量化实现 FAST-lite；保留全部频率系数，因此只验证表示与训练接口，不复现官方压缩率； +- Euler sampler 用固定步长，不包含面向部署的实时控制与安全系统。 + +因此，本项目适合验证架构接口和数学链路，不适合比较论文指标。 + +## 安装 + +建议使用 Python 3.10 或更新版本： + +```bash +cd pi05-minimal-validation +python -m venv .venv +source .venv/bin/activate +python -m pip install -U pip +python -m pip install -e '.[dev]' +``` + +默认实现只依赖 PyTorch、NumPy、Pillow、Matplotlib 和 PyYAML，不需要下载预训练权重。默认配置会读取： + +```text +../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl +``` + +如果该路径不同，请先修改 `configs/overfit.yaml` 的 `data.source_manifest`。 + +## 运行测试 + +```bash +cd pi05-minimal-validation +pytest -q +``` + +`tests/test_flow.py` 单独验证： + +- 直线路径 `x_t=(1-t)·noise+t·action`； +- 目标速度 `v=action-noise`； +- velocity MSE； +- 固定目标的精确 vector field 经 Euler 积分后能到达目标； +- 使用相同 `torch.Generator` 时采样可复现。 + +## 过拟合训练 + +```bash +cd pi05-minimal-validation +pi05mini-train --config configs/overfit.yaml +``` + +也可以直接使用模块入口： + +```bash +python -m pi05mini.train --config configs/overfit.yaml +``` + +配置默认先进行 300 步离散 VLM/FAST-lite 预训练，再进行 500 步联合 post-training。Stage 1 不更新 flow action expert；Stage 2 联合优化 token CE 与 flow loss。CPU 可以运行；Apple Silicon 可自动选择 MPS，CUDA 可自动选择 GPU。 + +## 推理与效果图 + +```bash +pi05mini-infer --checkpoint outputs/overfit/checkpoint.pt +``` + +输出图 `outputs/overfit/trajectory_comparison.png` 会叠加 ground truth、VLM 离散动作 token 解码轨迹和 diffusion/flow 轨迹；逐样本指标写入 `evaluation.json`。训练过程另存为 `training_history.json`，训练摘要为 `summary.json`。 + +## 配置重点 + +`configs/overfit.yaml` 中最常调整的参数： + +- `data.num_variants`：默认 4，刻意保持很小； +- `model.d_model`、`n_layers`：控制最小模型容量； +- `train.steps`：过拟合不充分时增加； +- `train.semantic_loss_weight` 与 `flow_loss_weight`：两个训练目标的权重； +- `inference.flow_steps`：Euler 积分步数，越大通常越慢且更平滑。 + +## 安全限制 + +- 该模型只在合成样本上验证记忆能力，输出不具备现实世界有效性。 +- 不得用于车辆、机器人、机械臂或任何会造成物理影响的设备控制。 +- 不得把 loss 下降或单样本轨迹重合解释为安全性、鲁棒性或泛化能力。 +- 真实部署至少还需要独立验证集、分布外和扰动测试、控制约束、碰撞检查、故障降级、急停以及人工监督。 +- checkpoint 和配置可能加载本地路径；只运行可信来源的文件，不加载未知 PyTorch pickle checkpoint。 + +## 项目结构 + +```text +pi05-minimal-validation/ +├── configs/overfit.yaml # 4 样本过拟合配置 +├── src/pi05mini/ # 模型、flow、数据、训练与推理 +├── tests/test_flow.py # flow matching 数学测试 +├── pyproject.toml +└── README.md +``` + +## 预期结论 + +若最小实验成功，它只能说明：π0.5 风格的语义主干、blockwise 双专家交互和 flow action head 在此实现中能够端到端反向传播并记忆几个样本。它不能证明实现与官方模型数值等价，也不能证明在真实驾驶或机器人任务中有效。 diff --git a/configs/overfit.yaml b/configs/overfit.yaml new file mode 100644 index 0000000..7abdc7c --- /dev/null +++ b/configs/overfit.yaml @@ -0,0 +1,32 @@ +seed: 7 +device: auto +output_dir: outputs/overfit + +data: + source_manifest: ../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl + num_variants: 4 + image_size: 64 + action_horizon: 16 + +model: + d_model: 96 + n_heads: 4 + n_layers: 3 + mlp_ratio: 3 + patch_size: 16 + max_text_tokens: 64 + dropout: 0.0 + +train: + pretrain_steps: 300 + posttrain_steps: 500 + batch_size: 4 + learning_rate: 0.001 + weight_decay: 0.0 + grad_clip_norm: 1.0 + semantic_loss_weight: 1.0 + flow_loss_weight: 2.0 + log_every: 20 + +inference: + flow_steps: 10 diff --git a/outputs/overfit/evaluation.json b/outputs/overfit/evaluation.json new file mode 100644 index 0000000..efb31eb --- /dev/null +++ b/outputs/overfit/evaluation.json @@ -0,0 +1,34 @@ +[ + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug0", + "target_subtask": "steer_left", + "predicted_subtask": "steer_left", + "vlm_action_tokens_valid": true, + "vlm_trajectory_mse": 0.07854792475700378, + "diffusion_trajectory_mse": 0.27604275941848755 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug1", + "target_subtask": "steer_left", + "predicted_subtask": "steer_left", + "vlm_action_tokens_valid": true, + "vlm_trajectory_mse": 0.025470610707998276, + "diffusion_trajectory_mse": 0.20049458742141724 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug2", + "target_subtask": "steer_left", + "predicted_subtask": "steer_left", + "vlm_action_tokens_valid": true, + "vlm_trajectory_mse": 0.013774838298559189, + "diffusion_trajectory_mse": 0.03129641339182854 + }, + { + "sample_id": "314f59c2-e420-4fcd-9df2-a51d572e0c20_1503899:aug3", + "target_subtask": "steer_left", + "predicted_subtask": "steer_left", + "vlm_action_tokens_valid": true, + "vlm_trajectory_mse": 0.059195056557655334, + "diffusion_trajectory_mse": 0.11321215331554413 + } +] \ No newline at end of file diff --git a/outputs/overfit/summary.json b/outputs/overfit/summary.json new file mode 100644 index 0000000..887b413 --- /dev/null +++ b/outputs/overfit/summary.json @@ -0,0 +1,19 @@ +{ + "device": "mps", + "samples": 4, + "parameters": 816482, + "initial": { + "stage": "pretrain", + "step": 1, + "loss": 78.87449645996094, + "text_loss": 78.87449645996094, + "flow_loss": null + }, + "final": { + "stage": "posttrain", + "step": 500, + "loss": 0.12353827804327011, + "text_loss": 3.1691649837739533e-06, + "flow_loss": 0.06176755577325821 + } +} \ No newline at end of file diff --git a/outputs/overfit/training_history.json b/outputs/overfit/training_history.json new file mode 100644 index 0000000..fdb3275 --- /dev/null +++ b/outputs/overfit/training_history.json @@ -0,0 +1,296 @@ +[ + { + "stage": "pretrain", + "step": 1, + "loss": 78.87449645996094, + "text_loss": 78.87449645996094, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 20, + "loss": 2.475086212158203, + "text_loss": 2.475086212158203, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 40, + "loss": 1.1025534868240356, + "text_loss": 1.1025534868240356, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 60, + "loss": 0.601564884185791, + "text_loss": 0.601564884185791, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 80, + "loss": 0.25571808218955994, + "text_loss": 0.25571808218955994, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 100, + "loss": 0.03314117714762688, + "text_loss": 0.03314117714762688, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 120, + "loss": 0.005310459993779659, + "text_loss": 0.005310459993779659, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 140, + "loss": 3.0251239877543412e-05, + "text_loss": 3.0251239877543412e-05, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 160, + "loss": 1.1354763955750968e-05, + "text_loss": 1.1354763955750968e-05, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 180, + "loss": 6.638075319642667e-06, + "text_loss": 6.638075319642667e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 200, + "loss": 5.348983449948719e-06, + "text_loss": 5.348983449948719e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 220, + "loss": 4.736786195280729e-06, + "text_loss": 4.736786195280729e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 240, + "loss": 4.316162630857434e-06, + "text_loss": 4.316162630857434e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 260, + "loss": 3.975574600190157e-06, + "text_loss": 3.975574600190157e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 280, + "loss": 3.6877775073662633e-06, + "text_loss": 3.6877775073662633e-06, + "flow_loss": null + }, + { + "stage": "pretrain", + "step": 300, + "loss": 3.4408508327032905e-06, + "text_loss": 3.4408508327032905e-06, + "flow_loss": null + }, + { + "stage": "posttrain", + "step": 1, + "loss": 8.707193374633789, + "text_loss": 3.428078571232618e-06, + "flow_loss": 4.353594779968262 + }, + { + "stage": "posttrain", + "step": 20, + "loss": 3.533726453781128, + "text_loss": 0.14151602983474731, + "flow_loss": 1.6961052417755127 + }, + { + "stage": "posttrain", + "step": 40, + "loss": 3.0850143432617188, + "text_loss": 0.00014649856893811375, + "flow_loss": 1.5424339771270752 + }, + { + "stage": "posttrain", + "step": 60, + "loss": 1.3556954860687256, + "text_loss": 0.00011281041224719957, + "flow_loss": 0.6777913570404053 + }, + { + "stage": "posttrain", + "step": 80, + "loss": 1.4165724515914917, + "text_loss": 6.599419066333212e-06, + "flow_loss": 0.7082829475402832 + }, + { + "stage": "posttrain", + "step": 100, + "loss": 0.8609042167663574, + "text_loss": 8.11900736152893e-06, + "flow_loss": 0.430448055267334 + }, + { + "stage": "posttrain", + "step": 120, + "loss": 0.5642163753509521, + "text_loss": 7.576879397674929e-06, + "flow_loss": 0.28210440278053284 + }, + { + "stage": "posttrain", + "step": 140, + "loss": 0.7110214233398438, + "text_loss": 6.249698344618082e-06, + "flow_loss": 0.35550758242607117 + }, + { + "stage": "posttrain", + "step": 160, + "loss": 0.18725578486919403, + "text_loss": 2.876289045161684e-06, + "flow_loss": 0.09362645447254181 + }, + { + "stage": "posttrain", + "step": 180, + "loss": 0.33923929929733276, + "text_loss": 6.748496616637567e-06, + "flow_loss": 0.16961628198623657 + }, + { + "stage": "posttrain", + "step": 200, + "loss": 0.3003162443637848, + "text_loss": 9.495882295595948e-06, + "flow_loss": 0.15015336871147156 + }, + { + "stage": "posttrain", + "step": 220, + "loss": 0.23716998100280762, + "text_loss": 5.660519946104614e-06, + "flow_loss": 0.11858215928077698 + }, + { + "stage": "posttrain", + "step": 240, + "loss": 0.42717865109443665, + "text_loss": 8.400202204938978e-06, + "flow_loss": 0.21358512341976166 + }, + { + "stage": "posttrain", + "step": 260, + "loss": 0.11702287942171097, + "text_loss": 4.122769951209193e-06, + "flow_loss": 0.058509379625320435 + }, + { + "stage": "posttrain", + "step": 280, + "loss": 0.19245445728302002, + "text_loss": 5.941245035501197e-06, + "flow_loss": 0.09622425585985184 + }, + { + "stage": "posttrain", + "step": 300, + "loss": 0.18475297093391418, + "text_loss": 4.744139005197212e-06, + "flow_loss": 0.09237411618232727 + }, + { + "stage": "posttrain", + "step": 320, + "loss": 0.38481318950653076, + "text_loss": 3.1461640901397914e-06, + "flow_loss": 0.19240501523017883 + }, + { + "stage": "posttrain", + "step": 340, + "loss": 0.2163764089345932, + "text_loss": 4.1934181354008615e-06, + "flow_loss": 0.10818611085414886 + }, + { + "stage": "posttrain", + "step": 360, + "loss": 0.23533868789672852, + "text_loss": 3.3019753118423978e-06, + "flow_loss": 0.11766768991947174 + }, + { + "stage": "posttrain", + "step": 380, + "loss": 0.08397223055362701, + "text_loss": 8.292082384286914e-06, + "flow_loss": 0.04198196902871132 + }, + { + "stage": "posttrain", + "step": 400, + "loss": 0.13249839842319489, + "text_loss": 1.0871653103095014e-05, + "flow_loss": 0.06624376028776169 + }, + { + "stage": "posttrain", + "step": 420, + "loss": 0.4479130506515503, + "text_loss": 1.0128642315976322e-05, + "flow_loss": 0.22395145893096924 + }, + { + "stage": "posttrain", + "step": 440, + "loss": 0.22416232526302338, + "text_loss": 2.6114355478057405e-06, + "flow_loss": 0.11207985877990723 + }, + { + "stage": "posttrain", + "step": 460, + "loss": 0.06346385926008224, + "text_loss": 3.1035808660817565e-06, + "flow_loss": 0.031730376183986664 + }, + { + "stage": "posttrain", + "step": 480, + "loss": 0.385486900806427, + "text_loss": 2.184032382501755e-06, + "flow_loss": 0.19274236261844635 + }, + { + "stage": "posttrain", + "step": 500, + "loss": 0.12353827804327011, + "text_loss": 3.1691649837739533e-06, + "flow_loss": 0.06176755577325821 + } +] \ No newline at end of file diff --git a/outputs/overfit/trajectory_comparison.png b/outputs/overfit/trajectory_comparison.png new file mode 100644 index 0000000..dde41ea Binary files /dev/null and b/outputs/overfit/trajectory_comparison.png differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..65f0665 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,38 @@ +[build-system] +requires = ["setuptools>=69", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "pi05-minimal-validation" +version = "0.1.0" +description = "A tiny, educational pi0.5-style hierarchical VLA overfitting experiment" +readme = "README.md" +requires-python = ">=3.10" +license = { text = "MIT" } +authors = [{ name = "pi05 minimal validation contributors" }] +dependencies = [ + "matplotlib>=3.8", + "numpy>=1.26", + "Pillow>=10.0", + "PyYAML>=6.0", + "torch>=2.2", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0", +] + +[project.scripts] +pi05mini-train = "pi05mini.train:main" +pi05mini-infer = "pi05mini.infer:main" + +[tool.setuptools] +package-dir = { "" = "src" } + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra" diff --git a/src/pi05mini/__init__.py b/src/pi05mini/__init__.py new file mode 100644 index 0000000..9be415e --- /dev/null +++ b/src/pi05mini/__init__.py @@ -0,0 +1,5 @@ +"""Minimal pi0.5 architecture validation package.""" + +from .model import Pi05Config, Pi05Mini, Pi05Output, flow_matching_sample + +__all__ = ["Pi05Config", "Pi05Mini", "Pi05Output", "flow_matching_sample"] diff --git a/src/pi05mini/data.py b/src/pi05mini/data.py new file mode 100644 index 0000000..6080d5d --- /dev/null +++ b/src/pi05mini/data.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +import json +import shutil +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image, ImageEnhance +from torch.utils.data import Dataset + +from .tokenizer import MiniTokenizer + + +@dataclass(frozen=True) +class SampleRecord: + clip_id: str + t0_us: int + reasoning: str + image_path: Path + trajectory_path: Path + + @property + def sample_id(self) -> str: + return f"{self.clip_id}_{self.t0_us}" + + +@dataclass(frozen=True) +class ActionStatistics: + mean: torch.Tensor + scale: torch.Tensor + + def as_dict(self) -> dict[str, list[float]]: + return { + "mean": [float(value) for value in self.mean], + "scale": [float(value) for value in self.scale], + } + + +class ActionNormalizer: + """Normalize 2-D waypoint deltas and retain invertible dataset statistics.""" + + def __init__(self, mean: torch.Tensor, scale: torch.Tensor) -> None: + self.mean = torch.as_tensor(mean, dtype=torch.float32).reshape(2) + self.scale = torch.as_tensor(scale, dtype=torch.float32).reshape(2).clamp_min(0.05) + + @classmethod + def fit(cls, histories: Sequence[torch.Tensor], futures: Sequence[torch.Tensor]) -> ActionNormalizer: + if not histories or len(histories) != len(futures): + raise ValueError("histories and futures must be non-empty sequences of equal length") + chunks = [] + for history, future in zip(histories, futures, strict=True): + chunks.append(cls.to_deltas(history, future).reshape(-1, 2)) + flat = torch.cat(chunks, dim=0).float() + return cls(flat.mean(dim=0), flat.std(dim=0, unbiased=False)) + + @classmethod + def from_dict(cls, values: dict[str, Sequence[float]]) -> ActionNormalizer: + return cls(torch.tensor(values["mean"]), torch.tensor(values["scale"])) + + @property + def statistics(self) -> ActionStatistics: + return ActionStatistics(self.mean.clone(), self.scale.clone()) + + @staticmethod + def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor: + previous = torch.cat((history_xy[-1:].float(), future_xy[:-1].float()), dim=0) + return future_xy.float() - previous + + def encode(self, history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor: + return (self.to_deltas(history_xy, future_xy) - self.mean) / self.scale + + def decode(self, actions: torch.Tensor, history_xy: torch.Tensor) -> torch.Tensor: + deltas = actions.float() * self.scale + self.mean + return history_xy[-1:].float() + torch.cumsum(deltas, dim=0) + + def as_dict(self) -> dict[str, list[float]]: + return self.statistics.as_dict() + + def save(self, path: str | Path) -> Path: + output = Path(path).expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(self.as_dict(), indent=2) + "\n", encoding="utf-8") + return output + + +def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path: + path = Path(raw_path).expanduser() + if path.is_absolute(): + return path + candidates = [manifest_path.parent / path] + if len(manifest_path.parents) >= 3: + candidates.insert(0, manifest_path.parents[2] / path) + for candidate in candidates: + if candidate.exists(): + return candidate.resolve() + return candidates[0].resolve() + + +def load_manifest(path: str | Path) -> list[SampleRecord]: + manifest_path = Path(path).expanduser().resolve() + records: list[SampleRecord] = [] + for line in manifest_path.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + row = json.loads(line) + records.append( + SampleRecord( + clip_id=str(row["clip_id"]), + t0_us=int(row["t0_us"]), + reasoning=str(row.get("reasoning", "")), + image_path=_resolve_data_path(manifest_path, row["image_path"]), + trajectory_path=_resolve_data_path(manifest_path, row["trajectory_path"]), + ) + ) + if not records: + raise ValueError(f"no samples found in {manifest_path}") + return records + + +def materialize_subset( + source_manifest: str | Path, + output_dir: str | Path, + max_samples: int = 3, + num_variants: int | None = None, +) -> Path: + """Copy a self-contained PAI overfit subset out of the Alpamayo project.""" + + # ``num_variants`` is accepted as a compatibility alias used by the + # training entry point. Actual variants stay virtual and deterministic in + # Pi05OverfitDataset, avoiding redundant copies of the large JPEG. + if num_variants is not None: + if num_variants < 1: + raise ValueError("num_variants must be at least one") + max_samples = min(max_samples, num_variants) + if max_samples < 1: + raise ValueError("max_samples must be at least one") + records = load_manifest(source_manifest)[:max_samples] + destination = Path(output_dir).expanduser().resolve() + image_dir = destination / "images" + trajectory_dir = destination / "trajectories" + image_dir.mkdir(parents=True, exist_ok=True) + trajectory_dir.mkdir(parents=True, exist_ok=True) + + rows: list[dict[str, Any]] = [] + for index, record in enumerate(records): + stem = f"sample_{index:02d}_{record.sample_id}" + image_path = image_dir / f"{stem}{record.image_path.suffix.lower()}" + trajectory_path = trajectory_dir / f"{stem}.npz" + shutil.copy2(record.image_path, image_path) + shutil.copy2(record.trajectory_path, trajectory_path) + rows.append( + { + "clip_id": record.clip_id, + "t0_us": record.t0_us, + "reasoning": record.reasoning, + "image_path": str(image_path.relative_to(destination)), + "trajectory_path": str(trajectory_path.relative_to(destination)), + "history_steps": 8, + "expert_waypoints": 16, + } + ) + + manifest_path = destination / "overfit_manifest.jsonl" + manifest_path.write_text( + "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows), + encoding="utf-8", + ) + return manifest_path + + +def _resample_waypoints(xy: torch.Tensor, steps: int = 16) -> torch.Tensor: + if xy.shape[0] == steps: + return xy.float() + if xy.shape[0] < 2: + raise ValueError("future trajectory must contain at least two waypoints") + signal = xy.float().T.unsqueeze(0) + return F.interpolate(signal, size=steps, mode="linear", align_corners=True)[0].T + + +class Pi05OverfitDataset(Dataset[dict[str, Any]]): + """Tiny deterministic PAI dataset for validating a pi0.5-style training loop.""" + + def __init__( + self, + manifest_path: str | Path, + *, + image_size: int | tuple[int, int] = (128, 224), + augmentation_factor: int = 4, + action_horizon: int = 16, + seed: int = 7, + tokenizer: MiniTokenizer | None = None, + action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None, + ) -> None: + if augmentation_factor < 1: + raise ValueError("augmentation_factor must be at least one") + if isinstance(image_size, int): + image_size = (image_size, image_size) + if len(image_size) != 2 or min(image_size) < 1: + raise ValueError("image_size must be a positive (height, width) pair") + if action_horizon < 2: + raise ValueError("action_horizon must be at least two") + self.records = load_manifest(manifest_path) + self.image_size = (int(image_size[0]), int(image_size[1])) + self.augmentation_factor = int(augmentation_factor) + self.action_horizon = int(action_horizon) + self.seed = int(seed) + self.tokenizer = tokenizer or MiniTokenizer() + self.state_dim = 16 + + if action_statistics is None: + geometries = [ + self._load_augmented_geometry(base_index, variant) + for base_index in range(len(self.records)) + for variant in range(self.augmentation_factor) + ] + histories = [geometry[0] for geometry in geometries] + futures = [geometry[1] for geometry in geometries] + self.action_normalizer = ActionNormalizer.fit(histories, futures) + elif isinstance(action_statistics, ActionStatistics): + self.action_normalizer = ActionNormalizer( + action_statistics.mean, action_statistics.scale + ) + else: + self.action_normalizer = ActionNormalizer.from_dict(action_statistics) + + @property + def action_statistics(self) -> ActionStatistics: + return self.action_normalizer.statistics + + @property + def action_mean(self) -> torch.Tensor: + return self.action_normalizer.mean.clone() + + @property + def action_std(self) -> torch.Tensor: + return self.action_normalizer.scale.clone() + + def save_action_statistics(self, path: str | Path) -> Path: + return self.action_normalizer.save(path) + + def __len__(self) -> int: + return len(self.records) * self.augmentation_factor + + def _indices(self, index: int) -> tuple[int, int]: + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + return divmod(index, self.augmentation_factor) + + def _rng(self, base_index: int, variant: int) -> np.random.Generator: + return np.random.default_rng( + np.random.SeedSequence([self.seed, base_index, variant]) + ) + + def _load_augmented_geometry( + self, base_index: int, variant: int + ) -> tuple[torch.Tensor, torch.Tensor]: + record = self.records[base_index] + with np.load(record.trajectory_path) as arrays: + history = torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2] + future = torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2] + future = _resample_waypoints(future, steps=self.action_horizon) + if variant == 0: + return history, future + + rng = self._rng(base_index, variant) + angle = float(rng.normal(0.0, 0.008)) + scale = float(rng.normal(1.0, 0.006)) + rotation = history.new_tensor( + [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]] + ) + history = scale * (history @ rotation.T) + future = scale * (future @ rotation.T) + # A tiny smooth lateral perturbation adds distinct trajectories while retaining the label. + curvature = float(rng.normal(0.0, 0.012)) + progress = torch.linspace(0.0, 1.0, future.shape[0], dtype=future.dtype) + future[:, 1] += curvature * progress.square() + return history, future + + def _load_image(self, base_index: int, variant: int) -> torch.Tensor: + with Image.open(self.records[base_index].image_path) as source: + image = source.convert("RGB") + if variant: + rng = self._rng(base_index, variant) + brightness = float(rng.uniform(0.97, 1.03)) + image = ImageEnhance.Brightness(image).enhance(brightness) + height, width = self.image_size + image = image.resize((width, height), resample=Image.Resampling.BILINEAR) + array = np.asarray(image, dtype=np.float32).copy() / 255.0 + return torch.from_numpy(array).permute(2, 0, 1).contiguous() + + def __getitem__(self, index: int) -> dict[str, Any]: + base_index, variant = self._indices(index) + record = self.records[base_index] + history_xy, future_xy = self._load_augmented_geometry(base_index, variant) + reasoning = record.reasoning + target = self.tokenizer.encode_subtask(reasoning) + actions = self.action_normalizer.encode(history_xy, future_xy) + return { + "image": self._load_image(base_index, variant), + "prefix_tokens": self.tokenizer.build_prefix(history_xy), + "subtask_target": target, + "subtask_label": self.tokenizer.subtasks.decode(target), + "actions": actions, + "action_tokens": self.tokenizer.encode_actions(actions), + "state": history_xy.reshape(-1), + "history_xy": history_xy, + "future_xy": future_xy, + "reasoning": reasoning, + "sample_id": f"{record.sample_id}:aug{variant}", + } diff --git a/src/pi05mini/flow.py b/src/pi05mini/flow.py new file mode 100644 index 0000000..ad24cb9 --- /dev/null +++ b/src/pi05mini/flow.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from collections.abc import Callable + +import torch +from torch import Tensor + + +class FlowMatching: + """Rectified-flow objective used by the minimal action expert.""" + + def __init__(self, inference_steps: int = 10, cutoff: float = 0.999) -> None: + self.inference_steps = inference_steps + self.cutoff = cutoff + + def training_batch( + self, target: Tensor, generator: torch.Generator | None = None + ) -> tuple[Tensor, Tensor, Tensor]: + batch = target.shape[0] + uniform = torch.rand( + batch, device=target.device, dtype=target.dtype, generator=generator + ) + # Paper sampling: z~Beta(1.5,1), tau=s*(1-z), emphasizing noisy/low tau. + timestep = self.cutoff * (1.0 - uniform.pow(1.0 / 1.5)) + noise = torch.randn( + target.shape, + device=target.device, + dtype=target.dtype, + generator=generator, + ) + t = timestep[:, None, None] + noisy = t * target + (1.0 - t) * noise + return noisy, timestep, target - noise + + @staticmethod + def loss(prediction: Tensor, target_velocity: Tensor) -> Tensor: + return torch.nn.functional.mse_loss(prediction.float(), target_velocity.float()) + + @torch.no_grad() + def sample( + self, + model_fn: Callable[[Tensor, Tensor], Tensor], + shape: tuple[int, ...], + device: torch.device | str, + steps: int | None = None, + generator: torch.Generator | None = None, + ) -> Tensor: + count = steps or self.inference_steps + x = torch.randn(shape, device=device, generator=generator) + grid = torch.linspace(0.0, 1.0, count + 1, device=device) + for index in range(count): + timestep = grid[index].expand(shape[0]) + x = x + (grid[index + 1] - grid[index]) * model_fn(x, timestep) + return x diff --git a/src/pi05mini/infer.py b/src/pi05mini/infer.py new file mode 100644 index 0000000..6dd9b2f --- /dev/null +++ b/src/pi05mini/infer.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import argparse +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import torch + +from .data import Pi05OverfitDataset +from .flow import FlowMatching +from .model import Pi05Config, Pi05Mini +from .runtime import choose_device, write_json + + +@torch.no_grad() +def evaluate(checkpoint_path: str, manifest: str | None = None) -> Path: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + cfg = checkpoint["config"] + root = Path(__file__).resolve().parents[2] + device = choose_device(cfg.get("device", "auto")) + data_manifest = Path(manifest) if manifest else root / "data/overfit/overfit_manifest.jsonl" + dataset = Pi05OverfitDataset( + data_manifest, image_size=cfg["data"]["image_size"], + action_horizon=cfg["data"]["action_horizon"], + action_statistics=checkpoint["action_statistics"] + ) + model = Pi05Mini(Pi05Config(**checkpoint["model_config"])) + model.load_state_dict(checkpoint["model"]) + model.to(device).eval() + flow = FlowMatching(cfg["inference"]["flow_steps"]) + rows = [] + figure, axes = plt.subplots(1, len(dataset), figsize=(4 * len(dataset), 4), squeeze=False) + for index, sample in enumerate(dataset): + image = sample["image"][None].to(device) + tokens = sample["prefix_tokens"][None].to(device) + # Autoregressively predict semantic subtask followed by the complete + # discretized action stream (ACTION_START, 32 scalar tokens, ACTION_END). + zero_action = torch.zeros(1, dataset.action_horizon, 2, device=device) + generated: list[int] = [] + expected_tokens = 1 + 2 + dataset.action_horizon * 2 + current_tokens = tokens + for _ in range(expected_tokens): + semantic = model( + image, current_tokens, zero_action, + torch.zeros(1, device=device), + generated_token_start=tokens.shape[1] + ) + next_token = int(semantic.text_logits[0, -1].argmax().cpu()) + generated.append(next_token) + current_tokens = torch.cat( + (current_tokens, torch.tensor([[next_token]], device=device)), dim=1 + ) + predicted_token = generated[0] + vlm_action_tokens = torch.tensor(generated[1:], dtype=torch.long) + try: + vlm_normalized = dataset.tokenizer.decode_actions(vlm_action_tokens) + vlm_xy = dataset.action_normalizer.decode( + vlm_normalized, sample["history_xy"] + ).numpy() + vlm_valid = True + except ValueError: + vlm_xy = np.full_like(sample["future_xy"].numpy(), np.nan) + vlm_valid = False + conditioned_tokens = torch.cat( + (tokens, torch.tensor([[predicted_token]], device=device)), dim=1 + ) + def velocity(x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + return model( + image, conditioned_tokens, x, t, + generated_token_start=tokens.shape[1] + ).vector_field + normalized = flow.sample( + velocity, (1, dataset.action_horizon, 2), device, + steps=cfg["inference"]["flow_steps"], + generator=None, + )[0].cpu() + predicted_xy = dataset.action_normalizer.decode( + normalized, sample["history_xy"] + ).numpy() + target_xy = sample["future_xy"].numpy() + diffusion_mse = float(np.mean((predicted_xy - target_xy) ** 2)) + vlm_mse = float(np.mean((vlm_xy - target_xy) ** 2)) if vlm_valid else None + target_label = dataset.tokenizer.subtasks.decode(sample["subtask_target"]) + try: + predicted_label = dataset.tokenizer.subtasks.decode(predicted_token) + except ValueError: + predicted_label = f"token_{predicted_token}" + rows.append({ + "sample_id": sample["sample_id"], "target_subtask": target_label, + "predicted_subtask": predicted_label, + "vlm_action_tokens_valid": vlm_valid, + "vlm_trajectory_mse": vlm_mse, + "diffusion_trajectory_mse": diffusion_mse, + }) + ax = axes[0, index] + ax.plot(target_xy[:, 0], target_xy[:, 1], "o-", label="target") + ax.plot(vlm_xy[:, 0], vlm_xy[:, 1], "s--", label="VLM tokens") + ax.plot(predicted_xy[:, 0], predicted_xy[:, 1], "x--", label="diffusion") + ax.scatter([0], [0], c="black", marker="s", label="ego") + vlm_text = "invalid" if vlm_mse is None else f"{vlm_mse:.3f}" + ax.set_title( + f"{predicted_label}\nVLM={vlm_text}, diffusion={diffusion_mse:.3f}" + ) + ax.set_xlabel("forward x (m)") + ax.set_ylabel("left y (m)") + ax.axis("equal") + ax.grid(True, alpha=.3) + ax.legend(fontsize=8) + figure.tight_layout() + output_dir = Path(checkpoint_path).parent + plot_path = output_dir / "trajectory_comparison.png" + figure.savefig(plot_path, dpi=160) + plt.close(figure) + write_json(output_dir / "evaluation.json", rows) + print(rows) + return plot_path + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt") + parser.add_argument("--manifest") + args = parser.parse_args() + evaluate(args.checkpoint, args.manifest) + + +if __name__ == "__main__": + main() diff --git a/src/pi05mini/model.py b/src/pi05mini/model.py new file mode 100644 index 0000000..0496397 --- /dev/null +++ b/src/pi05mini/model.py @@ -0,0 +1,396 @@ +"""A compact, pure-PyTorch approximation of the pi0.5 dual-expert model. + +The important architectural detail is that prefix and action tokens have +different parameters at every Transformer layer. Prefix queries only attend +to prefix keys, while action queries attend to the complete prefix and to all +action tokens. This is deliberately implemented explicitly (instead of +concatenating the streams into an ordinary Transformer) so the information +flow is easy to inspect in a minimal reproduction. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional + +import torch +from torch import Tensor, nn +import torch.nn.functional as F + + +@dataclass +class Pi05Config: + vocab_size: int = 512 + action_dim: int = 2 + image_channels: int = 3 + image_size: int = 64 + patch_size: int = 16 + d_model: int = 128 + n_heads: int = 4 + n_layers: int = 4 + mlp_ratio: int = 4 + max_text_tokens: int = 64 + max_action_tokens: int = 32 + dropout: float = 0.0 + rms_eps: float = 1e-6 + + def __post_init__(self) -> None: + if self.d_model % self.n_heads: + raise ValueError("d_model must be divisible by n_heads") + if self.image_size % self.patch_size: + raise ValueError("image_size must be divisible by patch_size") + + @property + def max_image_tokens(self) -> int: + return (self.image_size // self.patch_size) ** 2 + + @property + def max_prefix_tokens(self) -> int: + return self.max_image_tokens + self.max_text_tokens + + +@dataclass +class Pi05Output: + text_logits: Tensor + vector_field: Tensor + prefix_hidden: Tensor + action_hidden: Tensor + text_token_offset: int + loss: Optional[Tensor] = None + text_loss: Optional[Tensor] = None + flow_loss: Optional[Tensor] = None + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: Tensor) -> Tensor: + # Accumulating the variance in fp32 avoids overflows under autocast. + rms = x.float().pow(2).mean(dim=-1, keepdim=True) + normalized = x * torch.rsqrt(rms.to(x.dtype) + self.eps) + return normalized * self.weight + + +class AdaRMSNorm(nn.Module): + """RMSNorm modulated by a flow-time embedding.""" + + def __init__(self, dim: int, cond_dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.norm = RMSNorm(dim, eps) + self.modulation = nn.Linear(cond_dim, 2 * dim) + # Start as an ordinary RMSNorm, which makes tiny-data training stable. + nn.init.zeros_(self.modulation.weight) + nn.init.zeros_(self.modulation.bias) + + def forward(self, x: Tensor, condition: Tensor) -> Tensor: + shift, scale = self.modulation(condition).chunk(2, dim=-1) + return self.norm(x) * (1 + scale[:, None, :]) + shift[:, None, :] + + +class TimeEmbedding(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + fourier_dim = max(8, dim // 2) + if fourier_dim % 2: + fourier_dim += 1 + self.fourier_dim = fourier_dim + self.mlp = nn.Sequential( + nn.Linear(fourier_dim, dim), nn.SiLU(), nn.Linear(dim, dim) + ) + + def forward(self, timestep: Tensor) -> Tensor: + timestep = timestep.reshape(-1).float() + half = self.fourier_dim // 2 + frequencies = torch.exp( + torch.linspace(math.log(1.0), math.log(1000.0), half, + device=timestep.device, dtype=timestep.dtype) + ) + angles = 2 * math.pi * timestep[:, None] * frequencies[None, :] + embedding = torch.cat((angles.sin(), angles.cos()), dim=-1) + return self.mlp(embedding) + + +class AttentionWeights(nn.Module): + """One expert's independent Q/K/V/output projections.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.q = nn.Linear(dim, dim, bias=False) + self.k = nn.Linear(dim, dim, bias=False) + self.v = nn.Linear(dim, dim, bias=False) + self.out = nn.Linear(dim, dim, bias=False) + + +class MLP(nn.Module): + def __init__(self, dim: int, ratio: int, dropout: float) -> None: + super().__init__() + hidden = dim * ratio + self.net = nn.Sequential( + nn.Linear(dim, hidden), + nn.GELU(approximate="tanh"), + nn.Dropout(dropout), + nn.Linear(hidden, dim), + nn.Dropout(dropout), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.net(x) + + +def _attention( + query: Tensor, + key: Tensor, + value: Tensor, + n_heads: int, + allowed: Optional[Tensor], + dropout: float, + training: bool, +) -> Tensor: + batch, q_len, width = query.shape + k_len = key.shape[1] + head_dim = width // n_heads + + def split(x: Tensor) -> Tensor: + return x.reshape(batch, -1, n_heads, head_dim).transpose(1, 2) + + q, k, v = split(query), split(key), split(value) + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + if allowed is not None: + # allowed is [B or 1, Q, K]; True means the edge exists. + scores = scores.masked_fill(~allowed[:, None, :, :], torch.finfo(scores.dtype).min) + probabilities = scores.float().softmax(dim=-1).to(scores.dtype) + probabilities = F.dropout(probabilities, p=dropout, training=training) + output = torch.matmul(probabilities, v) + return output.transpose(1, 2).contiguous().reshape(batch, q_len, width) + + +class DualExpertBlock(nn.Module): + def __init__(self, config: Pi05Config) -> None: + super().__init__() + d = config.d_model + self.n_heads = config.n_heads + self.dropout = config.dropout + + # The two streams never share Transformer parameters. + self.vlm_norm1 = RMSNorm(d, config.rms_eps) + self.vlm_attention = AttentionWeights(d) + self.vlm_norm2 = RMSNorm(d, config.rms_eps) + self.vlm_mlp = MLP(d, config.mlp_ratio, config.dropout) + + self.action_norm1 = AdaRMSNorm(d, d, config.rms_eps) + self.action_attention = AttentionWeights(d) + self.action_norm2 = AdaRMSNorm(d, d, config.rms_eps) + self.action_mlp = MLP(d, config.mlp_ratio, config.dropout) + + def forward( + self, + prefix: Tensor, + action: Tensor, + time_condition: Tensor, + prefix_edges: Tensor, + action_edges: Tensor, + ) -> tuple[Tensor, Tensor]: + prefix_input = self.vlm_norm1(prefix) + p = self.vlm_attention + prefix_update = _attention( + p.q(prefix_input), p.k(prefix_input), p.v(prefix_input), + self.n_heads, prefix_edges, self.dropout, self.training + ) + prefix = prefix + p.out(prefix_update) + prefix = prefix + self.vlm_mlp(self.vlm_norm2(prefix)) + + action_input = self.action_norm1(action, time_condition) + a = self.action_attention + # Prefix KV uses the VLM expert's parameters; action KV uses the action + # expert's parameters. Thus action can read prefix without allowing a + # reverse information path from action into the prefix stream. + key = torch.cat((p.k(prefix_input), a.k(action_input)), dim=1) + value = torch.cat((p.v(prefix_input), a.v(action_input)), dim=1) + action_update = _attention( + a.q(action_input), key, value, self.n_heads, action_edges, + self.dropout, self.training + ) + action = action + a.out(action_update) + action = action + self.action_mlp(self.action_norm2(action, time_condition)) + return prefix, action + + +class Pi05Mini(nn.Module): + """Minimal pi0.5-style model for architecture validation/overfitting.""" + + def __init__(self, config: Optional[Pi05Config] = None) -> None: + super().__init__() + self.config = config or Pi05Config() + c = self.config + self.patch_encoder = nn.Conv2d( + c.image_channels, c.d_model, c.patch_size, stride=c.patch_size + ) + self.token_embedding = nn.Embedding(c.vocab_size, c.d_model) + self.action_projection = nn.Linear(c.action_dim, c.d_model) + self.prefix_position = nn.Parameter( + torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01 + ) + self.action_position = nn.Parameter( + torch.randn(1, c.max_action_tokens, c.d_model) * 0.01 + ) + self.time_embedding = TimeEmbedding(c.d_model) + self.layers = nn.ModuleList(DualExpertBlock(c) for _ in range(c.n_layers)) + self.prefix_norm = RMSNorm(c.d_model, c.rms_eps) + self.action_norm = AdaRMSNorm(c.d_model, c.d_model, c.rms_eps) + self.text_head = nn.Linear(c.d_model, c.vocab_size, bias=False) + self.vector_field_head = nn.Linear(c.d_model, c.action_dim) + self.text_head.weight = self.token_embedding.weight + + def encode_prefix(self, pixel_values: Tensor, input_ids: Tensor) -> tuple[Tensor, int]: + patches = self.patch_encoder(pixel_values).flatten(2).transpose(1, 2) + text_tokens = self.token_embedding(input_ids) + # pi0.5 represents state in the VLM vocabulary. There is deliberately + # no separate continuous-state projection in this model. + prefix = torch.cat((patches, text_tokens), dim=1) + if prefix.shape[1] > self.config.max_prefix_tokens: + raise ValueError("prefix exceeds configured maximum length") + prefix = prefix + self.prefix_position[:, : prefix.shape[1]] + return prefix, patches.shape[1] + + def _masks( + self, + batch: int, + prefix_length: int, + action_length: int, + text_offset: int, + attention_mask: Optional[Tensor], + action_prefix_mask: Optional[Tensor], + generated_token_start: Optional[int], + device: torch.device, + ) -> tuple[Tensor, Tensor]: + valid = torch.ones(batch, prefix_length, dtype=torch.bool, device=device) + if attention_mask is not None: + valid[:, text_offset:] = attention_mask.bool() + # Image plus observation/instruction/state tokens form a fully + # bidirectional observation block. Tokens generated after that block + # (subtask and FAST actions) read the observation and earlier generated + # tokens causally; observation queries cannot read generated tokens. + generated_start = prefix_length if generated_token_start is None else text_offset + generated_token_start + if not text_offset <= generated_start <= prefix_length: + raise ValueError("generated_token_start must index input_ids") + q = torch.arange(prefix_length, device=device)[:, None] + k = torch.arange(prefix_length, device=device)[None, :] + structural = torch.where(q < generated_start, k < generated_start, k <= q) + prefix_edges = structural[None] & valid[:, None, :] + # Padded prefix queries retain a harmless self edge to avoid all -inf + # rows; their outputs are not used by the text loss. + query_valid = valid[:, :, None] + identity = torch.eye(prefix_length, dtype=torch.bool, device=device)[None] + prefix_edges = torch.where(query_valid, prefix_edges, identity) + + action_valid = valid + if action_prefix_mask is not None: + supplied = action_prefix_mask.bool() + if supplied.shape[1] == prefix_length - text_offset: + supplied = torch.cat( + (torch.ones(batch, text_offset, dtype=torch.bool, device=device), supplied), + dim=1, + ) + if supplied.shape != valid.shape: + raise ValueError("action_prefix_mask must cover text tokens or the full prefix") + action_valid = action_valid & supplied + prefix_keys = action_valid[:, None, :].expand(batch, action_length, prefix_length) + action_keys = torch.ones( + batch, action_length, action_length, dtype=torch.bool, device=device + ) + action_edges = torch.cat((prefix_keys, action_keys), dim=-1) + return prefix_edges, action_edges + + def forward( + self, + pixel_values: Tensor, + input_ids: Tensor, + noisy_actions: Tensor, + timesteps: Tensor, + attention_mask: Optional[Tensor] = None, + action_prefix_mask: Optional[Tensor] = None, + generated_token_start: Optional[int] = None, + text_labels: Optional[Tensor] = None, + target_velocity: Optional[Tensor] = None, + text_loss_weight: float = 1.0, + flow_loss_weight: float = 1.0, + ) -> Pi05Output: + """Run both experts. + + ``timesteps`` is a scalar per batch item in [0, 1]. ``text_labels`` + corresponds only to ``input_ids`` (not image prefix positions) + and may contain -100. ``target_velocity`` has the same shape as + ``noisy_actions``. + """ + if noisy_actions.ndim != 3: + raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]") + if noisy_actions.shape[1] > self.config.max_action_tokens: + raise ValueError("action horizon exceeds configured maximum") + + prefix, text_offset = self.encode_prefix(pixel_values, input_ids) + horizon = noisy_actions.shape[1] + action = self.action_projection(noisy_actions) + action = action + self.action_position[:, :horizon] + condition = self.time_embedding(timesteps).to(action.dtype) + prefix_edges, action_edges = self._masks( + prefix.shape[0], prefix.shape[1], horizon, text_offset, + attention_mask, action_prefix_mask, generated_token_start, prefix.device + ) + for layer in self.layers: + prefix, action = layer( + prefix, action, condition, prefix_edges, action_edges + ) + + prefix_hidden = self.prefix_norm(prefix) + action_hidden = self.action_norm(action, condition) + # Return logits only for actual text positions, which keeps downstream + # next-token loss/generation code uncomplicated. + text_logits = self.text_head(prefix_hidden[:, text_offset:]) + vector_field = self.vector_field_head(action_hidden) + + text_loss = None + if text_labels is not None: + text_loss = F.cross_entropy( + text_logits.reshape(-1, text_logits.shape[-1]), + text_labels.reshape(-1), ignore_index=-100 + ) + flow_loss = None + if target_velocity is not None: + flow_loss = F.mse_loss(vector_field, target_velocity) + loss = None + if text_loss is not None or flow_loss is not None: + loss = sum( + component for component in ( + None if text_loss is None else text_loss_weight * text_loss, + None if flow_loss is None else flow_loss_weight * flow_loss, + ) if component is not None + ) + return Pi05Output( + text_logits=text_logits, + vector_field=vector_field, + prefix_hidden=prefix_hidden, + action_hidden=action_hidden, + text_token_offset=text_offset, + loss=loss, + text_loss=text_loss, + flow_loss=flow_loss, + ) + + +def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]: + """Construct pi-style conditional-flow training targets. + + The path is x_t = (1-t) * noise + t * action, with constant target vector + field ``action - noise``. + """ + batch = actions.shape[0] + timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype) + noise = torch.randn_like(actions) + broadcast_t = timestep.reshape(batch, *([1] * (actions.ndim - 1))) + noisy_actions = (1 - broadcast_t) * noise + broadcast_t * actions + target_velocity = actions - noise + return noisy_actions, timestep, target_velocity diff --git a/src/pi05mini/runtime.py b/src/pi05mini/runtime.py new file mode 100644 index 0000000..d015228 --- /dev/null +++ b/src/pi05mini/runtime.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import json +import random +from pathlib import Path + +import numpy as np +import torch +import yaml + + +def load_config(path: str | Path) -> dict: + with Path(path).open("r", encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def seed_everything(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def choose_device(requested: str = "auto") -> torch.device: + if requested != "auto": + return torch.device(requested) + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + +def write_json(path: str | Path, value: object) -> None: + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(json.dumps(value, indent=2, ensure_ascii=False), encoding="utf-8") diff --git a/src/pi05mini/tokenizer.py b/src/pi05mini/tokenizer.py new file mode 100644 index 0000000..461b644 --- /dev/null +++ b/src/pi05mini/tokenizer.py @@ -0,0 +1,285 @@ +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Iterable + +import torch + + +@dataclass(frozen=True) +class CoordinateRange: + minimum: float + maximum: float + + def __post_init__(self) -> None: + if self.maximum <= self.minimum: + raise ValueError("coordinate maximum must be greater than minimum") + + +class CoordinateTokenizer: + """Uniformly discretize XY state values into the literal token range 0..255.""" + + bins = 256 + + def __init__( + self, + x_range: tuple[float, float] = (-10.0, 10.0), + y_range: tuple[float, float] = (-5.0, 5.0), + ) -> None: + self.ranges = ( + CoordinateRange(*x_range), + CoordinateRange(*y_range), + ) + + def encode(self, coordinates: torch.Tensor | Iterable[Iterable[float]]) -> torch.Tensor: + values = torch.as_tensor(coordinates, dtype=torch.float32) + if values.shape[-1] != 2: + raise ValueError("coordinates must end in an XY dimension of size 2") + low = values.new_tensor([axis.minimum for axis in self.ranges]) + high = values.new_tensor([axis.maximum for axis in self.ranges]) + unit = ((values - low) / (high - low)).clamp(0.0, 1.0) + return torch.round(unit * (self.bins - 1)).to(torch.long) + + def decode(self, tokens: torch.Tensor | Iterable[Iterable[int]]) -> torch.Tensor: + encoded = torch.as_tensor(tokens, dtype=torch.long) + if encoded.shape[-1] != 2: + raise ValueError("coordinate tokens must end in an XY dimension of size 2") + if encoded.numel() and (encoded.min() < 0 or encoded.max() >= self.bins): + raise ValueError("coordinate tokens must be in the inclusive range 0..255") + low = torch.tensor([axis.minimum for axis in self.ranges], dtype=torch.float32) + high = torch.tensor([axis.maximum for axis in self.ranges], dtype=torch.float32) + return low + encoded.float() / (self.bins - 1) * (high - low) + + +class SubtaskVocabulary: + """Small semantic vocabulary standing in for pi0.5's generated language subtask.""" + + labels = ( + "keep_lane", + "steer_left", + "steer_right", + "slow_down", + "stop", + "accelerate", + "unknown", + ) + + def __init__(self, first_token_id: int = 260) -> None: + self.first_token_id = int(first_token_id) + self._label_to_token = { + label: self.first_token_id + index for index, label in enumerate(self.labels) + } + self._token_to_label = {token: label for label, token in self._label_to_token.items()} + + @property + def size(self) -> int: + return len(self.labels) + + @property + def token_ids(self) -> tuple[int, ...]: + return tuple(self._label_to_token[label] for label in self.labels) + + def classify(self, reasoning: str) -> str: + text = reasoning.casefold() + # More specific commands precede generic lane/follow language. + if any(term in text for term in ("stop", "halt", "stationary")): + return "stop" + if any(term in text for term in ("slow", "deceler", "reduce speed", "brake")): + return "slow_down" + if any(term in text for term in ("accelerat", "speed up")): + return "accelerate" + if any(term in text for term in ("left", "counterclockwise")): + return "steer_left" + if any(term in text for term in ("right", "clockwise")): + return "steer_right" + if any(term in text for term in ("lane", "straight", "follow", "continue", "proceed")): + return "keep_lane" + return "unknown" + + def encode(self, label_or_reasoning: str) -> int: + label = ( + label_or_reasoning + if label_or_reasoning in self._label_to_token + else self.classify(label_or_reasoning) + ) + return self._label_to_token[label] + + def decode(self, token_id: int | torch.Tensor) -> str: + token = int(token_id) + try: + return self._token_to_label[token] + except KeyError as error: + raise ValueError(f"unknown subtask token: {token}") from error + + +class MiniTokenizer: + """Compose image/state, semantic, and FAST-lite action tokens. + + Actions are z-score-normalized by the dataset before reaching this + tokenizer. A complete ``[horizon, 2]`` chunk is transformed along time + with an orthonormal DCT-II. All coefficients are retained and serialized + in frequency-major XY order, so low-frequency motion appears first in the + token stream. The coefficients are divided by ``sqrt(horizon)`` before + being uniformly quantized over ``[-4, 4]`` into 64 shared tokens. This + keeps the fixed quantization interval valid for any horizon when the input + actions have first been clipped to the same interval. + + This intentionally implements only the central transform/order idea from + FAST. It does not learn dataset-specific coefficient bins or drop trailing + high-frequency coefficients. + """ + + BOS = 256 + IMAGE = 257 + STATE = 258 + SUBTASK_QUERY = 259 + ACTION_START = SUBTASK_QUERY + 1 + len(SubtaskVocabulary.labels) + ACTION_END = ACTION_START + 1 + ACTION_TOKEN_START = ACTION_END + 1 + ACTION_BINS = 64 + ACTION_TOKEN_END = ACTION_TOKEN_START + ACTION_BINS - 1 + ACTION_MINIMUM = -4.0 + ACTION_MAXIMUM = 4.0 + + def __init__( + self, + coordinate_tokenizer: CoordinateTokenizer | None = None, + subtasks: SubtaskVocabulary | None = None, + ) -> None: + self.coordinates = coordinate_tokenizer or CoordinateTokenizer() + self.subtasks = subtasks or SubtaskVocabulary(first_token_id=self.SUBTASK_QUERY + 1) + + @property + def vocab_size(self) -> int: + return max(max(self.subtasks.token_ids), self.ACTION_TOKEN_END) + 1 + + @property + def action_token_ids(self) -> range: + """Contiguous vocabulary interval reserved for quantized DCT coefficients.""" + + return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1) + + def build_prefix(self, history_xy: torch.Tensor) -> torch.Tensor: + coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1) + specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long) + query = torch.tensor([self.SUBTASK_QUERY], dtype=torch.long) + return torch.cat((specials, coordinate_tokens.cpu(), query), dim=0) + + def encode_subtask(self, reasoning: str) -> torch.Tensor: + return torch.tensor(self.subtasks.encode(reasoning), dtype=torch.long) + + @staticmethod + def _orthonormal_dct_matrix( + horizon: int, + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Return the orthonormal DCT-II matrix ``[frequency, time]``.""" + + if horizon <= 0: + raise ValueError("DCT horizon must be positive") + time = torch.arange(horizon, device=device, dtype=dtype) + frequency = torch.arange(horizon, device=device, dtype=dtype).unsqueeze(1) + matrix = torch.cos(math.pi / horizon * (time + 0.5) * frequency) + matrix[0] *= math.sqrt(1.0 / horizon) + if horizon > 1: + matrix[1:] *= math.sqrt(2.0 / horizon) + return matrix + + def dct_actions( + self, actions: torch.Tensor | Iterable[Iterable[float]] + ) -> torch.Tensor: + """Apply an orthonormal DCT-II along an action chunk's time axis.""" + + values = torch.as_tensor(actions, dtype=torch.float32) + if values.ndim != 2 or values.shape[-1] != 2: + raise ValueError("actions must have shape [horizon, 2]") + if values.shape[0] == 0: + raise ValueError("actions must contain at least one timestep") + transform = self._orthonormal_dct_matrix( + values.shape[0], device=values.device, dtype=values.dtype + ) + return transform @ values + + def idct_actions( + self, coefficients: torch.Tensor | Iterable[Iterable[float]] + ) -> torch.Tensor: + """Invert :meth:`dct_actions` using the transpose of its basis.""" + + values = torch.as_tensor(coefficients, dtype=torch.float32) + if values.ndim != 2 or values.shape[-1] != 2: + raise ValueError("DCT coefficients must have shape [horizon, 2]") + if values.shape[0] == 0: + raise ValueError("DCT coefficients must contain at least one frequency") + transform = self._orthonormal_dct_matrix( + values.shape[0], device=values.device, dtype=values.dtype + ) + return transform.transpose(0, 1) @ values + + def encode_actions( + self, actions: torch.Tensor | Iterable[Iterable[float]] + ) -> torch.Tensor: + """DCT-transform and encode normalized actions as a marked token stream.""" + + values = torch.as_tensor(actions, dtype=torch.float32) + if values.ndim != 2 or values.shape[-1] != 2: + raise ValueError("actions must have shape [horizon, 2]") + if values.shape[0] == 0: + raise ValueError("actions must contain at least one timestep") + values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM) + coefficients = self.dct_actions(values) / math.sqrt(values.shape[0]) + # Exact zero is useful for a stable shared center bin. Tiny non-zero + # values here are only floating-point residue from the DCT basis. + coefficients = torch.where( + coefficients.abs() < 1e-6, + torch.zeros_like(coefficients), + coefficients, + ) + unit = ( + ( + coefficients.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM) + - self.ACTION_MINIMUM + ) + / (self.ACTION_MAXIMUM - self.ACTION_MINIMUM) + ) + quantized = torch.round(unit * (self.ACTION_BINS - 1)).to(torch.long) + # Row-major flattening is frequency-major: DC XY, frequency-1 XY, ... + action_tokens = quantized.reshape(-1) + self.ACTION_TOKEN_START + boundaries = torch.tensor( + [self.ACTION_START, self.ACTION_END], + dtype=torch.long, + device=action_tokens.device, + ) + return torch.cat((boundaries[:1], action_tokens, boundaries[1:])) + + def decode_actions(self, tokens: torch.Tensor | Iterable[int]) -> torch.Tensor: + """Decode FAST-lite tokens back to normalized ``[horizon, 2]`` actions.""" + + encoded = torch.as_tensor(tokens, dtype=torch.long) + if encoded.ndim != 1: + raise ValueError("action tokens must be a one-dimensional sequence") + if ( + encoded.numel() < 4 + or encoded[0] != self.ACTION_START + or encoded[-1] != self.ACTION_END + ): + raise ValueError("action tokens must be enclosed by ACTION_START and ACTION_END") + payload = encoded[1:-1] + if payload.numel() % 2: + raise ValueError("action token payload must contain complete XY pairs") + if payload.min() < self.ACTION_TOKEN_START or payload.max() > self.ACTION_TOKEN_END: + raise ValueError( + "action token payload is outside the reserved action token interval" + ) + quantized = payload - self.ACTION_TOKEN_START + normalized_coefficients = ( + quantized.float() + / (self.ACTION_BINS - 1) + * (self.ACTION_MAXIMUM - self.ACTION_MINIMUM) + + self.ACTION_MINIMUM + ) + horizon = payload.numel() // 2 + coefficients = normalized_coefficients.reshape(horizon, 2) * math.sqrt(horizon) + return self.idct_actions(coefficients) diff --git a/src/pi05mini/train.py b/src/pi05mini/train.py new file mode 100644 index 0000000..2524517 --- /dev/null +++ b/src/pi05mini/train.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from torch.utils.data import DataLoader + +from .data import Pi05OverfitDataset, materialize_subset +from .flow import FlowMatching +from .model import Pi05Config, Pi05Mini +from .runtime import choose_device, load_config, seed_everything, write_json + + +def _collate(rows: list[dict]) -> dict: + tensor_keys = ( + "image", "prefix_tokens", "subtask_target", "actions", "action_tokens" + ) + result = {key: torch.stack([row[key] for row in rows]) for key in tensor_keys} + result["sample_id"] = [row["sample_id"] for row in rows] + return result + + +def build_model(config: dict, dataset: Pi05OverfitDataset) -> Pi05Mini: + model_cfg = config["model"] + return Pi05Mini(Pi05Config( + vocab_size=dataset.tokenizer.vocab_size, + action_dim=2, + image_size=config["data"]["image_size"], + patch_size=model_cfg["patch_size"], + d_model=model_cfg["d_model"], + n_heads=model_cfg["n_heads"], + n_layers=model_cfg["n_layers"], + mlp_ratio=model_cfg["mlp_ratio"], + max_text_tokens=model_cfg["max_text_tokens"], + max_action_tokens=config["data"]["action_horizon"], + dropout=model_cfg["dropout"], + )) + + +def train(config_path: str) -> Path: + cfg = load_config(config_path) + root = Path(__file__).resolve().parents[2] + seed_everything(cfg["seed"]) + device = choose_device(cfg.get("device", "auto")) + manifest = materialize_subset( + source_manifest=cfg["data"]["source_manifest"], + output_dir=root / "data" / "overfit", + num_variants=cfg["data"]["num_variants"], + ) + dataset = Pi05OverfitDataset( + manifest, image_size=cfg["data"]["image_size"], + action_horizon=cfg["data"]["action_horizon"] + ) + loader = DataLoader( + dataset, batch_size=min(cfg["train"]["batch_size"], len(dataset)), + shuffle=True, collate_fn=_collate + ) + model = build_model(cfg, dataset).to(device) + flow = FlowMatching(cfg["inference"]["flow_steps"]) + output_dir = root / cfg["output_dir"] + output_dir.mkdir(parents=True, exist_ok=True) + history = [] + + def is_action_parameter(name: str) -> bool: + return any(part in name for part in ( + "action_", "time_embedding", "vector_field_head" + )) + + # Stage 1 follows pi0.5 pre-training: train the VLM/token pathway with + # discrete state and FAST-lite actions, while the flow expert is untouched. + stage1_params = [p for name, p in model.named_parameters() if not is_action_parameter(name)] + optimizer = torch.optim.AdamW( + stage1_params, lr=cfg["train"]["learning_rate"], + weight_decay=cfg["train"]["weight_decay"] + ) + + def next_batch(iterator): + try: + return next(iterator), iterator + except StopIteration: + iterator = iter(loader) + return next(iterator), iterator + + iterator = iter(loader) + model.train() + for step in range(1, cfg["train"]["pretrain_steps"] + 1): + batch, iterator = next_batch(iterator) + image = batch["image"].to(device) + tokens = batch["prefix_tokens"].to(device) + actions = batch["actions"].to(device) + subtask = batch["subtask_target"].to(device)[:, None] + action_tokens = batch["action_tokens"].to(device) + targets = torch.cat((subtask, action_tokens), dim=1) + # Teacher-forced autoregressive stream. The final target is predicted + # from the preceding token and therefore is not part of model_tokens. + model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1) + model_labels = torch.full_like(model_tokens, -100) + query_index = tokens.shape[1] - 1 + model_labels[:, query_index:] = targets + # Flow tokens may read observation/state and the generated semantic + # subtask, but not teacher-forced discrete action tokens. + zero_action = torch.zeros_like(actions) + output = model( + image, model_tokens, zero_action, + torch.zeros(image.shape[0], device=device), + generated_token_start=tokens.shape[1], + text_labels=model_labels, + ) + optimizer.zero_grad(set_to_none=True) + output.loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"]) + optimizer.step() + if step == 1 or step % cfg["train"]["log_every"] == 0: + record = { + "stage": "pretrain", + "step": step, + "loss": float(output.loss.detach().cpu()), + "text_loss": float(output.text_loss.detach().cpu()), + "flow_loss": None, + } + history.append(record) + print(json.dumps(record)) + + torch.save(model.state_dict(), output_dir / "stage1_vlm.pt") + + # Stage 2 follows post-training: add the still-random action expert and + # jointly retain token prediction while optimizing continuous flow. + optimizer = torch.optim.AdamW( + model.parameters(), lr=cfg["train"]["learning_rate"], + weight_decay=cfg["train"]["weight_decay"] + ) + iterator = iter(loader) + for step in range(1, cfg["train"]["posttrain_steps"] + 1): + batch, iterator = next_batch(iterator) + image = batch["image"].to(device) + tokens = batch["prefix_tokens"].to(device) + actions = batch["actions"].to(device) + subtask = batch["subtask_target"].to(device)[:, None] + action_tokens = batch["action_tokens"].to(device) + targets = torch.cat((subtask, action_tokens), dim=1) + model_tokens = torch.cat((tokens, targets[:, :-1]), dim=1) + model_labels = torch.full_like(model_tokens, -100) + model_labels[:, tokens.shape[1] - 1:] = targets + action_prefix_mask = torch.zeros_like(model_tokens, dtype=torch.bool) + action_prefix_mask[:, : tokens.shape[1] + 1] = True + noisy, timestep, velocity = flow.training_batch(actions) + output = model( + image, model_tokens, noisy, timestep, + action_prefix_mask=action_prefix_mask, + generated_token_start=tokens.shape[1], + text_labels=model_labels, target_velocity=velocity, + text_loss_weight=cfg["train"]["semantic_loss_weight"], + flow_loss_weight=cfg["train"]["flow_loss_weight"], + ) + optimizer.zero_grad(set_to_none=True) + output.loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), cfg["train"]["grad_clip_norm"]) + optimizer.step() + if step == 1 or step % cfg["train"]["log_every"] == 0: + record = { + "stage": "posttrain", "step": step, + "loss": float(output.loss.detach().cpu()), + "text_loss": float(output.text_loss.detach().cpu()), + "flow_loss": float(output.flow_loss.detach().cpu()), + } + history.append(record) + print(json.dumps(record)) + checkpoint = { + "model": model.state_dict(), "config": cfg, + "model_config": model.config.__dict__, + "action_statistics": dataset.action_normalizer.as_dict(), + } + torch.save(checkpoint, output_dir / "checkpoint.pt") + write_json(output_dir / "training_history.json", history) + write_json(output_dir / "summary.json", { + "device": str(device), "samples": len(dataset), + "parameters": sum(p.numel() for p in model.parameters()), + "initial": history[0], "final": history[-1], + }) + return output_dir / "checkpoint.pt" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--config", default="configs/overfit.yaml") + args = parser.parse_args() + train(args.config) + + +if __name__ == "__main__": + main() diff --git a/tests/test_action_tokens.py b/tests/test_action_tokens.py new file mode 100644 index 0000000..9025dbe --- /dev/null +++ b/tests/test_action_tokens.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import pytest +import torch + +from pi05mini.tokenizer import MiniTokenizer + + +def test_orthonormal_dct_round_trip() -> None: + tokenizer = MiniTokenizer() + actions = torch.tensor( + [ + [-4.0, -2.75], + [-0.25, 0.0], + [1.375, 3.9], + [4.0, 2.125], + ] + ) + + coefficients = tokenizer.dct_actions(actions) + reconstructed = tokenizer.idct_actions(coefficients) + + assert coefficients.shape == actions.shape + torch.testing.assert_close(reconstructed, actions, atol=1e-5, rtol=1e-5) + + +def test_fast_lite_tokens_round_trip_with_bounded_quantization_error() -> None: + tokenizer = MiniTokenizer() + actions = torch.tensor( + [ + [-3.5, -2.75], + [-0.25, 0.0], + [1.375, 3.5], + [3.75, 2.125], + ] + ) + + tokens = tokenizer.encode_actions(actions) + reconstructed = tokenizer.decode_actions(tokens) + + bin_width = ( + tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM + ) / (tokenizer.ACTION_BINS - 1) + assert reconstructed.shape == actions.shape + # Orthonormal IDCT preserves total coefficient error. Coefficients are + # quantized after division by sqrt(horizon), yielding this RMSE bound. + rmse = torch.sqrt(torch.mean((reconstructed - actions).square())) + assert float(rmse) <= bin_width / 2 * actions.shape[0] ** 0.5 + 1e-6 + + +def test_action_tokens_are_serialized_low_frequency_first() -> None: + tokenizer = MiniTokenizer() + horizon = 4 + coefficients = torch.zeros(horizon, 2) + coefficients[0, 0] = 1.5 * horizon**0.5 # positive X DC coefficient + coefficients[1, 1] = -1.5 * horizon**0.5 # negative Y first harmonic + actions = tokenizer.idct_actions(coefficients) + + payload = tokenizer.encode_actions(actions)[1:-1] + zero_token = tokenizer.ACTION_TOKEN_START + round( + (0.0 - tokenizer.ACTION_MINIMUM) + / (tokenizer.ACTION_MAXIMUM - tokenizer.ACTION_MINIMUM) + * (tokenizer.ACTION_BINS - 1) + ) + + # Frequency-major XY order is [DC-x, DC-y, f1-x, f1-y, ...]. + assert payload[0].item() > zero_token + assert payload[1].item() == zero_token + assert payload[2].item() == zero_token + assert payload[3].item() < zero_token + assert torch.all(payload[4:] == zero_token) + + +def test_action_token_ids_are_contiguous_reserved_vocabulary() -> None: + tokenizer = MiniTokenizer() + edge_actions = torch.tensor([[-4.0, 4.0], [0.0, 1.0]]) + tokens = tokenizer.encode_actions(edge_actions) + + assert tokens[0].item() == tokenizer.ACTION_START + assert tokens[-1].item() == tokenizer.ACTION_END + assert tokens[1:-1].min().item() >= tokenizer.ACTION_TOKEN_START + assert tokens[1:-1].max().item() <= tokenizer.ACTION_TOKEN_END + assert tuple(tokenizer.action_token_ids) == tuple( + range(tokenizer.ACTION_TOKEN_START, tokenizer.ACTION_TOKEN_END + 1) + ) + assert tokenizer.vocab_size == tokenizer.ACTION_TOKEN_END + 1 + + +def test_decode_actions_rejects_missing_boundaries_and_non_action_ids() -> None: + tokenizer = MiniTokenizer() + + with pytest.raises(ValueError, match="enclosed"): + tokenizer.decode_actions(torch.tensor([tokenizer.ACTION_TOKEN_START] * 2)) + with pytest.raises(ValueError, match="reserved"): + tokenizer.decode_actions( + torch.tensor( + [ + tokenizer.ACTION_START, + tokenizer.BOS, + tokenizer.ACTION_TOKEN_START, + tokenizer.ACTION_END, + ] + ) + ) diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..7b23be7 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from pi05mini.data import Pi05OverfitDataset, load_manifest, materialize_subset +from pi05mini.tokenizer import CoordinateTokenizer, MiniTokenizer, SubtaskVocabulary + + +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) + + image_path = image_dir / "sample.jpg" + 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_path) + + 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) + trajectory_path = trajectory_dir / "sample.npz" + np.savez_compressed(trajectory_path, history_xyz=history, future_xyz=future) + + manifest = processed / "overfit_manifest.jsonl" + row = { + "clip_id": "clip-a", + "t0_us": 1234, + "reasoning": "Steer left through the construction lane.", + "image_path": "data/processed/images/sample.jpg", + "trajectory_path": "data/processed/trajectories/sample.npz", + } + manifest.write_text(json.dumps(row) + "\n", encoding="utf-8") + return manifest + + +def test_coordinate_tokens_are_literal_bytes_and_round_trip() -> None: + tokenizer = CoordinateTokenizer(x_range=(-10.0, 10.0), y_range=(-5.0, 5.0)) + coordinates = torch.tensor([[-10.0, -5.0], [0.0, 0.0], [10.0, 5.0]]) + tokens = tokenizer.encode(coordinates) + + assert tokens.dtype == torch.long + assert tokens.tolist()[0] == [0, 0] + assert tokens.tolist()[-1] == [255, 255] + assert 0 <= int(tokens.min()) <= int(tokens.max()) <= 255 + decoded = tokenizer.decode(tokens) + assert torch.allclose(decoded, coordinates, atol=10.1 / 255.0) + + +def test_subtask_vocabulary_and_prefix_contract() -> None: + tokenizer = MiniTokenizer() + history = torch.zeros(8, 2) + prefix = tokenizer.build_prefix(history) + + assert prefix.shape == (20,) + assert prefix[:3].tolist() == [tokenizer.BOS, tokenizer.IMAGE, tokenizer.STATE] + assert prefix[-1].item() == tokenizer.SUBTASK_QUERY + assert prefix[3:-1].min() >= 0 and prefix[3:-1].max() <= 255 + target = tokenizer.encode_subtask("Steer left around the cones") + assert tokenizer.subtasks.decode(target) == "steer_left" + assert SubtaskVocabulary().classify("Slow down for a pedestrian") == "slow_down" + + +def test_materialize_and_dataset_are_self_contained_and_deterministic(tmp_path: Path) -> None: + source_manifest = _make_source_project(tmp_path / "source") + copied_manifest = materialize_subset(source_manifest, tmp_path / "target", max_samples=1) + records = load_manifest(copied_manifest) + + assert len(records) == 1 + assert records[0].image_path.is_file() + assert records[0].trajectory_path.is_file() + + dataset = Pi05OverfitDataset( + copied_manifest, + image_size=(32, 48), + augmentation_factor=3, + seed=19, + ) + same_dataset = Pi05OverfitDataset( + copied_manifest, + image_size=(32, 48), + augmentation_factor=3, + seed=19, + action_statistics=dataset.action_normalizer.as_dict(), + ) + + assert len(dataset) == 3 + original = dataset[0] + augmented = dataset[1] + repeated = same_dataset[1] + assert original["image"].shape == (3, 32, 48) + assert original["prefix_tokens"].shape == (20,) + assert original["actions"].shape == (16, 2) + assert original["subtask_target"].shape == () + assert original["subtask_label"] == "steer_left" + assert original["reasoning"].startswith("Steer left") + 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["future_xy"], repeated["future_xy"]) + assert torch.equal(augmented["image"], repeated["image"]) + + reconstructed = dataset.action_normalizer.decode( + original["actions"], original["history_xy"] + ) + assert torch.allclose(reconstructed, original["future_xy"], atol=1e-5) + + statistics_path = dataset.save_action_statistics(tmp_path / "action_stats.json") + statistics = json.loads(statistics_path.read_text(encoding="utf-8")) + assert len(statistics["mean"]) == 2 + assert len(statistics["scale"]) == 2 + assert min(statistics["scale"]) >= 0.05 diff --git a/tests/test_flow.py b/tests/test_flow.py new file mode 100644 index 0000000..a537242 --- /dev/null +++ b/tests/test_flow.py @@ -0,0 +1,85 @@ +import torch + +from pi05mini.flow import FlowMatching + + +def _expand_batch_time(t: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Expand a per-example time tensor over non-batch dimensions.""" + return t.reshape(t.shape[0], *([1] * (target.ndim - 1))) + + +def test_training_batch_obeys_linear_flow_equations() -> None: + flow = FlowMatching() + target = torch.tensor( + [ + [[1.0, -2.0], [0.5, 3.0]], + [[-1.0, 0.0], [2.5, -0.5]], + ] + ) + generator = torch.Generator().manual_seed(123) + + x_t, t, target_velocity = flow.training_batch(target, generator=generator) + + assert x_t.shape == target.shape + assert target_velocity.shape == target.shape + assert t.shape[0] == target.shape[0] + assert torch.all((t >= 0.0) & (t <= 1.0)) + + t_view = _expand_batch_time(t, target) + # v = target - noise, so the sampled noise is exactly recoverable. + noise = target - target_velocity + expected_x_t = (1.0 - t_view) * noise + t_view * target + torch.testing.assert_close(x_t, expected_x_t) + + +def test_training_batch_is_reproducible_with_generator() -> None: + flow = FlowMatching() + target = torch.randn(3, 5, 2) + + batch_a = flow.training_batch( + target, generator=torch.Generator().manual_seed(99) + ) + batch_b = flow.training_batch( + target, generator=torch.Generator().manual_seed(99) + ) + + for tensor_a, tensor_b in zip(batch_a, batch_b): + torch.testing.assert_close(tensor_a, tensor_b) + + +def test_loss_is_mean_squared_velocity_error() -> None: + flow = FlowMatching() + target_velocity = torch.tensor([[[1.0, -1.0], [0.0, 2.0]]]) + prediction = target_velocity + torch.tensor([[[1.0, 1.0], [2.0, 0.0]]]) + + # Squared errors are [1, 1, 4, 0], whose mean is 1.5. + torch.testing.assert_close(flow.loss(prediction, target_velocity), torch.tensor(1.5)) + torch.testing.assert_close( + flow.loss(target_velocity, target_velocity), torch.tensor(0.0) + ) + + +def test_euler_sampler_reaches_fixed_target_for_exact_vector_field() -> None: + flow = FlowMatching() + target = torch.tensor( + [ + [[0.25, -0.5], [1.0, 2.0]], + [[-2.0, 0.75], [0.0, 0.5]], + ] + ) + + def exact_model(x_t: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + t_view = _expand_batch_time(t, x_t) + # Straight conditional path to target. The sampler evaluates t < 1. + return (target.to(x_t) - x_t) / (1.0 - t_view) + + torch.manual_seed(5) + sampled = flow.sample( + exact_model, + shape=tuple(target.shape), + device=target.device, + steps=8, + ) + + torch.testing.assert_close(sampled, target, atol=1e-5, rtol=1e-5) + diff --git a/tests/test_model.py b/tests/test_model.py new file mode 100644 index 0000000..511734f --- /dev/null +++ b/tests/test_model.py @@ -0,0 +1,61 @@ +import torch + +from pi05mini.model import Pi05Config, Pi05Mini + + +def tiny_model() -> Pi05Mini: + return Pi05Mini(Pi05Config( + vocab_size=300, + image_size=16, + patch_size=16, + d_model=32, + n_heads=4, + n_layers=1, + max_text_tokens=16, + max_action_tokens=4, + )) + + +def test_pi05_has_no_continuous_state_projection() -> None: + model = tiny_model() + assert not hasattr(model, "state_projection") + output = model( + torch.randn(2, 3, 16, 16), + torch.randint(0, 300, (2, 7)), + torch.randn(2, 4, 2), + torch.rand(2), + generated_token_start=4, + ) + assert output.text_logits.shape == (2, 7, 300) + assert output.vector_field.shape == (2, 4, 2) + + +def test_observation_generated_and_flow_attention_partition() -> None: + model = tiny_model() + # One image patch followed by four observation tokens and three generated + # tokens. Flow may read observation plus the first generated (subtask) + # token, but not the two teacher-forced FAST tokens. + action_prefix_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0, 0]], dtype=torch.bool) + prefix, action = model._masks( + batch=1, + prefix_length=8, + action_length=2, + text_offset=1, + attention_mask=torch.ones(1, 7, dtype=torch.bool), + action_prefix_mask=action_prefix_mask, + generated_token_start=4, + device=torch.device("cpu"), + ) + + generated_start = 5 # image patch + four observation text tokens + # Observation block is fully bidirectional and cannot see generated tokens. + assert prefix[0, 0, :generated_start].all() + assert not prefix[0, 0, generated_start:].any() + assert prefix[0, 4, :generated_start].all() + # Generated tokens see all observation tokens and only prior/self generated tokens. + assert prefix[0, 5, :6].all() and not prefix[0, 5, 6:].any() + assert prefix[0, 6, :7].all() and not prefix[0, 6, 7] + # Flow reads allowed prefix keys and all flow-action keys bidirectionally. + assert action[0, :, :6].all() + assert not action[0, :, 6:8].any() + assert action[0, :, 8:].all()