Initial pi0.7 minimal validation
This commit is contained in:
seeker
@@ -0,0 +1,9 @@
|
|||||||
|
__pycache__/
|
||||||
|
*.py[cod]
|
||||||
|
.pytest_cache/
|
||||||
|
.ruff_cache/
|
||||||
|
*.egg-info/
|
||||||
|
.venv/
|
||||||
|
data/
|
||||||
|
outputs/**/*.pt
|
||||||
|
.DS_Store
|
||||||
@@ -0,0 +1,246 @@
|
|||||||
|
# π0.7 Minimal Validation
|
||||||
|
|
||||||
|
这是一个面向**结构验证和小样本过拟合**的 π0.7-style 教学项目。它在前面的
|
||||||
|
π0.6/RECAP 最小实验之上，重点验证一条可解释的 steerable policy 链路：利用
|
||||||
|
短时历史形成 memory，在高层预测视觉 subgoal，通过可独立缺省的 policy
|
||||||
|
metadata 控制动作风格，并在存在控制延迟时用 RTC 对齐动作 chunk。
|
||||||
|
|
||||||
|
本项目不是 Physical Intelligence 官方实现，也不能用于真实车辆或机器人控制。
|
||||||
|
|
||||||
|
## 公开资料边界
|
||||||
|
|
||||||
|
截至本项目实现时，作者没有在 Physical Intelligence 官方网站、公开论文索引或
|
||||||
|
`openpi` 仓库中找到可核验的 π0.7 原始论文、官方训练源码、checkpoint 或完整
|
||||||
|
model card。因此这里不能声称“逐层复现论文”，也不会虚构官方模块尺寸、损失、
|
||||||
|
数据配方或实验数字。
|
||||||
|
|
||||||
|
本仓库所说的 **paper-faithful toy** 有严格边界：若 π0.7 的公开描述强调 memory、
|
||||||
|
goal generation、high-level control、steerable metadata、CFG 和 real-time
|
||||||
|
chunking，本项目保留这些模块的输入输出关系与必要消融；所有未公开细节都使用
|
||||||
|
最小、明确标注且有测试的本地替代。`TinyMEMEncoder`、`TinyGoalGenerator`、
|
||||||
|
`HighLevelPolicy` 和 `Pi07Mini` 都是本仓库命名，不代表官方内部类名。
|
||||||
|
|
||||||
|
若后续公开原始论文或源码，应重新审计本项目，而不是把当前实现当成事实来源。
|
||||||
|
|
||||||
|
## 模型阶段与输入输出
|
||||||
|
|
||||||
|
| 阶段 | 输入 | 模块 | 输出 | 训练或验证目标 |
|
||||||
|
|---|---|---|---|---|
|
||||||
|
| 1. 历史构造 | 当前图像、过去 XY；不接受 future | deterministic history proxy | `6` 帧 oldest-to-newest 图像、历史状态、共享当前状态的 ambiguous history pair | 检查无未来轨迹泄漏 |
|
||||||
|
| 2. Temporal memory | 历史图像与历史状态 | `TinyMEMEncoder` | memory tokens / pooled memory embedding | 区分当前状态相同、过去路径不同的样本 |
|
||||||
|
| 3. Goal generation | 当前 observation、memory、任务文本 | `TinyGoalGenerator` | endpoint heatmap / subgoal XY | 只监督 future endpoint，不把完整真值轨迹输入 policy |
|
||||||
|
| 4. High-level policy | observation、memory、任务与可见 metadata | `HighLevelPolicy` | high-level context、goal/metadata logits | goal loss 与 high-level classification loss |
|
||||||
|
| 5. Steerable flow policy | high-level context、预测 goal、metadata、noisy action、flow time | `Pi07Mini` low-level action expert | 连续 action vector field `[16,2]` | flow-matching loss |
|
||||||
|
| 6. RTC alignment | action chunk、随机 delay `d∈[0,12]`、已执行 action prefix | real-time chunking target alignment | 从 `action[d]` 开始的 target 与有效位置 mask | 仅在有效位置计算 masked flow loss |
|
||||||
|
| 7. CFG inference | 同一 observation 的 conditional 与 metadata-omitted 两次前向 | classifier-free guidance | guided action chunk | 与普通 conditional、unconditional 和真值同图比较 |
|
||||||
|
| 8. 轨迹恢复 | normalized action deltas、历史终点、训练统计 | denormalize + cumulative sum | 未来 XY 轨迹 | 轨迹 MSE 与可视化，不等同安全指标 |
|
||||||
|
|
||||||
|
核心信息流：
|
||||||
|
|
||||||
|
```text
|
||||||
|
history frames + history states
|
||||||
|
↓
|
||||||
|
TinyMEMEncoder ───────────────┐
|
||||||
|
│
|
||||||
|
current image + task ────────────────┼→ TinyGoalGenerator → predicted subgoal
|
||||||
|
│
|
||||||
|
optional speed/quality/mistake/control metadata
|
||||||
|
↓
|
||||||
|
HighLevelPolicy
|
||||||
|
↓
|
||||||
|
noisy action + flow time + RTC prefix → Pi07Mini action expert
|
||||||
|
↓
|
||||||
|
continuous action chunk
|
||||||
|
```
|
||||||
|
|
||||||
|
## 重点结构
|
||||||
|
|
||||||
|
### 1. MEM：短时历史而非单帧
|
||||||
|
|
||||||
|
PAI 的极小 manifest 只有一张相机图像，本项目依据已观测的过去 XY 对该图像做
|
||||||
|
确定性的微小平移和曝光变化，构造短时视觉 history proxy。它不是实际多帧视频，
|
||||||
|
只能验证 temporal memory 的接口。
|
||||||
|
|
||||||
|
`ambiguous_history_pair` 刻意构造两个“当前状态相同、历史路径不同”的序列。
|
||||||
|
完整模型应比 `no_history` 消融更能区分它们。这里的 MEM 仅表示本仓库的 temporal
|
||||||
|
memory encoder；在官方资料缺失时，不推断该缩写的官方全称。
|
||||||
|
|
||||||
|
### 2. Goal generator：只用 endpoint 监督
|
||||||
|
|
||||||
|
训练目标是 future trajectory 的最后一个 XY endpoint，经高斯渲染得到 heatmap。
|
||||||
|
API 只接受 endpoint，不接受整条 future trajectory，从而避免把低层动作答案作为
|
||||||
|
视觉 subgoal 条件泄漏给 action expert。推理和最终评估必须使用**模型预测的 goal**；
|
||||||
|
真值 heatmap 只可用于训练监督和 oracle 诊断。
|
||||||
|
|
||||||
|
### 3. High-level 与 steerable metadata
|
||||||
|
|
||||||
|
最小数据由一条 expert demonstration 确定性构造五种 policy mode：
|
||||||
|
|
||||||
|
| Mode | speed | quality | mistake | control |
|
||||||
|
|---|---|---|---|---|
|
||||||
|
| expert | normal | expert | none | human |
|
||||||
|
| fast | fast | good | none | autonomous |
|
||||||
|
| slow | slow | good | none | autonomous |
|
||||||
|
| understeer | normal | poor | understeer | autonomous |
|
||||||
|
| stop | slow | poor | stop | autonomous |
|
||||||
|
|
||||||
|
训练时 `speed`、`quality`、`mistake`、`control` 独立做 metadata dropout，任务文本
|
||||||
|
始终保留。这既检验属性是否能单独控制行为，也提供 CFG 所需的 omitted condition。
|
||||||
|
这些合成标签不是人类偏好数据，也不能证明真实世界可控性。
|
||||||
|
|
||||||
|
### 4. RTC：控制延迟下的 action chunk 对齐
|
||||||
|
|
||||||
|
若推理耗时等价于已经执行 `d` 个 action，模型不应再次预测旧 chunk 的开头：
|
||||||
|
|
||||||
|
```text
|
||||||
|
original target: a0 a1 ... a15
|
||||||
|
delay d=4: a4 a5 ... a15 PAD PAD PAD PAD
|
||||||
|
valid mask: 1 1 ... 1 0 0 0 0
|
||||||
|
```
|
||||||
|
|
||||||
|
默认随机采样 `d=0..12`，返回已执行 prefix，并对 padding tail 屏蔽损失。RTC 在
|
||||||
|
这里表示训练/评估时的离线对齐机制，不包含实际控制器、异步推理线程或硬实时保证。
|
||||||
|
|
||||||
|
### 5. CFG
|
||||||
|
|
||||||
|
推理时在相同 observation、goal、memory 和初始噪声下计算 metadata conditional
|
||||||
|
与 metadata-omitted 预测：
|
||||||
|
|
||||||
|
```text
|
||||||
|
v_cfg = v_uncond + beta * (v_cond - v_uncond)
|
||||||
|
```
|
||||||
|
|
||||||
|
默认 `beta=2.0`。必须同时报告 `beta=0,1,2`，因为高 guidance 可能放大过拟合、
|
||||||
|
产生分布外动作，不能只展示最好看的曲线。
|
||||||
|
|
||||||
|
## 数据与防泄漏约束
|
||||||
|
|
||||||
|
- history proxy 和 ambiguous pair 只依赖当前图像与过去状态；
|
||||||
|
- subgoal condition 只使用 endpoint target，不能编码完整 action chunk；
|
||||||
|
- flow expert 训练时不能读取 teacher-forced 真值动作 token；
|
||||||
|
- 所有模型对比使用同一 initial noise、归一化统计和 RTC delay；
|
||||||
|
- 数据增强的随机数由 sample index 和 seed 决定，保证复现实验；
|
||||||
|
- 仓库不包含或重新分发原始 PAI 数据。
|
||||||
|
|
||||||
|
## 安装
|
||||||
|
|
||||||
|
```bash
|
||||||
|
cd pi07-minimal-validation
|
||||||
|
python -m venv .venv
|
||||||
|
source .venv/bin/activate
|
||||||
|
python -m pip install -U pip
|
||||||
|
python -m pip install -e '.[dev]'
|
||||||
|
```
|
||||||
|
|
||||||
|
默认数据路径：
|
||||||
|
|
||||||
|
```text
|
||||||
|
../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl
|
||||||
|
```
|
||||||
|
|
||||||
|
若路径不同，请修改 `configs/overfit.yaml` 的 `data.source_manifest`。
|
||||||
|
|
||||||
|
## 测试
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pytest -q
|
||||||
|
```
|
||||||
|
|
||||||
|
测试至少应覆盖：history/goal 无 future 泄漏、ambiguous history、metadata 独立
|
||||||
|
dropout、五类策略构造、action normalization、flow matching、RTC shift/mask、
|
||||||
|
MEM/goal/high-level/policy 前向形状、CFG 公式和 checkpoint reload。
|
||||||
|
|
||||||
|
## 三阶段过拟合训练
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pi07mini-train --config configs/overfit.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
或：
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m pi07mini.train --config configs/overfit.yaml
|
||||||
|
```
|
||||||
|
|
||||||
|
| 阶段 | 默认步数 | 更新内容 | 目的 |
|
||||||
|
|---|---:|---|---|
|
||||||
|
| goal | 300 | MEM + goal generator | 先证明 history-conditioned endpoint heatmap 可被记忆 |
|
||||||
|
| high-level | 200 | MEM + high-level policy | 记住五类 metadata/goal 高层条件 |
|
||||||
|
| policy | 900 | 完整 `Pi07Mini` | 联合训练 metadata dropout、flow action 与 RTC masked target |
|
||||||
|
|
||||||
|
默认预算面向极小样本过拟合，不代表官方训练 schedule。训练完成后应保存实际
|
||||||
|
训练会生成 loss history、checkpoint、evaluation JSON 和消融效果图。
|
||||||
|
|
||||||
|
## 当前最小实验结果
|
||||||
|
|
||||||
|
固定相同 flow 初始噪声时，高层策略预测 `steer_left`。轨迹 MSE 如下：
|
||||||
|
|
||||||
|
| 条件 | MSE |
|
||||||
|
|---|---:|
|
||||||
|
| Full context + oracle endpoint goal | 0.470 |
|
||||||
|
| Full context + generated goal | 0.491 |
|
||||||
|
| No metadata | 2.458 |
|
||||||
|
| No goal | 4.043 |
|
||||||
|
| No history | 2.159 |
|
||||||
|
| Metadata CFG beta=2 | 10.576 |
|
||||||
|
|
||||||
|
完整上下文明显优于移除 metadata、goal 或 history 的消融。TinyGoalGenerator 的像素级结果仍较粗糙；生成 goal 在当前 policy 上可提供有效条件，但不能视为真实未来图像生成能力。较强 CFG 在此极小数据实验中造成退化。
|
||||||
|
|
||||||
|
## 推理、画图与必要消融
|
||||||
|
|
||||||
|
```bash
|
||||||
|
pi07mini-infer --checkpoint outputs/overfit/checkpoint.pt
|
||||||
|
```
|
||||||
|
|
||||||
|
或：
|
||||||
|
|
||||||
|
```bash
|
||||||
|
python -m pi07mini.infer --checkpoint outputs/overfit/checkpoint.pt
|
||||||
|
```
|
||||||
|
|
||||||
|
最终展示不能只有 full model。相同样本、noise seed 和 RTC delay 下至少比较：
|
||||||
|
|
||||||
|
| 实验 | 移除内容 | 要回答的问题 |
|
||||||
|
|---|---|---|
|
||||||
|
| `full` | 无 | 完整链路能否过拟合并重建对应 mode 的轨迹？ |
|
||||||
|
| `no_metadata` | speed/quality/mistake/control | policy mode 差异是否真的来自 steerable condition？ |
|
||||||
|
| `no_goal` | predicted goal | endpoint condition 是否改善几何方向和终点？ |
|
||||||
|
| `no_history` | MEM history，只保留当前 observation | temporal history 是否解决相同当前状态的歧义？ |
|
||||||
|
|
||||||
|
还必须提供：
|
||||||
|
|
||||||
|
- predicted goal heatmap 与 ground-truth endpoint heatmap；
|
||||||
|
- expert/fast/slow/understeer/stop 的真值与预测轨迹；
|
||||||
|
- CFG `beta=0,1,2`；
|
||||||
|
- RTC delay `0,4,8,12` 的 masked trajectory error；
|
||||||
|
- 每个 ablation 的轨迹 MSE，以及 goal endpoint error；
|
||||||
|
- full/no-history 在 ambiguous history pair 上的差异指标。
|
||||||
|
|
||||||
|
小样本过拟合中，某个消融偶然更好并不证明模块无效；应展示全部固定设置结果，
|
||||||
|
不调换 noise seed 或只挑选最佳样本。
|
||||||
|
|
||||||
|
## 安全限制
|
||||||
|
|
||||||
|
- 严禁连接车辆、机械臂、移动机器人或任何会造成物理影响的执行器。
|
||||||
|
- 合成 fast/understeer/stop 轨迹只用于离线可控性探针，不是安全驾驶示范。
|
||||||
|
- endpoint 命中、轨迹 MSE 下降、loss 收敛和 CFG 改善都不代表闭环稳定性。
|
||||||
|
- history proxy 不是视频；RTC target shift 也不是硬实时系统验证。
|
||||||
|
- metadata 可控性可能放大危险属性，尤其是 `fast`、`poor` 或高 CFG beta。
|
||||||
|
- 真实部署至少需要真实时序传感器、闭环 rollout、OOD 与延迟压力测试、动力学
|
||||||
|
约束、碰撞检查、冗余控制器、故障降级、急停和人工监督。
|
||||||
|
- PyTorch checkpoint 使用 pickle 机制，只加载可信来源的文件。
|
||||||
|
|
||||||
|
## 能得出的结论
|
||||||
|
|
||||||
|
成功实验最多能说明：本仓库定义的 memory、endpoint goal、high-level metadata、
|
||||||
|
conditional flow、CFG 与 RTC 对齐可以在几个合成 PAI trajectory mode 上形成一条
|
||||||
|
可训练、可消融的最小链路。它不能证明与未公开 π0.7 的结构或数值等价，不能
|
||||||
|
证明 emergent capability，也不能支持任何真实世界安全或泛化主张。
|
||||||
|
|
||||||
|
## 可核验背景资料
|
||||||
|
|
||||||
|
- [Physical Intelligence](https://www.pi.website/)
|
||||||
|
- [Physical-Intelligence/openpi](https://github.com/Physical-Intelligence/openpi)
|
||||||
|
|
||||||
|
在原始 π0.7 论文或官方源码可核验前，本 README 不提供推测性的论文链接、arXiv
|
||||||
|
编号或官方结果表。
|
||||||
@@ -0,0 +1,31 @@
|
|||||||
|
# Third-party notices
|
||||||
|
|
||||||
|
This repository is an independent educational implementation. It does not
|
||||||
|
contain code, weights, datasets, or other assets copied from Physical
|
||||||
|
Intelligence.
|
||||||
|
|
||||||
|
At the time this project was prepared, no original pi0.7 paper, official
|
||||||
|
pi0.7 training implementation, checkpoint, or complete technical model card
|
||||||
|
was available in the public Physical Intelligence/openpi materials inspected
|
||||||
|
by the authors. Consequently, names such as `TinyMEMEncoder`,
|
||||||
|
`TinyGoalGenerator`, `HighLevelPolicy`, and `Pi07Mini` describe this
|
||||||
|
repository's own teaching implementation. They must not be interpreted as
|
||||||
|
official class names or proof of the unpublished system's internals.
|
||||||
|
|
||||||
|
The project uses the earlier public pi-family literature and public openpi
|
||||||
|
repository only as background for vision-language-action policies, action
|
||||||
|
chunking, and flow matching. No source file from openpi is vendored here.
|
||||||
|
|
||||||
|
Runtime dependencies retain their own licenses and copyright notices:
|
||||||
|
|
||||||
|
- PyTorch
|
||||||
|
- NumPy
|
||||||
|
- Pillow
|
||||||
|
- Matplotlib
|
||||||
|
- PyYAML
|
||||||
|
- pytest (development only)
|
||||||
|
|
||||||
|
Users are responsible for complying with the licenses and terms of any
|
||||||
|
external images, trajectories, datasets, or checkpoints they supply. The
|
||||||
|
default configuration points to a local manifest and this repository does not
|
||||||
|
redistribute the source PAI data.
|
||||||
@@ -0,0 +1,98 @@
|
|||||||
|
seed: 7
|
||||||
|
device: auto
|
||||||
|
output_dir: outputs/overfit
|
||||||
|
|
||||||
|
data:
|
||||||
|
source_manifest: ../../alpamayo/alpamayo-qwen05b-validation/data/processed/overfit_manifest.jsonl
|
||||||
|
max_samples: 1
|
||||||
|
augmentation_factor: 1
|
||||||
|
image_size: 64
|
||||||
|
history_frames: 6
|
||||||
|
action_horizon: 16
|
||||||
|
policy_variants:
|
||||||
|
- expert
|
||||||
|
- fast
|
||||||
|
- slow
|
||||||
|
- understeer
|
||||||
|
- stop
|
||||||
|
|
||||||
|
model:
|
||||||
|
patch_size: 16
|
||||||
|
d_model: 96
|
||||||
|
mem_d_model: 64
|
||||||
|
goal_d_model: 64
|
||||||
|
high_level_d_model: 64
|
||||||
|
n_heads: 4
|
||||||
|
n_layers: 3
|
||||||
|
mlp_ratio: 3
|
||||||
|
dropout: 0.0
|
||||||
|
max_text_tokens: 96
|
||||||
|
action_dim: 2
|
||||||
|
action_horizon: 16
|
||||||
|
goal_heatmap_size: 32
|
||||||
|
|
||||||
|
memory:
|
||||||
|
enabled: true
|
||||||
|
history_frames: 6
|
||||||
|
ambiguous_pair_weight: 0.25
|
||||||
|
|
||||||
|
goal:
|
||||||
|
enabled: true
|
||||||
|
heatmap_size: 32
|
||||||
|
heatmap_sigma: 1.5
|
||||||
|
xy_extent: [-8.0, 64.0, -16.0, 16.0]
|
||||||
|
endpoint_only_target: true
|
||||||
|
|
||||||
|
high_level:
|
||||||
|
enabled: true
|
||||||
|
predict_goal: true
|
||||||
|
predict_metadata: true
|
||||||
|
metadata_fields:
|
||||||
|
- speed
|
||||||
|
- quality
|
||||||
|
- mistake
|
||||||
|
- control
|
||||||
|
|
||||||
|
steering:
|
||||||
|
metadata_dropout: 0.25
|
||||||
|
preserve_task_prompt: true
|
||||||
|
cfg_training_dropout: 0.25
|
||||||
|
cfg_beta: 2.0
|
||||||
|
|
||||||
|
rtc:
|
||||||
|
enabled: true
|
||||||
|
max_delay_steps: 12
|
||||||
|
include_executed_prefix: true
|
||||||
|
masked_loss: true
|
||||||
|
|
||||||
|
flow:
|
||||||
|
timestep_beta_alpha: 1.5
|
||||||
|
timestep_beta_beta: 1.0
|
||||||
|
integration_steps: 10
|
||||||
|
|
||||||
|
train:
|
||||||
|
goal_steps: 300
|
||||||
|
high_level_steps: 200
|
||||||
|
policy_steps: 900
|
||||||
|
batch_size: 5
|
||||||
|
learning_rate: 0.001
|
||||||
|
weight_decay: 0.0
|
||||||
|
grad_clip_norm: 1.0
|
||||||
|
goal_loss_weight: 1.0
|
||||||
|
high_level_loss_weight: 1.0
|
||||||
|
flow_loss_weight: 2.0
|
||||||
|
rtc_loss_weight: 1.0
|
||||||
|
log_every: 20
|
||||||
|
|
||||||
|
evaluation:
|
||||||
|
fixed_noise_seed: 1
|
||||||
|
cfg_betas: [0.0, 1.0, 2.0]
|
||||||
|
ablations:
|
||||||
|
- full
|
||||||
|
- no_metadata
|
||||||
|
- no_goal
|
||||||
|
- no_history
|
||||||
|
compare_rtc_delays: [0, 4, 8, 12]
|
||||||
|
save_trajectory_plot: true
|
||||||
|
save_goal_plot: true
|
||||||
|
save_ablation_plot: true
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 92 KiB |
@@ -0,0 +1,11 @@
|
|||||||
|
{
|
||||||
|
"predicted_subtask": "steer_left",
|
||||||
|
"trajectory_metrics": {
|
||||||
|
"full_oracle_goal": 0.47031369805336,
|
||||||
|
"generated_goal": 0.4905388057231903,
|
||||||
|
"no_metadata": 2.457864284515381,
|
||||||
|
"no_goal": 4.042560577392578,
|
||||||
|
"no_history": 2.1593737602233887,
|
||||||
|
"cfg_beta2": 10.576497077941895
|
||||||
|
}
|
||||||
|
}
|
||||||
Binary file not shown.
|
After Width: | Height: | Size: 50 KiB |
@@ -0,0 +1,7 @@
|
|||||||
|
{
|
||||||
|
"device": "mps",
|
||||||
|
"samples": 5,
|
||||||
|
"policy_parameters": 1088930,
|
||||||
|
"goal_parameters": 75731,
|
||||||
|
"high_level_parameters": 652800
|
||||||
|
}
|
||||||
@@ -0,0 +1,459 @@
|
|||||||
|
[
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 1,
|
||||||
|
"loss": 1.0040572881698608
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 20,
|
||||||
|
"loss": 0.34230467677116394
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 40,
|
||||||
|
"loss": 0.37451064586639404
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 60,
|
||||||
|
"loss": 0.26792070269584656
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 80,
|
||||||
|
"loss": 0.10219617933034897
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 100,
|
||||||
|
"loss": 0.2711445987224579
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 120,
|
||||||
|
"loss": 0.2289590686559677
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 140,
|
||||||
|
"loss": 0.36456483602523804
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 160,
|
||||||
|
"loss": 0.0842418223619461
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 180,
|
||||||
|
"loss": 0.9889341592788696
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 200,
|
||||||
|
"loss": 0.5776903629302979
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 220,
|
||||||
|
"loss": 0.44620272517204285
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 240,
|
||||||
|
"loss": 0.05634145438671112
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 260,
|
||||||
|
"loss": 0.05701519176363945
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 280,
|
||||||
|
"loss": 0.06764699518680573
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "goal",
|
||||||
|
"step": 300,
|
||||||
|
"loss": 0.10973378270864487
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 1,
|
||||||
|
"loss": 76.15934753417969
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 20,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 40,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 60,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 80,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 100,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 120,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 140,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 160,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 180,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "high_level",
|
||||||
|
"step": 200,
|
||||||
|
"loss": -0.0
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 1,
|
||||||
|
"loss": 51.04338836669922,
|
||||||
|
"text_loss": 44.21609115600586,
|
||||||
|
"flow_loss": 3.4136481285095215
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 20,
|
||||||
|
"loss": 26.20608901977539,
|
||||||
|
"text_loss": 5.128565311431885,
|
||||||
|
"flow_loss": 10.538762092590332
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 40,
|
||||||
|
"loss": 13.018914222717285,
|
||||||
|
"text_loss": 3.846059799194336,
|
||||||
|
"flow_loss": 4.586427211761475
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 60,
|
||||||
|
"loss": 11.344329833984375,
|
||||||
|
"text_loss": 3.335771083831787,
|
||||||
|
"flow_loss": 4.004279136657715
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 80,
|
||||||
|
"loss": 6.586862087249756,
|
||||||
|
"text_loss": 2.5455708503723145,
|
||||||
|
"flow_loss": 2.0206456184387207
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 100,
|
||||||
|
"loss": 11.200895309448242,
|
||||||
|
"text_loss": 2.1739449501037598,
|
||||||
|
"flow_loss": 4.513474941253662
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 120,
|
||||||
|
"loss": 17.273035049438477,
|
||||||
|
"text_loss": 2.1746408939361572,
|
||||||
|
"flow_loss": 7.549196720123291
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 140,
|
||||||
|
"loss": 11.353687286376953,
|
||||||
|
"text_loss": 2.2218563556671143,
|
||||||
|
"flow_loss": 4.565915584564209
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 160,
|
||||||
|
"loss": 14.110220909118652,
|
||||||
|
"text_loss": 2.3123977184295654,
|
||||||
|
"flow_loss": 5.898911476135254
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 180,
|
||||||
|
"loss": 19.999370574951172,
|
||||||
|
"text_loss": 2.749119997024536,
|
||||||
|
"flow_loss": 8.62512493133545
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 200,
|
||||||
|
"loss": 8.102465629577637,
|
||||||
|
"text_loss": 1.631801962852478,
|
||||||
|
"flow_loss": 3.2353317737579346
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 220,
|
||||||
|
"loss": 21.903854370117188,
|
||||||
|
"text_loss": 1.6063504219055176,
|
||||||
|
"flow_loss": 10.148752212524414
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 240,
|
||||||
|
"loss": 10.197904586791992,
|
||||||
|
"text_loss": 1.5640149116516113,
|
||||||
|
"flow_loss": 4.3169450759887695
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 260,
|
||||||
|
"loss": 6.899936676025391,
|
||||||
|
"text_loss": 1.3081820011138916,
|
||||||
|
"flow_loss": 2.795877456665039
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 280,
|
||||||
|
"loss": 10.051774978637695,
|
||||||
|
"text_loss": 1.2130903005599976,
|
||||||
|
"flow_loss": 4.419342517852783
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 300,
|
||||||
|
"loss": 7.963553428649902,
|
||||||
|
"text_loss": 1.1586928367614746,
|
||||||
|
"flow_loss": 3.402430295944214
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 320,
|
||||||
|
"loss": 5.177718639373779,
|
||||||
|
"text_loss": 1.2673476934432983,
|
||||||
|
"flow_loss": 1.9551855325698853
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 340,
|
||||||
|
"loss": 4.487771034240723,
|
||||||
|
"text_loss": 1.0361082553863525,
|
||||||
|
"flow_loss": 1.7258312702178955
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 360,
|
||||||
|
"loss": 4.356203556060791,
|
||||||
|
"text_loss": 1.3223515748977661,
|
||||||
|
"flow_loss": 1.5169259309768677
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 380,
|
||||||
|
"loss": 5.289548397064209,
|
||||||
|
"text_loss": 1.4994667768478394,
|
||||||
|
"flow_loss": 1.89504075050354
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 400,
|
||||||
|
"loss": 8.761286735534668,
|
||||||
|
"text_loss": 1.4391229152679443,
|
||||||
|
"flow_loss": 3.6610820293426514
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 420,
|
||||||
|
"loss": 2.400188446044922,
|
||||||
|
"text_loss": 1.0106937885284424,
|
||||||
|
"flow_loss": 0.694747269153595
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 440,
|
||||||
|
"loss": 9.683706283569336,
|
||||||
|
"text_loss": 1.044368863105774,
|
||||||
|
"flow_loss": 4.319668769836426
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 460,
|
||||||
|
"loss": 10.443679809570312,
|
||||||
|
"text_loss": 0.7045990228652954,
|
||||||
|
"flow_loss": 4.869540214538574
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 480,
|
||||||
|
"loss": 6.514384746551514,
|
||||||
|
"text_loss": 1.185524582862854,
|
||||||
|
"flow_loss": 2.6644301414489746
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 500,
|
||||||
|
"loss": 2.255455255508423,
|
||||||
|
"text_loss": 1.0136079788208008,
|
||||||
|
"flow_loss": 0.620923638343811
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 520,
|
||||||
|
"loss": 4.1700439453125,
|
||||||
|
"text_loss": 1.1829472780227661,
|
||||||
|
"flow_loss": 1.4935483932495117
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 540,
|
||||||
|
"loss": 6.956455230712891,
|
||||||
|
"text_loss": 1.1096038818359375,
|
||||||
|
"flow_loss": 2.9234256744384766
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 560,
|
||||||
|
"loss": 4.5965800285339355,
|
||||||
|
"text_loss": 1.5257400274276733,
|
||||||
|
"flow_loss": 1.5354200601577759
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 580,
|
||||||
|
"loss": 1.7948781251907349,
|
||||||
|
"text_loss": 0.835124135017395,
|
||||||
|
"flow_loss": 0.4798769950866699
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 600,
|
||||||
|
"loss": 11.411609649658203,
|
||||||
|
"text_loss": 0.7508105635643005,
|
||||||
|
"flow_loss": 5.330399513244629
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 620,
|
||||||
|
"loss": 14.173189163208008,
|
||||||
|
"text_loss": 0.553065836429596,
|
||||||
|
"flow_loss": 6.810061454772949
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 640,
|
||||||
|
"loss": 2.423476457595825,
|
||||||
|
"text_loss": 0.6566628813743591,
|
||||||
|
"flow_loss": 0.8834067583084106
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 660,
|
||||||
|
"loss": 7.33753776550293,
|
||||||
|
"text_loss": 0.5569484829902649,
|
||||||
|
"flow_loss": 3.3902945518493652
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 680,
|
||||||
|
"loss": 7.2152018547058105,
|
||||||
|
"text_loss": 0.6606324911117554,
|
||||||
|
"flow_loss": 3.277284622192383
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 700,
|
||||||
|
"loss": 2.5858616828918457,
|
||||||
|
"text_loss": 0.471678227186203,
|
||||||
|
"flow_loss": 1.0570917129516602
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 720,
|
||||||
|
"loss": 5.87769889831543,
|
||||||
|
"text_loss": 0.5711548924446106,
|
||||||
|
"flow_loss": 2.6532719135284424
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 740,
|
||||||
|
"loss": 6.787323951721191,
|
||||||
|
"text_loss": 0.5064371228218079,
|
||||||
|
"flow_loss": 3.1404433250427246
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 760,
|
||||||
|
"loss": 11.166234970092773,
|
||||||
|
"text_loss": 0.5802748203277588,
|
||||||
|
"flow_loss": 5.292980194091797
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 780,
|
||||||
|
"loss": 3.563460350036621,
|
||||||
|
"text_loss": 0.6946841478347778,
|
||||||
|
"flow_loss": 1.4343880414962769
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 800,
|
||||||
|
"loss": 3.2471487522125244,
|
||||||
|
"text_loss": 0.44556593894958496,
|
||||||
|
"flow_loss": 1.4007914066314697
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 820,
|
||||||
|
"loss": 3.1103086471557617,
|
||||||
|
"text_loss": 0.7132699489593506,
|
||||||
|
"flow_loss": 1.1985193490982056
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 840,
|
||||||
|
"loss": 4.753307342529297,
|
||||||
|
"text_loss": 0.48397523164749146,
|
||||||
|
"flow_loss": 2.1346659660339355
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 860,
|
||||||
|
"loss": 3.5492496490478516,
|
||||||
|
"text_loss": 0.366057813167572,
|
||||||
|
"flow_loss": 1.5915958881378174
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 880,
|
||||||
|
"loss": 4.0379414558410645,
|
||||||
|
"text_loss": 0.6511648893356323,
|
||||||
|
"flow_loss": 1.6933883428573608
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"stage": "policy",
|
||||||
|
"step": 900,
|
||||||
|
"loss": 1.7170453071594238,
|
||||||
|
"text_loss": 0.2673894464969635,
|
||||||
|
"flow_loss": 0.7248279452323914
|
||||||
|
}
|
||||||
|
]
|
||||||
@@ -0,0 +1,38 @@
|
|||||||
|
[build-system]
|
||||||
|
requires = ["setuptools>=69", "wheel"]
|
||||||
|
build-backend = "setuptools.build_meta"
|
||||||
|
|
||||||
|
[project]
|
||||||
|
name = "pi07-minimal-validation"
|
||||||
|
version = "0.1.0"
|
||||||
|
description = "A tiny, safety-scoped pi0.7-style steerable policy validation experiment"
|
||||||
|
readme = "README.md"
|
||||||
|
requires-python = ">=3.10"
|
||||||
|
license = { text = "MIT" }
|
||||||
|
authors = [{ name = "pi07 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]
|
||||||
|
pi07mini-train = "pi07mini.train:main"
|
||||||
|
pi07mini-infer = "pi07mini.infer:main"
|
||||||
|
|
||||||
|
[tool.setuptools]
|
||||||
|
package-dir = { "" = "src" }
|
||||||
|
|
||||||
|
[tool.setuptools.packages.find]
|
||||||
|
where = ["src"]
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
testpaths = ["tests"]
|
||||||
|
addopts = "-ra"
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
"""Minimal pi0.7 architecture validation package."""
|
||||||
|
|
||||||
|
from .memory import MEMConfig, MEMVideoEncoder
|
||||||
|
from .model import (
|
||||||
|
GoalGeneratorOutput,
|
||||||
|
HighLevelOutput,
|
||||||
|
HighLevelPolicy,
|
||||||
|
Pi07Config,
|
||||||
|
Pi07Mini,
|
||||||
|
Pi07Output,
|
||||||
|
PrefixLayout,
|
||||||
|
TinyGoalGenerator,
|
||||||
|
flow_matching_sample,
|
||||||
|
)
|
||||||
|
from .tokenizer import EpisodeMetadata, MiniTokenizer, PolicySequence
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"EpisodeMetadata",
|
||||||
|
"GoalGeneratorOutput",
|
||||||
|
"HighLevelOutput",
|
||||||
|
"HighLevelPolicy",
|
||||||
|
"MEMConfig",
|
||||||
|
"MEMVideoEncoder",
|
||||||
|
"MiniTokenizer",
|
||||||
|
"Pi07Config",
|
||||||
|
"Pi07Mini",
|
||||||
|
"Pi07Output",
|
||||||
|
"PolicySequence",
|
||||||
|
"PrefixLayout",
|
||||||
|
"TinyGoalGenerator",
|
||||||
|
"flow_matching_sample",
|
||||||
|
]
|
||||||
@@ -0,0 +1,209 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Mapping, Sequence
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
|
||||||
|
PROMPT_COMPONENT_ORDER = ("task", "speed", "quality", "mistake", "control")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PolicyMetadata:
|
||||||
|
"""Steerable policy attributes used by the tiny pi0.7 experiment."""
|
||||||
|
|
||||||
|
speed: str
|
||||||
|
quality: str
|
||||||
|
mistake: str
|
||||||
|
control: str
|
||||||
|
speed_steps: int = 1_000
|
||||||
|
quality_score: int = 3
|
||||||
|
control_mode: str = "ee"
|
||||||
|
source_type: str = "autonomous"
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if self.speed_steps <= 0:
|
||||||
|
raise ValueError("speed_steps must be positive")
|
||||||
|
if not 1 <= self.quality_score <= 5:
|
||||||
|
raise ValueError("quality_score must lie in [1,5]")
|
||||||
|
if self.control_mode not in {"joint", "ee"}:
|
||||||
|
raise ValueError("control_mode must be joint or ee")
|
||||||
|
|
||||||
|
def as_dict(self) -> dict[str, str]:
|
||||||
|
return {
|
||||||
|
"speed": self.speed,
|
||||||
|
"quality": self.quality,
|
||||||
|
"mistake": self.mistake,
|
||||||
|
"control": self.control,
|
||||||
|
}
|
||||||
|
|
||||||
|
def as_token_metadata(self) -> dict[str, int | bool | str]:
|
||||||
|
"""Return the scalar metadata contract consumed by MiniTokenizer."""
|
||||||
|
|
||||||
|
return {
|
||||||
|
"speed_steps": self.speed_steps,
|
||||||
|
"quality": self.quality_score,
|
||||||
|
"mistake": self.mistake != "none",
|
||||||
|
"control_mode": self.control_mode,
|
||||||
|
"source_type": self.source_type,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def prompt_components(
|
||||||
|
metadata: PolicyMetadata,
|
||||||
|
*,
|
||||||
|
task: str = "Follow the visible road and predict the next waypoint chunk.",
|
||||||
|
) -> dict[str, str]:
|
||||||
|
return {"task": task, **metadata.as_dict()}
|
||||||
|
|
||||||
|
|
||||||
|
def drop_prompt_components(
|
||||||
|
components: Mapping[str, str],
|
||||||
|
*,
|
||||||
|
drop_probability: float,
|
||||||
|
generator: torch.Generator | None = None,
|
||||||
|
preserve: Sequence[str] = ("task",),
|
||||||
|
) -> tuple[dict[str, str], torch.Tensor]:
|
||||||
|
"""Apply classifier-free-style dropout to independent prompt attributes.
|
||||||
|
|
||||||
|
The returned mask follows ``PROMPT_COMPONENT_ORDER`` and marks components
|
||||||
|
that remain visible. Task text is preserved by default, while metadata can
|
||||||
|
be independently omitted so the policy cannot require every attribute.
|
||||||
|
"""
|
||||||
|
|
||||||
|
if not 0.0 <= drop_probability <= 1.0:
|
||||||
|
raise ValueError("drop_probability must lie in [0, 1]")
|
||||||
|
preserved = set(preserve)
|
||||||
|
visible: dict[str, str] = {}
|
||||||
|
mask = torch.zeros(len(PROMPT_COMPONENT_ORDER), dtype=torch.bool)
|
||||||
|
for index, name in enumerate(PROMPT_COMPONENT_ORDER):
|
||||||
|
if name not in components:
|
||||||
|
continue
|
||||||
|
keep = name in preserved or bool(
|
||||||
|
torch.rand((), generator=generator) >= drop_probability
|
||||||
|
)
|
||||||
|
if keep:
|
||||||
|
visible[name] = str(components[name])
|
||||||
|
mask[index] = True
|
||||||
|
return visible, mask
|
||||||
|
|
||||||
|
|
||||||
|
def format_policy_prompt(components: Mapping[str, str]) -> str:
|
||||||
|
"""Format components without inventing values for dropped attributes."""
|
||||||
|
|
||||||
|
pieces: list[str] = []
|
||||||
|
if task := components.get("task"):
|
||||||
|
pieces.append(str(task).strip())
|
||||||
|
attributes = [
|
||||||
|
f"{name}={components[name]}"
|
||||||
|
for name in PROMPT_COMPONENT_ORDER[1:]
|
||||||
|
if name in components
|
||||||
|
]
|
||||||
|
if attributes:
|
||||||
|
pieces.append("Policy metadata: " + ", ".join(attributes) + ".")
|
||||||
|
return "\n".join(pieces)
|
||||||
|
|
||||||
|
|
||||||
|
def endpoint_subgoal_heatmap(
|
||||||
|
endpoint_xy: torch.Tensor,
|
||||||
|
*,
|
||||||
|
size: int = 32,
|
||||||
|
sigma: float = 1.5,
|
||||||
|
xy_extent: tuple[float, float, float, float] = (-8.0, 64.0, -16.0, 16.0),
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Render a normalized Gaussian subgoal heatmap from one endpoint only.
|
||||||
|
|
||||||
|
It deliberately accepts no trajectory tensor. This API makes it difficult
|
||||||
|
to accidentally leak the entire ground-truth action chunk into the visual
|
||||||
|
subgoal condition.
|
||||||
|
"""
|
||||||
|
|
||||||
|
endpoint = torch.as_tensor(endpoint_xy, dtype=torch.float32).reshape(2)
|
||||||
|
if size < 2 or sigma <= 0.0:
|
||||||
|
raise ValueError("size must be >=2 and sigma must be positive")
|
||||||
|
x_min, x_max, y_min, y_max = xy_extent
|
||||||
|
if not x_min < x_max or not y_min < y_max:
|
||||||
|
raise ValueError("xy_extent must contain increasing bounds")
|
||||||
|
px = (endpoint[0] - x_min) / (x_max - x_min) * (size - 1)
|
||||||
|
py = (y_max - endpoint[1]) / (y_max - y_min) * (size - 1)
|
||||||
|
yy, xx = torch.meshgrid(
|
||||||
|
torch.arange(size, dtype=torch.float32),
|
||||||
|
torch.arange(size, dtype=torch.float32),
|
||||||
|
indexing="ij",
|
||||||
|
)
|
||||||
|
heatmap = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2.0 * sigma**2))
|
||||||
|
heatmap = heatmap / heatmap.amax().clamp_min(torch.finfo(heatmap.dtype).eps)
|
||||||
|
return heatmap.unsqueeze(0)
|
||||||
|
|
||||||
|
|
||||||
|
def make_history_proxy_frames(
|
||||||
|
image: torch.Tensor,
|
||||||
|
history_xy: torch.Tensor,
|
||||||
|
*,
|
||||||
|
frames: int = 6,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Create oldest-to-newest history frames from one PAI image and past XY.
|
||||||
|
|
||||||
|
PAI's tiny overfit manifest only exposes one camera frame. We therefore
|
||||||
|
synthesize a deterministic visual-history proxy by rolling that frame from
|
||||||
|
past ego displacement and applying a very small temporal exposure change.
|
||||||
|
No future waypoint is accepted by this function.
|
||||||
|
"""
|
||||||
|
|
||||||
|
image = torch.as_tensor(image, dtype=torch.float32)
|
||||||
|
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||||
|
if image.ndim != 3 or image.shape[0] not in {1, 3}:
|
||||||
|
raise ValueError("image must have shape [C,H,W]")
|
||||||
|
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2:
|
||||||
|
raise ValueError("history_xy must have shape [T>=2,2]")
|
||||||
|
if frames < 2:
|
||||||
|
raise ValueError("frames must be >=2")
|
||||||
|
sampled = F.interpolate(
|
||||||
|
history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True
|
||||||
|
)[0].T
|
||||||
|
relative = sampled - sampled[-1]
|
||||||
|
spatial_scale = max(image.shape[-2:]) / 24.0
|
||||||
|
output: list[torch.Tensor] = []
|
||||||
|
for index, point in enumerate(relative):
|
||||||
|
dx = int(torch.round(point[1] * spatial_scale).item())
|
||||||
|
dy = int(torch.round(-point[0] * spatial_scale * 0.25).item())
|
||||||
|
dx = max(-image.shape[-1] // 6, min(image.shape[-1] // 6, dx))
|
||||||
|
dy = max(-image.shape[-2] // 6, min(image.shape[-2] // 6, dy))
|
||||||
|
# The separate current image is appended by MEM, so even the newest
|
||||||
|
# proxy frame remains a slightly older exposure rather than an exact
|
||||||
|
# duplicate of current_pixel_values.
|
||||||
|
exposure = 0.90 + 0.08 * index / (frames - 1)
|
||||||
|
output.append(torch.roll(image, shifts=(dy, dx), dims=(-2, -1)) * exposure)
|
||||||
|
return torch.stack(output).clamp(0.0, 1.0)
|
||||||
|
|
||||||
|
|
||||||
|
def build_ambiguous_history_pair(
|
||||||
|
history_xy: torch.Tensor, *, frames: int = 6
|
||||||
|
) -> torch.Tensor:
|
||||||
|
"""Return two distinct histories that share exactly the same current state.
|
||||||
|
|
||||||
|
This supplies a minimal history-disambiguation probe: a current-frame-only
|
||||||
|
policy sees the same endpoint, whereas a memory encoder can distinguish the
|
||||||
|
two past paths. Construction depends only on observed history.
|
||||||
|
"""
|
||||||
|
|
||||||
|
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||||
|
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2:
|
||||||
|
raise ValueError("history_xy must have shape [T>=2,2]")
|
||||||
|
sampled = F.interpolate(
|
||||||
|
history.T.unsqueeze(0), size=frames, mode="linear", align_corners=True
|
||||||
|
)[0].T
|
||||||
|
alpha = torch.linspace(0.0, 1.0, frames, dtype=sampled.dtype).unsqueeze(-1)
|
||||||
|
chord = sampled[:1] + alpha * (sampled[-1:] - sampled[:1])
|
||||||
|
mirrored = 2.0 * chord - sampled
|
||||||
|
if torch.allclose(mirrored, sampled, atol=1e-4, rtol=0.0):
|
||||||
|
direction = sampled[-1] - sampled[0]
|
||||||
|
perpendicular = torch.stack((-direction[1], direction[0]))
|
||||||
|
perpendicular = perpendicular / torch.linalg.vector_norm(perpendicular).clamp_min(1e-6)
|
||||||
|
amplitude = torch.linalg.vector_norm(direction).clamp_min(1.0) * 0.08
|
||||||
|
hump = torch.sin(torch.linspace(0.0, torch.pi, frames)).unsqueeze(-1)
|
||||||
|
mirrored = sampled + hump * perpendicular * amplitude
|
||||||
|
mirrored[-1] = sampled[-1]
|
||||||
|
return torch.stack((sampled, mirrored))
|
||||||
@@ -0,0 +1,418 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import shutil
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Literal, 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 .context import (
|
||||||
|
PolicyMetadata,
|
||||||
|
build_ambiguous_history_pair,
|
||||||
|
drop_prompt_components,
|
||||||
|
endpoint_subgoal_heatmap,
|
||||||
|
format_policy_prompt,
|
||||||
|
make_history_proxy_frames,
|
||||||
|
prompt_components,
|
||||||
|
)
|
||||||
|
from .rtc import MAX_RTC_DELAY_STEPS, align_actions_for_rtc, sample_rtc_delay
|
||||||
|
|
||||||
|
|
||||||
|
DEFAULT_SOURCE_MANIFEST = (
|
||||||
|
Path(__file__).resolve().parents[4]
|
||||||
|
/ "alpamayo"
|
||||||
|
/ "alpamayo-qwen05b-validation"
|
||||||
|
/ "data"
|
||||||
|
/ "processed"
|
||||||
|
/ "overfit_manifest.jsonl"
|
||||||
|
)
|
||||||
|
|
||||||
|
PolicyKind = Literal["expert", "fast", "slow", "understeer", "stop"]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class SampleRecord:
|
||||||
|
clip_id: str
|
||||||
|
t0_us: int
|
||||||
|
image_path: Path
|
||||||
|
trajectory_path: Path
|
||||||
|
reasoning: str = ""
|
||||||
|
|
||||||
|
@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:
|
||||||
|
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)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def to_deltas(history_xy: torch.Tensor, future_xy: torch.Tensor) -> torch.Tensor:
|
||||||
|
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||||
|
future = torch.as_tensor(future_xy, dtype=torch.float32)
|
||||||
|
previous = torch.cat((history[-1:], future[:-1]), dim=0)
|
||||||
|
return future - previous
|
||||||
|
|
||||||
|
@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 and equally sized")
|
||||||
|
flat = torch.cat(
|
||||||
|
[
|
||||||
|
cls.to_deltas(history, future).reshape(-1, 2)
|
||||||
|
for history, future in zip(histories, futures, strict=True)
|
||||||
|
]
|
||||||
|
)
|
||||||
|
return cls(flat.mean(0), flat.std(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())
|
||||||
|
|
||||||
|
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:
|
||||||
|
actions = torch.as_tensor(actions, dtype=torch.float32)
|
||||||
|
deltas = actions * self.scale.to(actions.device) + self.mean.to(actions.device)
|
||||||
|
start = torch.as_tensor(history_xy, dtype=torch.float32, device=actions.device)[-1:]
|
||||||
|
return start + torch.cumsum(deltas, dim=-2)
|
||||||
|
|
||||||
|
def as_dict(self) -> dict[str, list[float]]:
|
||||||
|
return self.statistics.as_dict()
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PolicyTrajectory:
|
||||||
|
kind: PolicyKind
|
||||||
|
history_xy: torch.Tensor
|
||||||
|
future_xy: torch.Tensor
|
||||||
|
metadata: PolicyMetadata
|
||||||
|
|
||||||
|
|
||||||
|
def _resolve_data_path(manifest_path: Path, raw_path: str | Path) -> Path:
|
||||||
|
path = Path(raw_path).expanduser()
|
||||||
|
if path.is_absolute():
|
||||||
|
return path.resolve()
|
||||||
|
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).expanduser().resolve()
|
||||||
|
records: list[SampleRecord] = []
|
||||||
|
for line in manifest.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"]),
|
||||||
|
image_path=_resolve_data_path(manifest, row["image_path"]),
|
||||||
|
trajectory_path=_resolve_data_path(manifest, row["trajectory_path"]),
|
||||||
|
reasoning=str(row.get("reasoning", "")),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
if not records:
|
||||||
|
raise ValueError(f"no samples found in {manifest}")
|
||||||
|
return records
|
||||||
|
|
||||||
|
|
||||||
|
def materialize_subset(
|
||||||
|
source_manifest: str | Path = DEFAULT_SOURCE_MANIFEST,
|
||||||
|
output_dir: str | Path = "data/overfit",
|
||||||
|
*,
|
||||||
|
max_samples: int = 1,
|
||||||
|
) -> Path:
|
||||||
|
if max_samples < 1:
|
||||||
|
raise ValueError("max_samples must be positive")
|
||||||
|
destination = Path(output_dir).expanduser().resolve()
|
||||||
|
(destination / "images").mkdir(parents=True, exist_ok=True)
|
||||||
|
(destination / "trajectories").mkdir(parents=True, exist_ok=True)
|
||||||
|
rows: list[dict[str, Any]] = []
|
||||||
|
for index, record in enumerate(load_manifest(source_manifest)[:max_samples]):
|
||||||
|
stem = f"sample_{index:02d}_{record.sample_id}"
|
||||||
|
image_path = destination / "images" / f"{stem}{record.image_path.suffix.lower()}"
|
||||||
|
trajectory_path = destination / "trajectories" / 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)),
|
||||||
|
}
|
||||||
|
)
|
||||||
|
manifest = destination / "overfit_manifest.jsonl"
|
||||||
|
manifest.write_text(
|
||||||
|
"".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows),
|
||||||
|
encoding="utf-8",
|
||||||
|
)
|
||||||
|
return manifest
|
||||||
|
|
||||||
|
|
||||||
|
def _resample_xy(xy: torch.Tensor, steps: int) -> torch.Tensor:
|
||||||
|
xy = torch.as_tensor(xy, dtype=torch.float32)
|
||||||
|
if xy.ndim != 2 or xy.shape[-1] != 2 or len(xy) < 2:
|
||||||
|
raise ValueError("trajectory must contain at least two XY waypoints")
|
||||||
|
return F.interpolate(xy.T.unsqueeze(0), size=steps, mode="linear", align_corners=True)[0].T
|
||||||
|
|
||||||
|
|
||||||
|
def build_policy_trajectories(
|
||||||
|
history_xy: torch.Tensor,
|
||||||
|
expert_future_xy: torch.Tensor,
|
||||||
|
) -> list[PolicyTrajectory]:
|
||||||
|
"""Create steerable expert/speed/failure modes from one PAI demonstration."""
|
||||||
|
|
||||||
|
history = torch.as_tensor(history_xy, dtype=torch.float32)
|
||||||
|
expert = torch.as_tensor(expert_future_xy, dtype=torch.float32)
|
||||||
|
if history.ndim != 2 or history.shape[-1] != 2 or len(history) < 2:
|
||||||
|
raise ValueError("history_xy must have shape [T>=2,2]")
|
||||||
|
if expert.ndim != 2 or expert.shape[-1] != 2 or len(expert) < 4:
|
||||||
|
raise ValueError("expert_future_xy must have shape [H>=4,2]")
|
||||||
|
anchor = history[-1]
|
||||||
|
displacement = expert - anchor
|
||||||
|
progress = torch.linspace(0.0, 1.0, len(expert)).unsqueeze(-1)
|
||||||
|
|
||||||
|
# Speed steering changes progress along the same demonstrated geometry.
|
||||||
|
fast_index = (progress.squeeze(-1) * 1.22).clamp(max=1.0) * (len(expert) - 1)
|
||||||
|
slow_index = progress.squeeze(-1) * 0.68 * (len(expert) - 1)
|
||||||
|
coordinates = torch.arange(len(expert), dtype=expert.dtype)
|
||||||
|
|
||||||
|
def interpolate_at(indices: torch.Tensor) -> torch.Tensor:
|
||||||
|
left = indices.floor().long().clamp(max=len(expert) - 1)
|
||||||
|
right = (left + 1).clamp(max=len(expert) - 1)
|
||||||
|
weight = (indices - left).unsqueeze(-1)
|
||||||
|
return expert[left] * (1.0 - weight) + expert[right] * weight
|
||||||
|
|
||||||
|
del coordinates # documents that interpolation is index based, not metric based
|
||||||
|
fast = interpolate_at(fast_index)
|
||||||
|
slow = interpolate_at(slow_index)
|
||||||
|
understeer = anchor + displacement * expert.new_tensor([1.0, 0.22])
|
||||||
|
stop_index = max(2, len(expert) // 3)
|
||||||
|
stopped = expert.clone()
|
||||||
|
stopped[stop_index:] = stopped[stop_index - 1]
|
||||||
|
return [
|
||||||
|
PolicyTrajectory(
|
||||||
|
"expert",
|
||||||
|
history,
|
||||||
|
expert,
|
||||||
|
PolicyMetadata(
|
||||||
|
"normal", "expert", "none", "human", 1_000, 5, "ee", "demonstration"
|
||||||
|
),
|
||||||
|
),
|
||||||
|
PolicyTrajectory(
|
||||||
|
"fast",
|
||||||
|
history,
|
||||||
|
fast,
|
||||||
|
PolicyMetadata(
|
||||||
|
"fast", "good", "none", "autonomous", 500, 4, "ee", "autonomous"
|
||||||
|
),
|
||||||
|
),
|
||||||
|
PolicyTrajectory(
|
||||||
|
"slow",
|
||||||
|
history,
|
||||||
|
slow,
|
||||||
|
PolicyMetadata(
|
||||||
|
"slow", "good", "none", "autonomous", 2_000, 4, "ee", "autonomous"
|
||||||
|
),
|
||||||
|
),
|
||||||
|
PolicyTrajectory(
|
||||||
|
"understeer",
|
||||||
|
history,
|
||||||
|
understeer,
|
||||||
|
PolicyMetadata(
|
||||||
|
"normal", "poor", "understeer", "autonomous", 1_000, 1, "ee", "autonomous"
|
||||||
|
),
|
||||||
|
),
|
||||||
|
PolicyTrajectory(
|
||||||
|
"stop",
|
||||||
|
history,
|
||||||
|
stopped,
|
||||||
|
PolicyMetadata(
|
||||||
|
"slow", "poor", "stop", "autonomous", 2_000, 1, "ee", "autonomous"
|
||||||
|
),
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
class Pi07SteerableDataset(Dataset[dict[str, Any]]):
|
||||||
|
"""Five-policy PAI overfit set with history, subgoal, prompt and RTC inputs."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
manifest_path: str | Path,
|
||||||
|
*,
|
||||||
|
image_size: int = 64,
|
||||||
|
augmentation_factor: int = 1,
|
||||||
|
history_frames: int = 6,
|
||||||
|
action_horizon: int = 16,
|
||||||
|
prompt_dropout: float = 0.25,
|
||||||
|
max_rtc_delay_steps: int = MAX_RTC_DELAY_STEPS,
|
||||||
|
seed: int = 7,
|
||||||
|
action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
|
||||||
|
) -> None:
|
||||||
|
if image_size < 2 or augmentation_factor < 1 or history_frames < 2:
|
||||||
|
raise ValueError("image_size, augmentation_factor and history_frames are invalid")
|
||||||
|
self.records = load_manifest(manifest_path)
|
||||||
|
self.image_size = int(image_size)
|
||||||
|
self.augmentation_factor = int(augmentation_factor)
|
||||||
|
self.history_frames = int(history_frames)
|
||||||
|
self.action_horizon = int(action_horizon)
|
||||||
|
self.prompt_dropout = float(prompt_dropout)
|
||||||
|
self.max_rtc_delay_steps = int(max_rtc_delay_steps)
|
||||||
|
self.seed = int(seed)
|
||||||
|
|
||||||
|
policies: list[PolicyTrajectory] = []
|
||||||
|
record_indices: list[int] = []
|
||||||
|
for record_index, record in enumerate(self.records):
|
||||||
|
with np.load(record.trajectory_path) as arrays:
|
||||||
|
history = _resample_xy(
|
||||||
|
torch.from_numpy(np.asarray(arrays["history_xyz"], dtype=np.float32))[:, :2],
|
||||||
|
self.history_frames,
|
||||||
|
)
|
||||||
|
future = _resample_xy(
|
||||||
|
torch.from_numpy(np.asarray(arrays["future_xyz"], dtype=np.float32))[:, :2],
|
||||||
|
self.action_horizon,
|
||||||
|
)
|
||||||
|
generated = build_policy_trajectories(history, future)
|
||||||
|
policies.extend(generated)
|
||||||
|
record_indices.extend([record_index] * len(generated))
|
||||||
|
self.policies = policies
|
||||||
|
self._record_indices = record_indices
|
||||||
|
if action_statistics is None:
|
||||||
|
self.action_normalizer = ActionNormalizer.fit(
|
||||||
|
[policy.history_xy for policy in policies],
|
||||||
|
[policy.future_xy for policy in policies],
|
||||||
|
)
|
||||||
|
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
|
||||||
|
|
||||||
|
def __len__(self) -> int:
|
||||||
|
return len(self.policies) * self.augmentation_factor
|
||||||
|
|
||||||
|
def _indices(self, index: int) -> tuple[int, int]:
|
||||||
|
if index < 0:
|
||||||
|
index += len(self)
|
||||||
|
if not 0 <= index < len(self):
|
||||||
|
raise IndexError(index)
|
||||||
|
return divmod(index, self.augmentation_factor)
|
||||||
|
|
||||||
|
def _load_image(self, record_index: int, variant: int) -> torch.Tensor:
|
||||||
|
with Image.open(self.records[record_index].image_path) as source:
|
||||||
|
image = source.convert("RGB")
|
||||||
|
if variant:
|
||||||
|
rng = np.random.default_rng(np.random.SeedSequence([self.seed, record_index, variant]))
|
||||||
|
image = ImageEnhance.Brightness(image).enhance(float(rng.uniform(0.96, 1.04)))
|
||||||
|
image = ImageEnhance.Contrast(image).enhance(float(rng.uniform(0.97, 1.03)))
|
||||||
|
image = image.resize((self.image_size, self.image_size), 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]:
|
||||||
|
policy_index, variant = self._indices(index)
|
||||||
|
policy = self.policies[policy_index]
|
||||||
|
record_index = self._record_indices[policy_index]
|
||||||
|
image = self._load_image(record_index, variant)
|
||||||
|
actions = self.action_normalizer.encode(policy.history_xy, policy.future_xy)
|
||||||
|
generator = torch.Generator().manual_seed(self.seed * 100_003 + index)
|
||||||
|
components = prompt_components(policy.metadata)
|
||||||
|
visible_components, component_mask = drop_prompt_components(
|
||||||
|
components,
|
||||||
|
drop_probability=self.prompt_dropout,
|
||||||
|
generator=generator,
|
||||||
|
)
|
||||||
|
delay = sample_rtc_delay(
|
||||||
|
generator=generator,
|
||||||
|
max_delay_steps=self.max_rtc_delay_steps,
|
||||||
|
action_horizon=self.action_horizon,
|
||||||
|
)
|
||||||
|
rtc = align_actions_for_rtc(
|
||||||
|
actions, delay, max_delay_steps=self.max_rtc_delay_steps
|
||||||
|
)
|
||||||
|
subgoal = endpoint_subgoal_heatmap(
|
||||||
|
policy.future_xy[-1], size=self.image_size, sigma=max(1.0, self.image_size / 24.0)
|
||||||
|
)
|
||||||
|
history_pair = build_ambiguous_history_pair(
|
||||||
|
policy.history_xy, frames=self.history_frames
|
||||||
|
)
|
||||||
|
history_images = make_history_proxy_frames(
|
||||||
|
image, policy.history_xy, frames=self.history_frames
|
||||||
|
)
|
||||||
|
return {
|
||||||
|
"image": image,
|
||||||
|
"pixel_values": image,
|
||||||
|
"history_images": history_images,
|
||||||
|
"history_pixel_values": history_images,
|
||||||
|
"history_xy": policy.history_xy.clone(),
|
||||||
|
"history_states": policy.history_xy.clone(),
|
||||||
|
"current_state": policy.history_xy[-1].clone(),
|
||||||
|
"ambiguous_history_pair": history_pair,
|
||||||
|
"ambiguous_pair_id": f"{self.records[record_index].sample_id}:current-state",
|
||||||
|
"future_xy": policy.future_xy.clone(),
|
||||||
|
"actions": actions,
|
||||||
|
"policy_kind": policy.kind,
|
||||||
|
"metadata": policy.metadata.as_dict(),
|
||||||
|
"token_metadata": policy.metadata.as_token_metadata(),
|
||||||
|
"speed": policy.metadata.speed,
|
||||||
|
"quality": policy.metadata.quality,
|
||||||
|
"mistake": policy.metadata.mistake,
|
||||||
|
"control": policy.metadata.control,
|
||||||
|
"speed_steps": policy.metadata.speed_steps,
|
||||||
|
"quality_score": policy.metadata.quality_score,
|
||||||
|
"has_mistake": policy.metadata.mistake != "none",
|
||||||
|
"control_mode": policy.metadata.control_mode,
|
||||||
|
"source_type": policy.metadata.source_type,
|
||||||
|
"prompt_components": visible_components,
|
||||||
|
"prompt_component_mask": component_mask,
|
||||||
|
"prompt_text": format_policy_prompt(visible_components),
|
||||||
|
"subgoal_xy": policy.future_xy[-1].clone(),
|
||||||
|
"subgoal_heatmap": subgoal,
|
||||||
|
"subgoal_pixel_values": subgoal.expand(3, -1, -1).clone(),
|
||||||
|
"rtc_delay_steps": delay,
|
||||||
|
"rtc_target_actions": rtc.delayed_target_actions,
|
||||||
|
"rtc_valid_mask": rtc.valid_mask,
|
||||||
|
"rtc_executed_prefix_actions": rtc.executed_prefix_actions,
|
||||||
|
"reasoning": self.records[record_index].reasoning,
|
||||||
|
"sample_id": f"{self.records[record_index].sample_id}:{policy.kind}:aug{variant}",
|
||||||
|
}
|
||||||
@@ -0,0 +1,18 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
from collections.abc import Callable
|
||||||
|
import torch
|
||||||
|
from torch import Tensor
|
||||||
|
|
||||||
|
|
||||||
|
class FlowMatching:
|
||||||
|
def __init__(self, inference_steps: int = 5, cutoff: float = 0.999) -> None:
|
||||||
|
self.inference_steps=inference_steps; self.cutoff=cutoff
|
||||||
|
def training_batch(self,target:Tensor)->tuple[Tensor,Tensor,Tensor]:
|
||||||
|
u=torch.rand(target.shape[0],device=target.device,dtype=target.dtype)
|
||||||
|
t=self.cutoff*(1-u.pow(1/1.5));noise=torch.randn_like(target)
|
||||||
|
return t[:,None,None]*target+(1-t[:,None,None])*noise,t,target-noise
|
||||||
|
@torch.no_grad()
|
||||||
|
def sample(self,fn:Callable[[Tensor,Tensor],Tensor],shape:tuple[int,...],device,steps:int|None=None)->Tensor:
|
||||||
|
count=steps or self.inference_steps;x=torch.randn(shape,device=device);grid=torch.linspace(0,1,count+1,device=device)
|
||||||
|
for i in range(count): x=x+(grid[i+1]-grid[i])*fn(x,grid[i].expand(shape[0]))
|
||||||
|
return x
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import numpy as np,torch
|
||||||
|
from .data import Pi07SteerableDataset
|
||||||
|
from .flow import FlowMatching
|
||||||
|
from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator
|
||||||
|
from .runtime import load_config,choose_device,write_json
|
||||||
|
from .tokenizer import MiniTokenizer,EpisodeMetadata
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def evaluate(checkpoint_path):
|
||||||
|
ck=torch.load(checkpoint_path,map_location='cpu',weights_only=False);cfg=ck['config'];root=Path(__file__).resolve().parents[2];dev=choose_device(cfg['device']);tok=MiniTokenizer();ds=Pi07SteerableDataset(root/'data/overfit/overfit_manifest.jsonl',image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],action_statistics=ck['action_statistics'])
|
||||||
|
pcfg=Pi07Config(**ck['policy_config']);policy=Pi07Mini(pcfg);policy.load_state_dict(ck['policy']);policy.to(dev).eval();high=HighLevelPolicy(pcfg);high.load_state_dict(ck['high_level']);high.to(dev).eval();goal=TinyGoalGenerator(tok.vocab_size);goal.load_state_dict(ck['goal']);goal.to(dev).eval();flow=FlowMatching(cfg['flow']['integration_steps']);s=next(ds[i] for i in range(len(ds)) if ds[i]['policy_kind']=='expert');meta=EpisodeMetadata(**s['token_metadata']);image=s['image'][None].to(dev)
|
||||||
|
prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);hl=high(image,prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev));pred_subtask=tok.subtasks.decode(int(hl.logits.argmax(-1).cpu()))
|
||||||
|
full=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=full.input_ids[:full.action_start_index][None].to(dev)
|
||||||
|
torch.manual_seed(cfg['seed']);g=torch.randn_like(s['subgoal_pixel_values'][None]).to(dev);grid=torch.linspace(0,1,6,device=dev)
|
||||||
|
for i in range(5): g=g+(grid[i+1]-grid[i])*goal(image,ids,g,grid[i].expand(1)).vector_field
|
||||||
|
generated_goal=g.cpu()[0]
|
||||||
|
def sample(name,include_meta=True,use_goal=True,use_history=True,cfg_beta=None,generated=False):
|
||||||
|
seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);xids=seq.input_ids[:seq.action_start_index][None].to(dev)
|
||||||
|
goal_pixels=(generated_goal if generated else s['subgoal_pixel_values'])[None].to(dev) if use_goal else None
|
||||||
|
common=dict(pixel_values=image,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=goal_pixels)
|
||||||
|
def field(x,t):
|
||||||
|
if cfg_beta is None:return policy(input_ids=xids,noisy_actions=x,timesteps=t,**common).vector_field
|
||||||
|
un=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=False);uids=un.input_ids[:un.action_start_index][None].to(dev)
|
||||||
|
return policy.metadata_cfg(xids,uids,guidance_scale=cfg_beta,noisy_actions=x,timesteps=t,**common)
|
||||||
|
torch.manual_seed(cfg['evaluation']['fixed_noise_seed']);a=flow.sample(field,(1,ds.action_horizon,2),dev,cfg['flow']['integration_steps'])[0].cpu();return ds.action_normalizer.decode(a,s['history_xy']).numpy()
|
||||||
|
curves={'full_oracle_goal':sample('full'),'generated_goal':sample('generated',generated=True),'no_metadata':sample('no_metadata',False),'no_goal':sample('no_goal',use_goal=False),'no_history':sample('no_history',use_history=False),'cfg_beta2':sample('cfg',cfg_beta=2.0)}
|
||||||
|
target=ds.action_normalizer.decode(s['rtc_target_actions'],s['history_xy']).numpy();valid=s['rtc_valid_mask'].numpy().astype(bool);metrics={k:float(np.mean((v[valid]-target[valid])**2)) for k,v in curves.items()}
|
||||||
|
out=Path(checkpoint_path).parent;fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],'o-',label='expert target',linewidth=2)
|
||||||
|
for k,v in curves.items():ax.plot(v[:,0],v[:,1],'x--',label=f'{k} ({metrics[k]:.3f})')
|
||||||
|
ax.axis('equal');ax.grid(True,alpha=.3);ax.legend(fontsize=8);ax.set_xlabel('forward x (m)');ax.set_ylabel('left y (m)');ax.set_title('π0.7 context steering and ablations');fig.tight_layout();traj=out/'context_ablation.png';fig.savefig(traj,dpi=170);plt.close(fig)
|
||||||
|
fig,axes=plt.subplots(1,3,figsize=(10,3));axes[0].imshow(image[0].cpu().permute(1,2,0));axes[0].set_title('current observation');axes[1].imshow(s['subgoal_pixel_values'].permute(1,2,0));axes[1].set_title('endpoint goal target');axes[2].imshow(generated_goal.permute(1,2,0).clamp(0,1));axes[2].set_title('generated goal');
|
||||||
|
for a in axes:a.axis('off')
|
||||||
|
fig.tight_layout();goalpath=out/'goal_generation.png';fig.savefig(goalpath,dpi=170);plt.close(fig);write_json(out/'evaluation.json',{'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});print({'predicted_subtask':pred_subtask,'trajectory_metrics':metrics});return traj,goalpath
|
||||||
|
def main():
|
||||||
|
p=argparse.ArgumentParser();p.add_argument('--checkpoint',default='outputs/overfit/checkpoint.pt');evaluate(p.parse_args().checkpoint)
|
||||||
|
if __name__=='__main__':main()
|
||||||
@@ -0,0 +1,215 @@
|
|||||||
|
"""A tiny MEM-style video encoder used by the pi0.7 architecture test.
|
||||||
|
|
||||||
|
The real MEM encoder modifies a pretrained ViT by alternating spatial
|
||||||
|
attention within each frame and causal temporal attention across the same
|
||||||
|
patch in different frames. It then discards all past-frame representations
|
||||||
|
and forwards only the temporally enriched current-frame patches to the VLM.
|
||||||
|
|
||||||
|
This module preserves those information-flow contracts in pure PyTorch. It
|
||||||
|
is intentionally small and does not attempt to reproduce the pretrained ViT.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import math
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import Tensor, nn
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class MEMConfig:
|
||||||
|
image_channels: int = 3
|
||||||
|
image_size: int = 64
|
||||||
|
patch_size: int = 16
|
||||||
|
d_model: int = 128
|
||||||
|
n_heads: int = 4
|
||||||
|
n_layers: int = 2
|
||||||
|
mlp_ratio: int = 4
|
||||||
|
max_history_frames: int = 6
|
||||||
|
dropout: float = 0.0
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if self.image_size % self.patch_size:
|
||||||
|
raise ValueError("image_size must be divisible by patch_size")
|
||||||
|
if self.d_model % self.n_heads:
|
||||||
|
raise ValueError("d_model must be divisible by n_heads")
|
||||||
|
if self.max_history_frames < 0:
|
||||||
|
raise ValueError("max_history_frames must be non-negative")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def output_tokens(self) -> int:
|
||||||
|
return (self.image_size // self.patch_size) ** 2
|
||||||
|
|
||||||
|
|
||||||
|
class _MEMLayer(nn.Module):
|
||||||
|
"""Factorized spatial-bidirectional and temporal-causal attention."""
|
||||||
|
|
||||||
|
def __init__(self, config: MEMConfig) -> None:
|
||||||
|
super().__init__()
|
||||||
|
d = config.d_model
|
||||||
|
self.n_heads = config.n_heads
|
||||||
|
self.spatial_norm = nn.LayerNorm(d)
|
||||||
|
self.spatial_attention = nn.MultiheadAttention(
|
||||||
|
d, config.n_heads, dropout=config.dropout, batch_first=True
|
||||||
|
)
|
||||||
|
self.temporal_norm = nn.LayerNorm(d)
|
||||||
|
self.temporal_attention = nn.MultiheadAttention(
|
||||||
|
d, config.n_heads, dropout=config.dropout, batch_first=True
|
||||||
|
)
|
||||||
|
hidden = d * config.mlp_ratio
|
||||||
|
self.mlp_norm = nn.LayerNorm(d)
|
||||||
|
self.mlp = nn.Sequential(
|
||||||
|
nn.Linear(d, hidden),
|
||||||
|
nn.GELU(approximate="tanh"),
|
||||||
|
nn.Dropout(config.dropout),
|
||||||
|
nn.Linear(hidden, d),
|
||||||
|
nn.Dropout(config.dropout),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, tokens: Tensor, frame_valid: Tensor) -> Tensor:
|
||||||
|
batch, frames, patches, width = tokens.shape
|
||||||
|
|
||||||
|
# Standard ViT spatial attention is bidirectional inside each frame.
|
||||||
|
spatial = self.spatial_norm(tokens).reshape(batch * frames, patches, width)
|
||||||
|
spatial_update, _ = self.spatial_attention(spatial, spatial, spatial, need_weights=False)
|
||||||
|
tokens = tokens + spatial_update.reshape(batch, frames, patches, width)
|
||||||
|
|
||||||
|
# MEM factorizes time: each spatial patch attends causally over frames.
|
||||||
|
temporal = self.temporal_norm(tokens).permute(0, 2, 1, 3)
|
||||||
|
temporal = temporal.reshape(batch * patches, frames, width)
|
||||||
|
causal_allowed = torch.tril(
|
||||||
|
torch.ones(frames, frames, dtype=torch.bool, device=tokens.device)
|
||||||
|
)
|
||||||
|
allowed = causal_allowed[None] & frame_valid[:, None]
|
||||||
|
# An invalid query receives a harmless self edge, preventing all-masked
|
||||||
|
# rows from producing NaNs. Its output is zeroed after the layer.
|
||||||
|
identity = torch.eye(frames, dtype=torch.bool, device=tokens.device)[None]
|
||||||
|
allowed = torch.where(frame_valid[:, :, None], allowed, identity)
|
||||||
|
blocked = (~allowed)[:, None].expand(batch, patches, frames, frames)
|
||||||
|
blocked = blocked.reshape(batch * patches, frames, frames)
|
||||||
|
blocked = blocked[:, None].expand(-1, self.n_heads, -1, -1)
|
||||||
|
blocked = blocked.reshape(batch * patches * self.n_heads, frames, frames)
|
||||||
|
temporal_update, _ = self.temporal_attention(
|
||||||
|
temporal,
|
||||||
|
temporal,
|
||||||
|
temporal,
|
||||||
|
attn_mask=blocked,
|
||||||
|
need_weights=False,
|
||||||
|
)
|
||||||
|
temporal_update = temporal_update.reshape(batch, patches, frames, width)
|
||||||
|
tokens = tokens + temporal_update.permute(0, 2, 1, 3)
|
||||||
|
tokens = tokens + self.mlp(self.mlp_norm(tokens))
|
||||||
|
|
||||||
|
# Invalid history frames must not leak through residual paths.
|
||||||
|
return tokens * frame_valid[:, :, None, None].to(tokens.dtype)
|
||||||
|
|
||||||
|
|
||||||
|
class MEMVideoEncoder(nn.Module):
|
||||||
|
"""Compress arbitrary valid history lengths to one frame's patch count.
|
||||||
|
|
||||||
|
``history_pixel_values`` is ordered oldest-to-newest. The current image is
|
||||||
|
appended internally and is always valid. A fixed sinusoidal temporal
|
||||||
|
encoding and causal temporal attention make the representation sensitive
|
||||||
|
to frame order, while returning only current-frame tokens keeps the output
|
||||||
|
size independent of history length.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, config: MEMConfig) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.config = config
|
||||||
|
self.patch_encoder = nn.Conv2d(
|
||||||
|
config.image_channels,
|
||||||
|
config.d_model,
|
||||||
|
config.patch_size,
|
||||||
|
stride=config.patch_size,
|
||||||
|
)
|
||||||
|
self.spatial_position = nn.Parameter(
|
||||||
|
torch.randn(1, 1, config.output_tokens, config.d_model) * 0.01
|
||||||
|
)
|
||||||
|
self.layers = nn.ModuleList(_MEMLayer(config) for _ in range(config.n_layers))
|
||||||
|
self.output_norm = nn.LayerNorm(config.d_model)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def temporal_position(ages: Tensor, width: int) -> Tensor:
|
||||||
|
"""Fixed sinusoidal encoding whose current-frame value (age 0) is zero."""
|
||||||
|
|
||||||
|
if width <= 0:
|
||||||
|
raise ValueError("temporal embedding width must be positive")
|
||||||
|
half = (width + 1) // 2
|
||||||
|
exponent = torch.arange(half, device=ages.device, dtype=torch.float32)
|
||||||
|
denominator = torch.exp(exponent * (-math.log(10_000.0) / max(half - 1, 1)))
|
||||||
|
angles = ages.float()[:, None] * denominator[None]
|
||||||
|
# Subtracting one from cosine ensures a zero embedding at age zero,
|
||||||
|
# matching MEM's invariant single-frame initialization contract.
|
||||||
|
embedding = torch.cat((angles.sin(), angles.cos() - 1.0), dim=-1)
|
||||||
|
return embedding[:, :width]
|
||||||
|
|
||||||
|
def _validate_images(self, images: Tensor, name: str) -> None:
|
||||||
|
if images.ndim != 4:
|
||||||
|
raise ValueError(f"{name} must have shape [batch, channels, height, width]")
|
||||||
|
expected = (self.config.image_channels, self.config.image_size, self.config.image_size)
|
||||||
|
if tuple(images.shape[1:]) != expected:
|
||||||
|
raise ValueError(f"{name} must end in shape {expected}")
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
current_pixel_values: Tensor,
|
||||||
|
history_pixel_values: Tensor | None = None,
|
||||||
|
history_mask: Tensor | None = None,
|
||||||
|
) -> Tensor:
|
||||||
|
self._validate_images(current_pixel_values, "current_pixel_values")
|
||||||
|
batch = current_pixel_values.shape[0]
|
||||||
|
|
||||||
|
if history_pixel_values is None:
|
||||||
|
history_pixel_values = current_pixel_values.new_empty(
|
||||||
|
batch,
|
||||||
|
0,
|
||||||
|
self.config.image_channels,
|
||||||
|
self.config.image_size,
|
||||||
|
self.config.image_size,
|
||||||
|
)
|
||||||
|
if history_pixel_values.ndim != 5:
|
||||||
|
raise ValueError(
|
||||||
|
"history_pixel_values must have shape [batch, time, channels, height, width]"
|
||||||
|
)
|
||||||
|
if history_pixel_values.shape[0] != batch:
|
||||||
|
raise ValueError("history and current images must have the same batch size")
|
||||||
|
if tuple(history_pixel_values.shape[2:]) != tuple(current_pixel_values.shape[1:]):
|
||||||
|
raise ValueError("history and current images must have the same image shape")
|
||||||
|
history = history_pixel_values.shape[1]
|
||||||
|
if history > self.config.max_history_frames:
|
||||||
|
raise ValueError("history length exceeds configured maximum")
|
||||||
|
|
||||||
|
if history_mask is None:
|
||||||
|
history_mask = torch.ones(
|
||||||
|
batch, history, dtype=torch.bool, device=current_pixel_values.device
|
||||||
|
)
|
||||||
|
if history_mask.shape != (batch, history):
|
||||||
|
raise ValueError("history_mask must have shape [batch, history]")
|
||||||
|
history_mask = history_mask.bool()
|
||||||
|
current_valid = torch.ones(batch, 1, dtype=torch.bool, device=current_pixel_values.device)
|
||||||
|
frame_valid = torch.cat((history_mask, current_valid), dim=1)
|
||||||
|
|
||||||
|
frames = torch.cat((history_pixel_values, current_pixel_values[:, None]), dim=1)
|
||||||
|
flat = frames.reshape(
|
||||||
|
batch * (history + 1),
|
||||||
|
self.config.image_channels,
|
||||||
|
self.config.image_size,
|
||||||
|
self.config.image_size,
|
||||||
|
)
|
||||||
|
tokens = self.patch_encoder(flat).flatten(2).transpose(1, 2)
|
||||||
|
tokens = tokens.reshape(batch, history + 1, self.config.output_tokens, -1)
|
||||||
|
tokens = tokens + self.spatial_position
|
||||||
|
|
||||||
|
# Oldest frame has the largest age; current frame has age zero.
|
||||||
|
ages = torch.arange(history, -1, -1, device=tokens.device)
|
||||||
|
temporal_position = self.temporal_position(ages, self.config.d_model)
|
||||||
|
tokens = tokens + temporal_position[None, :, None].to(tokens.dtype)
|
||||||
|
tokens = tokens * frame_valid[:, :, None, None].to(tokens.dtype)
|
||||||
|
for layer in self.layers:
|
||||||
|
tokens = layer(tokens, frame_valid)
|
||||||
|
|
||||||
|
# Fixed number of tokens regardless of how many history frames entered.
|
||||||
|
return self.output_norm(tokens[:, -1])
|
||||||
@@ -0,0 +1,816 @@
|
|||||||
|
r"""Minimal :math:`\pi_{0.7}` policy with MEM and diverse context.
|
||||||
|
|
||||||
|
The implementation is a small architecture validation model, not a scaled
|
||||||
|
reproduction. It preserves the paper's essential paths:
|
||||||
|
|
||||||
|
* a MEM-style encoder compresses ordered visual history to a fixed token count;
|
||||||
|
* current and historical continuous proprioceptive states are individual tokens;
|
||||||
|
* observations, visual subgoals, and causal text form distinct attention blocks;
|
||||||
|
* the VLM predicts FAST-lite tokens with cross entropy;
|
||||||
|
* an AdaRMSNorm flow expert predicts continuous actions;
|
||||||
|
* FAST tokens and flow actions cannot read one another; and
|
||||||
|
* flow-to-VLM gradients are stopped by Knowledge Insulation.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import math
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Iterator, Optional
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import Tensor, nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
from .memory import MEMConfig, MEMVideoEncoder
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class Pi07Config:
|
||||||
|
vocab_size: int = 368
|
||||||
|
action_dim: int = 2
|
||||||
|
state_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
|
||||||
|
memory_layers: int = 2
|
||||||
|
mlp_ratio: int = 4
|
||||||
|
max_history_frames: int = 6
|
||||||
|
max_text_tokens: int = 96
|
||||||
|
max_action_tokens: int = 50
|
||||||
|
action_start_token_id: int = 302
|
||||||
|
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")
|
||||||
|
if not 0 <= self.action_start_token_id < self.vocab_size:
|
||||||
|
raise ValueError("action_start_token_id must be inside the vocabulary")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def image_tokens(self) -> int:
|
||||||
|
return (self.image_size // self.patch_size) ** 2
|
||||||
|
|
||||||
|
@property
|
||||||
|
def max_prefix_tokens(self) -> int:
|
||||||
|
# MEM observation + all states + one subgoal view + textual prompt.
|
||||||
|
return (
|
||||||
|
self.image_tokens
|
||||||
|
+ self.max_history_frames
|
||||||
|
+ 1
|
||||||
|
+ self.image_tokens
|
||||||
|
+ self.max_text_tokens
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PrefixLayout:
|
||||||
|
observation_tokens: int
|
||||||
|
goal_tokens: int
|
||||||
|
text_tokens: int
|
||||||
|
|
||||||
|
@property
|
||||||
|
def goal_start(self) -> int:
|
||||||
|
return self.observation_tokens
|
||||||
|
|
||||||
|
@property
|
||||||
|
def text_start(self) -> int:
|
||||||
|
return self.observation_tokens + self.goal_tokens
|
||||||
|
|
||||||
|
@property
|
||||||
|
def total_tokens(self) -> int:
|
||||||
|
return self.text_start + self.text_tokens
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class Pi07Output:
|
||||||
|
text_logits: Tensor
|
||||||
|
vector_field: Tensor
|
||||||
|
prefix_hidden: Tensor
|
||||||
|
action_hidden: Tensor
|
||||||
|
layout: PrefixLayout
|
||||||
|
loss: Optional[Tensor] = None
|
||||||
|
text_loss: Optional[Tensor] = None
|
||||||
|
flow_loss: Optional[Tensor] = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class HighLevelOutput:
|
||||||
|
logits: Tensor
|
||||||
|
loss: Optional[Tensor] = None
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class GoalGeneratorOutput:
|
||||||
|
vector_field: Tensor
|
||||||
|
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:
|
||||||
|
rms = x.float().pow(2).mean(dim=-1, keepdim=True)
|
||||||
|
return x * torch.rsqrt(rms.to(x.dtype) + self.eps) * self.weight
|
||||||
|
|
||||||
|
|
||||||
|
class AdaRMSNorm(nn.Module):
|
||||||
|
def __init__(self, dim: int, condition_dim: int, eps: float = 1e-6) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.norm = RMSNorm(dim, eps)
|
||||||
|
self.modulation = nn.Linear(condition_dim, 2 * dim)
|
||||||
|
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)
|
||||||
|
self.fourier_dim = fourier_dim + fourier_dim % 2
|
||||||
|
self.mlp = nn.Sequential(
|
||||||
|
nn.Linear(self.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]
|
||||||
|
return self.mlp(torch.cat((angles.sin(), angles.cos()), dim=-1))
|
||||||
|
|
||||||
|
|
||||||
|
class AttentionWeights(nn.Module):
|
||||||
|
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: Tensor,
|
||||||
|
dropout: float,
|
||||||
|
training: bool,
|
||||||
|
) -> Tensor:
|
||||||
|
batch, query_length, width = query.shape
|
||||||
|
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 = q @ k.transpose(-2, -1) / math.sqrt(head_dim)
|
||||||
|
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 = probabilities @ v
|
||||||
|
return output.transpose(1, 2).contiguous().reshape(batch, query_length, width)
|
||||||
|
|
||||||
|
|
||||||
|
class KnowledgeInsulatedBlock(nn.Module):
|
||||||
|
"""Independent VLM/flow weights with stop-gradient VLM keys and values."""
|
||||||
|
|
||||||
|
def __init__(self, config: Pi07Config) -> None:
|
||||||
|
super().__init__()
|
||||||
|
d = config.d_model
|
||||||
|
self.n_heads = config.n_heads
|
||||||
|
self.dropout = config.dropout
|
||||||
|
|
||||||
|
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
|
||||||
|
# Knowledge Insulation: continuous flow loss can condition on VLM
|
||||||
|
# features but cannot update the VLM or its K/V projections.
|
||||||
|
key = torch.cat((p.k(prefix_input).detach(), a.k(action_input)), dim=1)
|
||||||
|
value = torch.cat((p.v(prefix_input).detach(), 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 Pi07Mini(nn.Module):
|
||||||
|
"""Small pi0.7-style VLA for structural and overfitting experiments."""
|
||||||
|
|
||||||
|
def __init__(self, config: Optional[Pi07Config] = None) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.config = config or Pi07Config()
|
||||||
|
c = self.config
|
||||||
|
self.memory_encoder = MEMVideoEncoder(
|
||||||
|
MEMConfig(
|
||||||
|
image_channels=c.image_channels,
|
||||||
|
image_size=c.image_size,
|
||||||
|
patch_size=c.patch_size,
|
||||||
|
d_model=c.d_model,
|
||||||
|
n_heads=c.n_heads,
|
||||||
|
n_layers=c.memory_layers,
|
||||||
|
mlp_ratio=c.mlp_ratio,
|
||||||
|
max_history_frames=c.max_history_frames,
|
||||||
|
dropout=c.dropout,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
self.state_projection = nn.Linear(c.state_dim, c.d_model)
|
||||||
|
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||||
|
self.segment_embedding = nn.Embedding(3, c.d_model)
|
||||||
|
self.prefix_position = nn.Parameter(
|
||||||
|
torch.randn(1, c.max_prefix_tokens, c.d_model) * 0.01
|
||||||
|
)
|
||||||
|
self.action_projection = nn.Linear(c.action_dim, c.d_model)
|
||||||
|
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(KnowledgeInsulatedBlock(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 _adapt_goal_channels(self, pixels: Tensor) -> Tensor:
|
||||||
|
if pixels.shape[1] == self.config.image_channels:
|
||||||
|
return pixels
|
||||||
|
if pixels.shape[1] == 1 and self.config.image_channels == 3:
|
||||||
|
return pixels.expand(-1, 3, -1, -1)
|
||||||
|
raise ValueError("subgoal image has unsupported channel count")
|
||||||
|
|
||||||
|
def encode_prefix(
|
||||||
|
self,
|
||||||
|
pixel_values: Tensor,
|
||||||
|
history_pixel_values: Tensor | None,
|
||||||
|
history_states: Tensor | None,
|
||||||
|
current_state: Tensor,
|
||||||
|
input_ids: Tensor,
|
||||||
|
subgoal_pixel_values: Tensor | None = None,
|
||||||
|
history_mask: Tensor | None = None,
|
||||||
|
subgoal_mask: Tensor | None = None,
|
||||||
|
attention_mask: Tensor | None = None,
|
||||||
|
) -> tuple[Tensor, Tensor, PrefixLayout]:
|
||||||
|
observation_images = self.memory_encoder(
|
||||||
|
pixel_values, history_pixel_values, history_mask
|
||||||
|
)
|
||||||
|
batch = pixel_values.shape[0]
|
||||||
|
if current_state.shape != (batch, self.config.state_dim):
|
||||||
|
raise ValueError("current_state must have shape [batch, state_dim]")
|
||||||
|
|
||||||
|
if history_states is None:
|
||||||
|
history_states = current_state.new_empty(batch, 0, self.config.state_dim)
|
||||||
|
if history_states.ndim != 3 or history_states.shape[0] != batch:
|
||||||
|
raise ValueError("history_states must have shape [batch, time, state_dim]")
|
||||||
|
if history_states.shape[-1] != self.config.state_dim:
|
||||||
|
raise ValueError("history_states has the wrong state dimension")
|
||||||
|
history_length = history_states.shape[1]
|
||||||
|
if history_length > self.config.max_history_frames:
|
||||||
|
raise ValueError("state history exceeds configured maximum")
|
||||||
|
if history_pixel_values is not None and history_pixel_values.shape[1] != history_length:
|
||||||
|
raise ValueError("visual and proprioceptive history lengths must match")
|
||||||
|
if history_mask is None:
|
||||||
|
history_mask = torch.ones(
|
||||||
|
batch, history_length, dtype=torch.bool, device=pixel_values.device
|
||||||
|
)
|
||||||
|
if history_mask.shape != (batch, history_length):
|
||||||
|
raise ValueError("history_mask must have shape [batch, history]")
|
||||||
|
|
||||||
|
all_states = torch.cat((history_states, current_state[:, None]), dim=1)
|
||||||
|
state_tokens = self.state_projection(all_states)
|
||||||
|
ages = torch.arange(history_length, -1, -1, device=state_tokens.device)
|
||||||
|
state_tokens = state_tokens + MEMVideoEncoder.temporal_position(
|
||||||
|
ages, self.config.d_model
|
||||||
|
)[None].to(state_tokens.dtype)
|
||||||
|
state_valid = torch.cat(
|
||||||
|
(
|
||||||
|
history_mask.bool(),
|
||||||
|
torch.ones(batch, 1, dtype=torch.bool, device=pixel_values.device),
|
||||||
|
),
|
||||||
|
dim=1,
|
||||||
|
)
|
||||||
|
state_tokens = state_tokens * state_valid[:, :, None].to(state_tokens.dtype)
|
||||||
|
observation = torch.cat((observation_images, state_tokens), dim=1)
|
||||||
|
observation_valid = torch.cat(
|
||||||
|
(
|
||||||
|
torch.ones(
|
||||||
|
batch,
|
||||||
|
observation_images.shape[1],
|
||||||
|
dtype=torch.bool,
|
||||||
|
device=pixel_values.device,
|
||||||
|
),
|
||||||
|
state_valid,
|
||||||
|
),
|
||||||
|
dim=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
if subgoal_pixel_values is None:
|
||||||
|
goal = observation.new_empty(batch, 0, self.config.d_model)
|
||||||
|
goal_valid = torch.empty(batch, 0, dtype=torch.bool, device=pixel_values.device)
|
||||||
|
else:
|
||||||
|
goal_pixels = self._adapt_goal_channels(subgoal_pixel_values)
|
||||||
|
goal = self.memory_encoder(goal_pixels)
|
||||||
|
if subgoal_mask is None:
|
||||||
|
subgoal_mask = torch.ones(batch, dtype=torch.bool, device=pixel_values.device)
|
||||||
|
if subgoal_mask.shape != (batch,):
|
||||||
|
raise ValueError("subgoal_mask must have shape [batch]")
|
||||||
|
goal_valid = subgoal_mask.bool()[:, None].expand(batch, goal.shape[1])
|
||||||
|
goal = goal * goal_valid[:, :, None].to(goal.dtype)
|
||||||
|
|
||||||
|
if input_ids.ndim != 2 or input_ids.shape[0] != batch:
|
||||||
|
raise ValueError("input_ids must have shape [batch, text]")
|
||||||
|
text = self.token_embedding(input_ids)
|
||||||
|
if attention_mask is None:
|
||||||
|
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
|
||||||
|
if attention_mask.shape != input_ids.shape:
|
||||||
|
raise ValueError("attention_mask must match input_ids")
|
||||||
|
text_valid = attention_mask.bool()
|
||||||
|
|
||||||
|
layout = PrefixLayout(observation.shape[1], goal.shape[1], text.shape[1])
|
||||||
|
if layout.total_tokens > self.config.max_prefix_tokens:
|
||||||
|
raise ValueError("prefix exceeds configured maximum length")
|
||||||
|
prefix = torch.cat((observation, goal, text), dim=1)
|
||||||
|
segments = torch.cat(
|
||||||
|
(
|
||||||
|
torch.zeros(layout.observation_tokens, dtype=torch.long, device=prefix.device),
|
||||||
|
torch.ones(layout.goal_tokens, dtype=torch.long, device=prefix.device),
|
||||||
|
torch.full(
|
||||||
|
(layout.text_tokens,), 2, dtype=torch.long, device=prefix.device
|
||||||
|
),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
prefix = prefix + self.segment_embedding(segments)[None]
|
||||||
|
prefix = prefix + self.prefix_position[:, : layout.total_tokens]
|
||||||
|
valid = torch.cat((observation_valid, goal_valid, text_valid), dim=1)
|
||||||
|
return prefix, valid, layout
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def build_prefix_attention_mask(layout: PrefixLayout, valid: Tensor) -> Tensor:
|
||||||
|
"""Build obs -> goal -> causal-text block attention.
|
||||||
|
|
||||||
|
Observation queries only read observations. Goal-image queries read
|
||||||
|
observations and all goal tokens bidirectionally. Text queries read
|
||||||
|
both visual blocks and earlier text tokens causally.
|
||||||
|
"""
|
||||||
|
|
||||||
|
length = layout.total_tokens
|
||||||
|
if valid.shape[1] != length:
|
||||||
|
raise ValueError("valid mask does not match PrefixLayout")
|
||||||
|
device = valid.device
|
||||||
|
q = torch.arange(length, device=device)[:, None]
|
||||||
|
k = torch.arange(length, device=device)[None]
|
||||||
|
obs = q < layout.goal_start
|
||||||
|
goal = (q >= layout.goal_start) & (q < layout.text_start)
|
||||||
|
text = q >= layout.text_start
|
||||||
|
structural = (
|
||||||
|
(obs & (k < layout.goal_start))
|
||||||
|
| (goal & (k < layout.text_start))
|
||||||
|
| (text & ((k < layout.text_start) | ((k >= layout.text_start) & (k <= q))))
|
||||||
|
)
|
||||||
|
edges = structural[None] & valid[:, None]
|
||||||
|
identity = torch.eye(length, dtype=torch.bool, device=device)[None]
|
||||||
|
return torch.where(valid[:, :, None], edges, identity)
|
||||||
|
|
||||||
|
def _derive_action_prefix_mask(self, input_ids: Tensor) -> Tensor:
|
||||||
|
is_fast = input_ids.eq(self.config.action_start_token_id)
|
||||||
|
return is_fast.long().cumsum(dim=1).eq(0)
|
||||||
|
|
||||||
|
def _action_attention_mask(
|
||||||
|
self,
|
||||||
|
valid: Tensor,
|
||||||
|
layout: PrefixLayout,
|
||||||
|
action_length: int,
|
||||||
|
action_prefix_mask: Tensor,
|
||||||
|
) -> Tensor:
|
||||||
|
batch = valid.shape[0]
|
||||||
|
supplied = action_prefix_mask.bool()
|
||||||
|
if supplied.shape == (batch, layout.text_tokens):
|
||||||
|
supplied = torch.cat(
|
||||||
|
(
|
||||||
|
torch.ones(
|
||||||
|
batch,
|
||||||
|
layout.text_start,
|
||||||
|
dtype=torch.bool,
|
||||||
|
device=valid.device,
|
||||||
|
),
|
||||||
|
supplied,
|
||||||
|
),
|
||||||
|
dim=1,
|
||||||
|
)
|
||||||
|
if supplied.shape != valid.shape:
|
||||||
|
raise ValueError("action_prefix_mask must cover text or the full prefix")
|
||||||
|
readable = valid & supplied
|
||||||
|
prefix_keys = readable[:, None].expand(batch, action_length, layout.total_tokens)
|
||||||
|
action_keys = torch.ones(
|
||||||
|
batch, action_length, action_length, dtype=torch.bool, device=valid.device
|
||||||
|
)
|
||||||
|
return torch.cat((prefix_keys, action_keys), dim=-1)
|
||||||
|
|
||||||
|
def vlm_parameters(self) -> Iterator[nn.Parameter]:
|
||||||
|
modules: list[nn.Module] = [
|
||||||
|
self.memory_encoder,
|
||||||
|
self.state_projection,
|
||||||
|
self.token_embedding,
|
||||||
|
self.segment_embedding,
|
||||||
|
self.prefix_norm,
|
||||||
|
self.text_head,
|
||||||
|
]
|
||||||
|
for layer in self.layers:
|
||||||
|
modules.extend(
|
||||||
|
(layer.vlm_norm1, layer.vlm_attention, layer.vlm_norm2, layer.vlm_mlp)
|
||||||
|
)
|
||||||
|
seen: set[int] = set()
|
||||||
|
if id(self.prefix_position) not in seen:
|
||||||
|
seen.add(id(self.prefix_position))
|
||||||
|
yield self.prefix_position
|
||||||
|
for module in modules:
|
||||||
|
for parameter in module.parameters():
|
||||||
|
if id(parameter) not in seen:
|
||||||
|
seen.add(id(parameter))
|
||||||
|
yield parameter
|
||||||
|
|
||||||
|
def action_parameters(self) -> Iterator[nn.Parameter]:
|
||||||
|
modules: list[nn.Module] = [
|
||||||
|
self.action_projection,
|
||||||
|
self.time_embedding,
|
||||||
|
self.action_norm,
|
||||||
|
self.vector_field_head,
|
||||||
|
]
|
||||||
|
for layer in self.layers:
|
||||||
|
modules.extend(
|
||||||
|
(
|
||||||
|
layer.action_norm1,
|
||||||
|
layer.action_attention,
|
||||||
|
layer.action_norm2,
|
||||||
|
layer.action_mlp,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
yield self.action_position
|
||||||
|
for module in modules:
|
||||||
|
yield from module.parameters()
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
pixel_values: Tensor,
|
||||||
|
input_ids: Tensor,
|
||||||
|
noisy_actions: Tensor,
|
||||||
|
timesteps: Tensor,
|
||||||
|
history_pixel_values: Tensor | None = None,
|
||||||
|
history_states: Tensor | None = None,
|
||||||
|
current_state: Tensor | None = None,
|
||||||
|
subgoal_pixel_values: Tensor | None = None,
|
||||||
|
history_mask: Tensor | None = None,
|
||||||
|
subgoal_mask: Tensor | None = None,
|
||||||
|
attention_mask: Tensor | None = None,
|
||||||
|
action_prefix_mask: Tensor | None = None,
|
||||||
|
generated_token_start: int | None = None,
|
||||||
|
text_labels: Tensor | None = None,
|
||||||
|
target_velocity: Tensor | None = None,
|
||||||
|
text_loss_weight: float = 1.0,
|
||||||
|
flow_loss_weight: float = 1.0,
|
||||||
|
) -> Pi07Output:
|
||||||
|
if generated_token_start not in (None, 0):
|
||||||
|
raise ValueError("pi0.7 prompt text is causal from its first token")
|
||||||
|
if noisy_actions.ndim != 3:
|
||||||
|
raise ValueError("noisy_actions must have shape [batch, horizon, action_dim]")
|
||||||
|
if noisy_actions.shape[-1] != self.config.action_dim:
|
||||||
|
raise ValueError("noisy_actions has the wrong action dimension")
|
||||||
|
if noisy_actions.shape[1] > self.config.max_action_tokens:
|
||||||
|
raise ValueError("action horizon exceeds configured maximum")
|
||||||
|
if current_state is None:
|
||||||
|
current_state = noisy_actions.new_zeros(
|
||||||
|
noisy_actions.shape[0], self.config.state_dim
|
||||||
|
)
|
||||||
|
|
||||||
|
prefix, valid, layout = self.encode_prefix(
|
||||||
|
pixel_values,
|
||||||
|
history_pixel_values,
|
||||||
|
history_states,
|
||||||
|
current_state,
|
||||||
|
input_ids,
|
||||||
|
subgoal_pixel_values,
|
||||||
|
history_mask,
|
||||||
|
subgoal_mask,
|
||||||
|
attention_mask,
|
||||||
|
)
|
||||||
|
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)
|
||||||
|
if action_prefix_mask is None:
|
||||||
|
action_prefix_mask = self._derive_action_prefix_mask(input_ids)
|
||||||
|
prefix_edges = self.build_prefix_attention_mask(layout, valid)
|
||||||
|
action_edges = self._action_attention_mask(
|
||||||
|
valid, layout, horizon, action_prefix_mask
|
||||||
|
)
|
||||||
|
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)
|
||||||
|
text_logits = self.text_head(prefix_hidden[:, layout.text_start :])
|
||||||
|
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)
|
||||||
|
terms = [
|
||||||
|
term
|
||||||
|
for term 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 term is not None
|
||||||
|
]
|
||||||
|
loss = sum(terms) if terms else None
|
||||||
|
return Pi07Output(
|
||||||
|
text_logits=text_logits,
|
||||||
|
vector_field=vector_field,
|
||||||
|
prefix_hidden=prefix_hidden,
|
||||||
|
action_hidden=action_hidden,
|
||||||
|
layout=layout,
|
||||||
|
loss=loss,
|
||||||
|
text_loss=text_loss,
|
||||||
|
flow_loss=flow_loss,
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def combine_metadata_cfg(
|
||||||
|
conditioned: Tensor, unconditioned: Tensor, guidance_scale: float = 1.5
|
||||||
|
) -> Tensor:
|
||||||
|
"""pi0.7 metadata CFG: conditional + beta*(conditional-unconditional)."""
|
||||||
|
|
||||||
|
if guidance_scale < 0:
|
||||||
|
raise ValueError("guidance_scale must be non-negative")
|
||||||
|
if conditioned.shape != unconditioned.shape:
|
||||||
|
raise ValueError("conditioned and unconditioned predictions must match")
|
||||||
|
return conditioned + guidance_scale * (conditioned - unconditioned)
|
||||||
|
|
||||||
|
def metadata_cfg(
|
||||||
|
self,
|
||||||
|
conditioned_input_ids: Tensor,
|
||||||
|
unconditioned_input_ids: Tensor,
|
||||||
|
*,
|
||||||
|
guidance_scale: float = 1.5,
|
||||||
|
conditioned_attention_mask: Tensor | None = None,
|
||||||
|
unconditioned_attention_mask: Tensor | None = None,
|
||||||
|
conditioned_action_prefix_mask: Tensor | None = None,
|
||||||
|
unconditioned_action_prefix_mask: Tensor | None = None,
|
||||||
|
**common_inputs: Tensor | None,
|
||||||
|
) -> Tensor:
|
||||||
|
"""Run metadata-positive and metadata-dropped branches and combine them."""
|
||||||
|
|
||||||
|
conditioned = self(
|
||||||
|
input_ids=conditioned_input_ids,
|
||||||
|
attention_mask=conditioned_attention_mask,
|
||||||
|
action_prefix_mask=conditioned_action_prefix_mask,
|
||||||
|
**common_inputs,
|
||||||
|
).vector_field
|
||||||
|
unconditioned = self(
|
||||||
|
input_ids=unconditioned_input_ids,
|
||||||
|
attention_mask=unconditioned_attention_mask,
|
||||||
|
action_prefix_mask=unconditioned_action_prefix_mask,
|
||||||
|
**common_inputs,
|
||||||
|
).vector_field
|
||||||
|
return self.combine_metadata_cfg(conditioned, unconditioned, guidance_scale)
|
||||||
|
|
||||||
|
|
||||||
|
def flow_matching_sample(actions: Tensor) -> tuple[Tensor, Tensor, Tensor]:
|
||||||
|
"""Construct the linear conditional-flow target used by the action expert."""
|
||||||
|
|
||||||
|
batch = actions.shape[0]
|
||||||
|
timestep = torch.rand(batch, device=actions.device, dtype=actions.dtype)
|
||||||
|
noise = torch.randn_like(actions)
|
||||||
|
broadcast = timestep.reshape(batch, *([1] * (actions.ndim - 1)))
|
||||||
|
noisy_actions = (1 - broadcast) * noise + broadcast * actions
|
||||||
|
return noisy_actions, timestep, actions - noise
|
||||||
|
|
||||||
|
|
||||||
|
class HighLevelPolicy(nn.Module):
|
||||||
|
"""Tiny semantic policy that predicts the next subtask prompt token.
|
||||||
|
|
||||||
|
The paper uses a separate policy based on the same VLM family. This proxy
|
||||||
|
shares the same MEM input contract while replacing the 4B VLM with a small
|
||||||
|
Transformer encoder suitable for a few-example overfit test.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, config: Optional[Pi07Config] = None) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.config = config or Pi07Config()
|
||||||
|
c = self.config
|
||||||
|
self.memory_encoder = MEMVideoEncoder(
|
||||||
|
MEMConfig(
|
||||||
|
image_channels=c.image_channels,
|
||||||
|
image_size=c.image_size,
|
||||||
|
patch_size=c.patch_size,
|
||||||
|
d_model=c.d_model,
|
||||||
|
n_heads=c.n_heads,
|
||||||
|
n_layers=c.memory_layers,
|
||||||
|
mlp_ratio=c.mlp_ratio,
|
||||||
|
max_history_frames=c.max_history_frames,
|
||||||
|
dropout=c.dropout,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
self.state_projection = nn.Linear(c.state_dim, c.d_model)
|
||||||
|
self.token_embedding = nn.Embedding(c.vocab_size, c.d_model)
|
||||||
|
layer = nn.TransformerEncoderLayer(
|
||||||
|
c.d_model,
|
||||||
|
c.n_heads,
|
||||||
|
c.d_model * c.mlp_ratio,
|
||||||
|
c.dropout,
|
||||||
|
activation="gelu",
|
||||||
|
batch_first=True,
|
||||||
|
norm_first=True,
|
||||||
|
)
|
||||||
|
self.transformer = nn.TransformerEncoder(layer, c.n_layers)
|
||||||
|
self.norm = nn.LayerNorm(c.d_model)
|
||||||
|
self.head = nn.Linear(c.d_model, c.vocab_size, bias=False)
|
||||||
|
self.head.weight = self.token_embedding.weight
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
pixel_values: Tensor,
|
||||||
|
input_ids: Tensor,
|
||||||
|
history_pixel_values: Tensor | None = None,
|
||||||
|
history_states: Tensor | None = None,
|
||||||
|
current_state: Tensor | None = None,
|
||||||
|
history_mask: Tensor | None = None,
|
||||||
|
attention_mask: Tensor | None = None,
|
||||||
|
target_token: Tensor | None = None,
|
||||||
|
) -> HighLevelOutput:
|
||||||
|
batch = pixel_values.shape[0]
|
||||||
|
if current_state is None:
|
||||||
|
current_state = pixel_values.new_zeros(batch, self.config.state_dim)
|
||||||
|
visual = self.memory_encoder(pixel_values, history_pixel_values, history_mask)
|
||||||
|
if history_states is None:
|
||||||
|
history_states = current_state.new_empty(batch, 0, self.config.state_dim)
|
||||||
|
states = self.state_projection(torch.cat((history_states, current_state[:, None]), 1))
|
||||||
|
text = self.token_embedding(input_ids)
|
||||||
|
tokens = torch.cat((visual, states, text), dim=1)
|
||||||
|
observation_length = visual.shape[1] + states.shape[1]
|
||||||
|
length = tokens.shape[1]
|
||||||
|
q = torch.arange(length, device=tokens.device)[:, None]
|
||||||
|
k = torch.arange(length, device=tokens.device)[None]
|
||||||
|
allowed = torch.where(q < observation_length, k < observation_length, k <= q)
|
||||||
|
blocked = ~allowed
|
||||||
|
if attention_mask is None:
|
||||||
|
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
|
||||||
|
state_valid = torch.ones(batch, states.shape[1], dtype=torch.bool, device=tokens.device)
|
||||||
|
if history_mask is not None and states.shape[1] > 1:
|
||||||
|
state_valid[:, :-1] = history_mask.bool()
|
||||||
|
valid = torch.cat(
|
||||||
|
(
|
||||||
|
torch.ones(batch, visual.shape[1], dtype=torch.bool, device=tokens.device),
|
||||||
|
state_valid,
|
||||||
|
attention_mask.bool(),
|
||||||
|
),
|
||||||
|
dim=1,
|
||||||
|
)
|
||||||
|
hidden = self.transformer(tokens, mask=blocked, src_key_padding_mask=~valid)
|
||||||
|
text_hidden = self.norm(hidden[:, observation_length:])
|
||||||
|
lengths = attention_mask.long().sum(dim=1).clamp_min(1) - 1
|
||||||
|
pooled = text_hidden[torch.arange(batch, device=tokens.device), lengths]
|
||||||
|
logits = self.head(pooled)
|
||||||
|
loss = None if target_token is None else F.cross_entropy(logits, target_token.long())
|
||||||
|
return HighLevelOutput(logits, loss)
|
||||||
|
|
||||||
|
|
||||||
|
class TinyGoalGenerator(nn.Module):
|
||||||
|
"""Small conditional-flow image model standing in for the BAGEL world model."""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
vocab_size: int,
|
||||||
|
image_channels: int = 3,
|
||||||
|
hidden_dim: int = 48,
|
||||||
|
text_dim: int = 64,
|
||||||
|
) -> None:
|
||||||
|
super().__init__()
|
||||||
|
self.image_channels = image_channels
|
||||||
|
self.token_embedding = nn.Embedding(vocab_size, text_dim)
|
||||||
|
self.text_projection = nn.Linear(text_dim, hidden_dim)
|
||||||
|
self.time_embedding = TimeEmbedding(hidden_dim)
|
||||||
|
self.input_projection = nn.Conv2d(2 * image_channels, hidden_dim, 3, padding=1)
|
||||||
|
self.network = nn.Sequential(
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Conv2d(hidden_dim, hidden_dim, 3, padding=1),
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Conv2d(hidden_dim, image_channels, 3, padding=1),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
current_pixel_values: Tensor,
|
||||||
|
input_ids: Tensor,
|
||||||
|
noisy_subgoal: Tensor,
|
||||||
|
timesteps: Tensor,
|
||||||
|
attention_mask: Tensor | None = None,
|
||||||
|
target_velocity: Tensor | None = None,
|
||||||
|
) -> GoalGeneratorOutput:
|
||||||
|
if current_pixel_values.shape != noisy_subgoal.shape:
|
||||||
|
raise ValueError("current image and noisy subgoal must have identical shapes")
|
||||||
|
if current_pixel_values.ndim != 4 or current_pixel_values.shape[1] != self.image_channels:
|
||||||
|
raise ValueError("goal-generator images have the wrong shape")
|
||||||
|
if attention_mask is None:
|
||||||
|
attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
|
||||||
|
embedded = self.token_embedding(input_ids)
|
||||||
|
weights = attention_mask.to(embedded.dtype)[..., None]
|
||||||
|
pooled = (embedded * weights).sum(1) / weights.sum(1).clamp_min(1)
|
||||||
|
condition = self.text_projection(pooled) + self.time_embedding(timesteps)
|
||||||
|
hidden = self.input_projection(
|
||||||
|
torch.cat((current_pixel_values, noisy_subgoal), dim=1)
|
||||||
|
)
|
||||||
|
hidden = hidden + condition[:, :, None, None].to(hidden.dtype)
|
||||||
|
vector_field = self.network(hidden)
|
||||||
|
loss = None
|
||||||
|
if target_velocity is not None:
|
||||||
|
loss = F.mse_loss(vector_field, target_velocity)
|
||||||
|
return GoalGeneratorOutput(vector_field, loss)
|
||||||
@@ -0,0 +1,62 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
|
||||||
|
MAX_RTC_DELAY_STEPS = 12
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class RTCAlignment:
|
||||||
|
"""A fixed-width target aligned to an inference/control delay."""
|
||||||
|
|
||||||
|
delayed_target_actions: torch.Tensor
|
||||||
|
valid_mask: torch.Tensor
|
||||||
|
executed_prefix_actions: torch.Tensor
|
||||||
|
delay_steps: int
|
||||||
|
|
||||||
|
|
||||||
|
def align_actions_for_rtc(
|
||||||
|
actions: torch.Tensor,
|
||||||
|
delay_steps: int,
|
||||||
|
*,
|
||||||
|
max_delay_steps: int = MAX_RTC_DELAY_STEPS,
|
||||||
|
) -> RTCAlignment:
|
||||||
|
"""Shift targets left so prediction step zero starts after control delay.
|
||||||
|
|
||||||
|
For delay ``d``, target position ``i`` supervises original action ``i+d``.
|
||||||
|
Vacated tail positions are zero padded and excluded by ``valid_mask``.
|
||||||
|
The already-executed prefix is returned for optional action-prefix input.
|
||||||
|
"""
|
||||||
|
|
||||||
|
actions = torch.as_tensor(actions, dtype=torch.float32)
|
||||||
|
if actions.ndim < 2 or actions.shape[-2] < 1:
|
||||||
|
raise ValueError("actions must have shape [..., horizon, action_dim]")
|
||||||
|
if not 0 <= delay_steps <= max_delay_steps:
|
||||||
|
raise ValueError(f"delay_steps must lie in [0, {max_delay_steps}]")
|
||||||
|
horizon = actions.shape[-2]
|
||||||
|
if delay_steps >= horizon:
|
||||||
|
raise ValueError("delay_steps must be smaller than the action horizon")
|
||||||
|
aligned = torch.zeros_like(actions)
|
||||||
|
valid = torch.zeros(actions.shape[:-1], dtype=torch.bool, device=actions.device)
|
||||||
|
remaining = horizon - delay_steps
|
||||||
|
aligned[..., :remaining, :] = actions[..., delay_steps:, :]
|
||||||
|
valid[..., :remaining] = True
|
||||||
|
prefix = actions[..., :delay_steps, :].clone()
|
||||||
|
return RTCAlignment(aligned, valid, prefix, int(delay_steps))
|
||||||
|
|
||||||
|
|
||||||
|
def sample_rtc_delay(
|
||||||
|
*,
|
||||||
|
generator: torch.Generator | None = None,
|
||||||
|
max_delay_steps: int = MAX_RTC_DELAY_STEPS,
|
||||||
|
action_horizon: int,
|
||||||
|
) -> int:
|
||||||
|
if action_horizon < 1:
|
||||||
|
raise ValueError("action_horizon must be positive")
|
||||||
|
maximum = min(int(max_delay_steps), action_horizon - 1)
|
||||||
|
if maximum < 0:
|
||||||
|
raise ValueError("no valid delay for an empty action horizon")
|
||||||
|
return int(torch.randint(maximum + 1, (), generator=generator).item())
|
||||||
@@ -0,0 +1,14 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
import json,random
|
||||||
|
from pathlib import Path
|
||||||
|
import numpy as np,torch,yaml
|
||||||
|
def load_config(path):
|
||||||
|
with Path(path).open() as f:return yaml.safe_load(f)
|
||||||
|
def seed_everything(seed:int):random.seed(seed);np.random.seed(seed);torch.manual_seed(seed)
|
||||||
|
def choose_device(name="auto"):
|
||||||
|
if name!="auto":return torch.device(name)
|
||||||
|
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,value):
|
||||||
|
p=Path(path);p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(value,indent=2,ensure_ascii=False))
|
||||||
@@ -0,0 +1,393 @@
|
|||||||
|
"""Prompt and FAST-lite tokenization for the minimal pi0.7 policy.
|
||||||
|
|
||||||
|
pi0.7 conditions the low-level policy on a richer prompt than earlier pi
|
||||||
|
models: task and subtask language, episode speed/quality/mistake metadata,
|
||||||
|
control mode, and data provenance. This tokenizer represents those concepts
|
||||||
|
with explicit marker/value pairs so the architecture can be validated without
|
||||||
|
shipping a full language tokenizer.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import math
|
||||||
|
from dataclasses import dataclass
|
||||||
|
from typing import Iterable, Literal
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import Tensor
|
||||||
|
|
||||||
|
|
||||||
|
ControlMode = Literal["joint", "ee"]
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class EpisodeMetadata:
|
||||||
|
speed_steps: int | None = None
|
||||||
|
quality: int | None = None
|
||||||
|
mistake: bool | None = None
|
||||||
|
control_mode: ControlMode = "joint"
|
||||||
|
source_type: str | None = None
|
||||||
|
|
||||||
|
def __post_init__(self) -> None:
|
||||||
|
if self.speed_steps is not None and self.speed_steps <= 0:
|
||||||
|
raise ValueError("speed_steps must be positive")
|
||||||
|
if self.quality is not None and not 1 <= self.quality <= 5:
|
||||||
|
raise ValueError("quality must be in the inclusive range 1..5")
|
||||||
|
if self.control_mode not in ("joint", "ee"):
|
||||||
|
raise ValueError("control_mode must be 'joint' or 'ee'")
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass(frozen=True)
|
||||||
|
class PolicySequence:
|
||||||
|
input_ids: Tensor
|
||||||
|
labels: Tensor
|
||||||
|
action_prefix_mask: Tensor
|
||||||
|
generated_token_start: int
|
||||||
|
action_start_index: int
|
||||||
|
prompt_length: int
|
||||||
|
metadata_indices: tuple[int, ...]
|
||||||
|
|
||||||
|
|
||||||
|
class _LabelVocabulary:
|
||||||
|
def __init__(self, labels: tuple[str, ...], first_token_id: int) -> None:
|
||||||
|
self.labels = labels
|
||||||
|
self.first_token_id = first_token_id
|
||||||
|
self._to_token = {
|
||||||
|
label: first_token_id + index for index, label in enumerate(labels)
|
||||||
|
}
|
||||||
|
self._to_label = {token: label for label, token in self._to_token.items()}
|
||||||
|
|
||||||
|
@property
|
||||||
|
def token_ids(self) -> tuple[int, ...]:
|
||||||
|
return tuple(self._to_token[label] for label in self.labels)
|
||||||
|
|
||||||
|
def encode(self, label: str) -> int:
|
||||||
|
try:
|
||||||
|
return self._to_token[label]
|
||||||
|
except KeyError as error:
|
||||||
|
raise ValueError(f"unknown categorical label: {label}") from error
|
||||||
|
|
||||||
|
def decode(self, token: int | Tensor) -> str:
|
||||||
|
try:
|
||||||
|
return self._to_label[int(token)]
|
||||||
|
except KeyError as error:
|
||||||
|
raise ValueError(f"unknown categorical token: {int(token)}") from error
|
||||||
|
|
||||||
|
|
||||||
|
class MiniTokenizer:
|
||||||
|
"""Explicit rich-prompt tokens plus a DCT-based FAST-lite action stream."""
|
||||||
|
|
||||||
|
BOS = 256
|
||||||
|
TASK_MARKER = 257
|
||||||
|
SUBTASK_MARKER = 258
|
||||||
|
SPEED_MARKER = 259
|
||||||
|
QUALITY_MARKER = 260
|
||||||
|
MISTAKE_MARKER = 261
|
||||||
|
CONTROL_MARKER = 262
|
||||||
|
SOURCE_MARKER = 263
|
||||||
|
|
||||||
|
TASK_LABELS = (
|
||||||
|
"drive",
|
||||||
|
"pick_place",
|
||||||
|
"fold",
|
||||||
|
"open_close",
|
||||||
|
"clean",
|
||||||
|
"cook",
|
||||||
|
"unknown",
|
||||||
|
)
|
||||||
|
TASK_TOKEN_START = 264
|
||||||
|
TASK_TOKEN_END = TASK_TOKEN_START + len(TASK_LABELS) - 1
|
||||||
|
|
||||||
|
SUBTASK_LABELS = (
|
||||||
|
"keep_lane",
|
||||||
|
"steer_left",
|
||||||
|
"steer_right",
|
||||||
|
"slow_down",
|
||||||
|
"stop",
|
||||||
|
"accelerate",
|
||||||
|
"pick_place",
|
||||||
|
"fold",
|
||||||
|
"open_close",
|
||||||
|
"unknown",
|
||||||
|
)
|
||||||
|
SUBTASK_TOKEN_START = TASK_TOKEN_END + 1
|
||||||
|
SUBTASK_TOKEN_END = SUBTASK_TOKEN_START + len(SUBTASK_LABELS) - 1
|
||||||
|
|
||||||
|
SPEED_BINS = (500, 1000, 2000, 4000, 8000, 16000)
|
||||||
|
SPEED_TOKEN_START = SUBTASK_TOKEN_END + 1
|
||||||
|
SPEED_TOKEN_END = SPEED_TOKEN_START + len(SPEED_BINS) - 1
|
||||||
|
QUALITY_TOKEN_START = SPEED_TOKEN_END + 1
|
||||||
|
QUALITY_TOKEN_END = QUALITY_TOKEN_START + 4
|
||||||
|
MISTAKE_FALSE = QUALITY_TOKEN_END + 1
|
||||||
|
MISTAKE_TRUE = MISTAKE_FALSE + 1
|
||||||
|
CONTROL_JOINT = MISTAKE_TRUE + 1
|
||||||
|
CONTROL_EE = CONTROL_JOINT + 1
|
||||||
|
|
||||||
|
SOURCE_LABELS = (
|
||||||
|
"demonstration",
|
||||||
|
"autonomous",
|
||||||
|
"intervention",
|
||||||
|
"human_video",
|
||||||
|
"web",
|
||||||
|
"unknown",
|
||||||
|
)
|
||||||
|
SOURCE_TOKEN_START = CONTROL_EE + 1
|
||||||
|
SOURCE_TOKEN_END = SOURCE_TOKEN_START + len(SOURCE_LABELS) - 1
|
||||||
|
|
||||||
|
ACTION_START = SOURCE_TOKEN_END + 1
|
||||||
|
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) -> None:
|
||||||
|
self.tasks = _LabelVocabulary(self.TASK_LABELS, self.TASK_TOKEN_START)
|
||||||
|
self.subtasks = _LabelVocabulary(self.SUBTASK_LABELS, self.SUBTASK_TOKEN_START)
|
||||||
|
self.sources = _LabelVocabulary(self.SOURCE_LABELS, self.SOURCE_TOKEN_START)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def vocab_size(self) -> int:
|
||||||
|
return self.ACTION_TOKEN_END + 1
|
||||||
|
|
||||||
|
@property
|
||||||
|
def action_token_ids(self) -> range:
|
||||||
|
return range(self.ACTION_TOKEN_START, self.ACTION_TOKEN_END + 1)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _classify_task(description: str) -> str:
|
||||||
|
text = description.casefold()
|
||||||
|
if any(word in text for word in ("drive", "lane", "road", "steer", "vehicle")):
|
||||||
|
return "drive"
|
||||||
|
if any(word in text for word in ("fold", "shirt", "towel", "laundry")):
|
||||||
|
return "fold"
|
||||||
|
if any(word in text for word in ("open", "close", "door", "drawer", "fridge")):
|
||||||
|
return "open_close"
|
||||||
|
if any(word in text for word in ("clean", "wipe", "wash", "tidy")):
|
||||||
|
return "clean"
|
||||||
|
if any(word in text for word in ("cook", "food", "recipe", "microwave")):
|
||||||
|
return "cook"
|
||||||
|
if any(word in text for word in ("pick", "place", "move", "put")):
|
||||||
|
return "pick_place"
|
||||||
|
return description if description in MiniTokenizer.TASK_LABELS else "unknown"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _classify_subtask(description: str) -> str:
|
||||||
|
if description in MiniTokenizer.SUBTASK_LABELS:
|
||||||
|
return description
|
||||||
|
text = description.casefold()
|
||||||
|
if any(word in text for word in ("stop", "halt", "stationary")):
|
||||||
|
return "stop"
|
||||||
|
if any(word in text for word in ("slow", "brake", "deceler")):
|
||||||
|
return "slow_down"
|
||||||
|
if any(word in text for word in ("accelerat", "speed up")):
|
||||||
|
return "accelerate"
|
||||||
|
if "left" in text:
|
||||||
|
return "steer_left"
|
||||||
|
if "right" in text:
|
||||||
|
return "steer_right"
|
||||||
|
if any(word in text for word in ("lane", "straight", "continue", "follow")):
|
||||||
|
return "keep_lane"
|
||||||
|
if "fold" in text:
|
||||||
|
return "fold"
|
||||||
|
if any(word in text for word in ("open", "close")):
|
||||||
|
return "open_close"
|
||||||
|
if any(word in text for word in ("pick", "place", "put")):
|
||||||
|
return "pick_place"
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _classify_source(source: str) -> str:
|
||||||
|
if source in MiniTokenizer.SOURCE_LABELS:
|
||||||
|
return source
|
||||||
|
text = source.casefold()
|
||||||
|
if any(word in text for word in ("demo", "teleop")):
|
||||||
|
return "demonstration"
|
||||||
|
if any(word in text for word in ("autonomous", "rollout", "policy")):
|
||||||
|
return "autonomous"
|
||||||
|
if any(word in text for word in ("intervention", "correction")):
|
||||||
|
return "intervention"
|
||||||
|
if any(word in text for word in ("human", "ego", "video")):
|
||||||
|
return "human_video"
|
||||||
|
if any(word in text for word in ("web", "internet")):
|
||||||
|
return "web"
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
def encode_prompt(
|
||||||
|
self,
|
||||||
|
task: str,
|
||||||
|
subtask: str | None,
|
||||||
|
metadata: EpisodeMetadata,
|
||||||
|
*,
|
||||||
|
include_episode_metadata: bool = True,
|
||||||
|
include_subtask: bool = True,
|
||||||
|
) -> tuple[Tensor, tuple[int, ...]]:
|
||||||
|
"""Encode prompt and return positions controlled by metadata CFG.
|
||||||
|
|
||||||
|
Control mode is deliberately retained even in the unconditional branch,
|
||||||
|
matching pi0.7 training where control-mode dropout is not applied.
|
||||||
|
"""
|
||||||
|
|
||||||
|
ids = [self.BOS, self.TASK_MARKER, self.tasks.encode(self._classify_task(task))]
|
||||||
|
if include_subtask and subtask is not None:
|
||||||
|
ids.extend(
|
||||||
|
(self.SUBTASK_MARKER, self.subtasks.encode(self._classify_subtask(subtask)))
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata_indices: list[int] = []
|
||||||
|
|
||||||
|
def add_metadata(marker: int, value: int) -> None:
|
||||||
|
metadata_indices.extend((len(ids), len(ids) + 1))
|
||||||
|
ids.extend((marker, value))
|
||||||
|
|
||||||
|
if include_episode_metadata:
|
||||||
|
if metadata.speed_steps is not None:
|
||||||
|
bin_index = min(
|
||||||
|
range(len(self.SPEED_BINS)),
|
||||||
|
key=lambda index: abs(self.SPEED_BINS[index] - metadata.speed_steps),
|
||||||
|
)
|
||||||
|
add_metadata(self.SPEED_MARKER, self.SPEED_TOKEN_START + bin_index)
|
||||||
|
if metadata.quality is not None:
|
||||||
|
add_metadata(
|
||||||
|
self.QUALITY_MARKER, self.QUALITY_TOKEN_START + metadata.quality - 1
|
||||||
|
)
|
||||||
|
if metadata.mistake is not None:
|
||||||
|
add_metadata(
|
||||||
|
self.MISTAKE_MARKER,
|
||||||
|
self.MISTAKE_TRUE if metadata.mistake else self.MISTAKE_FALSE,
|
||||||
|
)
|
||||||
|
if metadata.source_type is not None:
|
||||||
|
source = self._classify_source(metadata.source_type)
|
||||||
|
add_metadata(self.SOURCE_MARKER, self.sources.encode(source))
|
||||||
|
|
||||||
|
control = self.CONTROL_JOINT if metadata.control_mode == "joint" else self.CONTROL_EE
|
||||||
|
ids.extend((self.CONTROL_MARKER, control))
|
||||||
|
return torch.tensor(ids, dtype=torch.long), tuple(metadata_indices)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _orthonormal_dct_matrix(
|
||||||
|
horizon: int, *, device: torch.device, dtype: torch.dtype
|
||||||
|
) -> Tensor:
|
||||||
|
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: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||||
|
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||||
|
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||||
|
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||||
|
transform = self._orthonormal_dct_matrix(
|
||||||
|
values.shape[0], device=values.device, dtype=values.dtype
|
||||||
|
)
|
||||||
|
return transform @ values
|
||||||
|
|
||||||
|
def idct_actions(self, coefficients: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||||
|
values = torch.as_tensor(coefficients, dtype=torch.float32)
|
||||||
|
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||||
|
raise ValueError("DCT coefficients must have non-empty shape [horizon, 2]")
|
||||||
|
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: Tensor | Iterable[Iterable[float]]) -> Tensor:
|
||||||
|
values = torch.as_tensor(actions, dtype=torch.float32)
|
||||||
|
if values.ndim != 2 or values.shape[-1] != 2 or values.shape[0] == 0:
|
||||||
|
raise ValueError("actions must have non-empty shape [horizon, 2]")
|
||||||
|
values = values.clamp(self.ACTION_MINIMUM, self.ACTION_MAXIMUM)
|
||||||
|
coefficients = self.dct_actions(values) / math.sqrt(values.shape[0])
|
||||||
|
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)
|
||||||
|
payload = torch.round(unit * (self.ACTION_BINS - 1)).long().reshape(-1)
|
||||||
|
payload = payload + self.ACTION_TOKEN_START
|
||||||
|
boundaries = torch.tensor(
|
||||||
|
(self.ACTION_START, self.ACTION_END), dtype=torch.long, device=payload.device
|
||||||
|
)
|
||||||
|
return torch.cat((boundaries[:1], payload, boundaries[1:]))
|
||||||
|
|
||||||
|
def decode_actions(self, tokens: Tensor | Iterable[int]) -> Tensor:
|
||||||
|
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 FAST-lite interval")
|
||||||
|
quantized = payload - self.ACTION_TOKEN_START
|
||||||
|
coefficients = (
|
||||||
|
quantized.float()
|
||||||
|
/ (self.ACTION_BINS - 1)
|
||||||
|
* (self.ACTION_MAXIMUM - self.ACTION_MINIMUM)
|
||||||
|
+ self.ACTION_MINIMUM
|
||||||
|
)
|
||||||
|
horizon = payload.numel() // 2
|
||||||
|
return self.idct_actions(coefficients.reshape(horizon, 2) * math.sqrt(horizon))
|
||||||
|
|
||||||
|
def build_policy_sequence(
|
||||||
|
self,
|
||||||
|
task: str,
|
||||||
|
subtask: str | None,
|
||||||
|
actions: Tensor,
|
||||||
|
metadata: EpisodeMetadata,
|
||||||
|
*,
|
||||||
|
include_episode_metadata: bool = True,
|
||||||
|
include_subtask: bool = True,
|
||||||
|
) -> PolicySequence:
|
||||||
|
prompt, metadata_indices = self.encode_prompt(
|
||||||
|
task,
|
||||||
|
subtask,
|
||||||
|
metadata,
|
||||||
|
include_episode_metadata=include_episode_metadata,
|
||||||
|
include_subtask=include_subtask,
|
||||||
|
)
|
||||||
|
action_tokens = self.encode_actions(actions).cpu()
|
||||||
|
input_ids = torch.cat((prompt, action_tokens))
|
||||||
|
action_start_index = prompt.numel()
|
||||||
|
labels = torch.full_like(input_ids, -100)
|
||||||
|
labels[action_start_index - 1 : -1] = input_ids[action_start_index:]
|
||||||
|
action_prefix_mask = torch.zeros_like(input_ids, dtype=torch.bool)
|
||||||
|
action_prefix_mask[:action_start_index] = True
|
||||||
|
return PolicySequence(
|
||||||
|
input_ids=input_ids,
|
||||||
|
labels=labels,
|
||||||
|
action_prefix_mask=action_prefix_mask,
|
||||||
|
generated_token_start=0,
|
||||||
|
action_start_index=action_start_index,
|
||||||
|
prompt_length=prompt.numel(),
|
||||||
|
metadata_indices=metadata_indices,
|
||||||
|
)
|
||||||
|
|
||||||
|
def build_metadata_cfg_pair(
|
||||||
|
self,
|
||||||
|
task: str,
|
||||||
|
subtask: str | None,
|
||||||
|
actions: Tensor,
|
||||||
|
metadata: EpisodeMetadata,
|
||||||
|
) -> tuple[PolicySequence, PolicySequence]:
|
||||||
|
"""Return conditioned and episode-metadata-dropped policy sequences."""
|
||||||
|
|
||||||
|
conditioned = self.build_policy_sequence(task, subtask, actions, metadata)
|
||||||
|
unconditional = self.build_policy_sequence(
|
||||||
|
task,
|
||||||
|
subtask,
|
||||||
|
actions,
|
||||||
|
metadata,
|
||||||
|
include_episode_metadata=False,
|
||||||
|
)
|
||||||
|
return conditioned, unconditional
|
||||||
|
|
||||||
@@ -0,0 +1,40 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
import argparse,copy,json
|
||||||
|
from pathlib import Path
|
||||||
|
import torch
|
||||||
|
from .data import Pi07SteerableDataset,materialize_subset
|
||||||
|
from .flow import FlowMatching
|
||||||
|
from .model import Pi07Config,Pi07Mini,HighLevelPolicy,TinyGoalGenerator
|
||||||
|
from .runtime import load_config,seed_everything,choose_device,write_json
|
||||||
|
from .tokenizer import MiniTokenizer,EpisodeMetadata
|
||||||
|
|
||||||
|
def build_config(cfg,tok):
|
||||||
|
m=cfg['model'];return Pi07Config(vocab_size=tok.vocab_size,action_dim=2,state_dim=2,image_size=cfg['data']['image_size'],patch_size=m['patch_size'],d_model=m['d_model'],n_heads=m['n_heads'],n_layers=m['n_layers'],mlp_ratio=m['mlp_ratio'],max_history_frames=cfg['data']['history_frames'],max_text_tokens=m['max_text_tokens'],max_action_tokens=cfg['data']['action_horizon'],action_start_token_id=tok.ACTION_START)
|
||||||
|
|
||||||
|
def train(config_path):
|
||||||
|
cfg=load_config(config_path);root=Path(__file__).resolve().parents[2];seed_everything(cfg['seed']);dev=choose_device(cfg['device'])
|
||||||
|
manifest=materialize_subset(cfg['data']['source_manifest'],root/'data/overfit',max_samples=cfg['data']['max_samples'])
|
||||||
|
ds=Pi07SteerableDataset(manifest,image_size=cfg['data']['image_size'],history_frames=cfg['data']['history_frames'],action_horizon=cfg['data']['action_horizon'],augmentation_factor=cfg['data']['augmentation_factor'])
|
||||||
|
tok=MiniTokenizer();pcfg=build_config(cfg,tok);policy=Pi07Mini(pcfg).to(dev);high=HighLevelPolicy(pcfg).to(dev);goal=TinyGoalGenerator(tok.vocab_size).to(dev);flow=FlowMatching(cfg['flow']['integration_steps']);out=root/cfg['output_dir'];out.mkdir(parents=True,exist_ok=True);history=[]
|
||||||
|
# Goal generator
|
||||||
|
opt=torch.optim.AdamW(goal.parameters(),lr=cfg['train']['learning_rate'])
|
||||||
|
for step in range(1,cfg['train']['goal_steps']+1):
|
||||||
|
s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta);ids=seq.input_ids[:seq.action_start_index][None].to(dev);target=s['subgoal_pixel_values'][None].to(dev);t=torch.rand(1,device=dev);noise=torch.randn_like(target);noisy=t[:,None,None,None]*target+(1-t[:,None,None,None])*noise;o=goal(s['image'][None].to(dev),ids,noisy,t,target_velocity=target-noise);opt.zero_grad();o.loss.backward();opt.step()
|
||||||
|
if step==1 or step%cfg['train']['log_every']==0:row={'stage':'goal','step':step,'loss':float(o.loss.detach().cpu())};history.append(row);print(json.dumps(row))
|
||||||
|
# High-level semantic policy
|
||||||
|
opt=torch.optim.AdamW(high.parameters(),lr=cfg['train']['learning_rate'])
|
||||||
|
for step in range(1,cfg['train']['high_level_steps']+1):
|
||||||
|
s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);prompt,_=tok.encode_prompt('drive safely',None,meta,include_subtask=False);target=torch.tensor([tok.subtasks.encode(tok._classify_subtask(s['reasoning']))],device=dev);o=high(s['image'][None].to(dev),prompt[None].to(dev),s['history_images'][None].to(dev),s['history_states'][None].to(dev),s['current_state'][None].to(dev),target_token=target);opt.zero_grad();o.loss.backward();opt.step()
|
||||||
|
if step==1 or step%cfg['train']['log_every']==0:row={'stage':'high_level','step':step,'loss':float(o.loss.detach().cpu())};history.append(row);print(json.dumps(row))
|
||||||
|
# Full steerable policy, alternating prompt-dropout ablations and RTC targets.
|
||||||
|
opt=torch.optim.AdamW(policy.parameters(),lr=cfg['train']['learning_rate'])
|
||||||
|
for step in range(1,cfg['train']['policy_steps']+1):
|
||||||
|
s=ds[(step-1)%len(ds)];meta=EpisodeMetadata(**s['token_metadata']);include_meta=(step%4)!=0;seq=tok.build_policy_sequence('drive safely',s['reasoning'],s['actions'],meta,include_episode_metadata=include_meta);ids=seq.input_ids[None].to(dev);labels=seq.labels[None].to(dev);target=s['rtc_target_actions'][None].to(dev);noisy,t,vel=flow.training_batch(target);use_goal=(step%4)!=1;use_history=(step%5)!=1
|
||||||
|
o=policy(s['image'][None].to(dev),ids,noisy,t,history_pixel_values=s['history_images'][None].to(dev) if use_history else None,history_states=s['history_states'][None].to(dev) if use_history else None,current_state=s['current_state'][None].to(dev),subgoal_pixel_values=s['subgoal_pixel_values'][None].to(dev) if use_goal else None,action_prefix_mask=seq.action_prefix_mask[None].to(dev),text_labels=labels)
|
||||||
|
valid=s['rtc_valid_mask'][None].to(dev)[:,:,None];flow_loss=((o.vector_field-vel).square()*valid).sum()/valid.sum().clamp_min(1);loss=o.text_loss+cfg['train']['flow_loss_weight']*flow_loss;opt.zero_grad();loss.backward();torch.nn.utils.clip_grad_norm_(policy.parameters(),cfg['train']['grad_clip_norm']);opt.step()
|
||||||
|
if step==1 or step%cfg['train']['log_every']==0:row={'stage':'policy','step':step,'loss':float(loss.detach().cpu()),'text_loss':float(o.text_loss.detach().cpu()),'flow_loss':float(flow_loss.detach().cpu())};history.append(row);print(json.dumps(row))
|
||||||
|
ck={'policy':policy.state_dict(),'high_level':high.state_dict(),'goal':goal.state_dict(),'policy_config':pcfg.__dict__,'config':cfg,'action_statistics':ds.action_normalizer.as_dict()};torch.save(ck,out/'checkpoint.pt');write_json(out/'training_history.json',history);write_json(out/'summary.json',{'device':str(dev),'samples':len(ds),'policy_parameters':sum(p.numel() for p in policy.parameters()),'goal_parameters':sum(p.numel() for p in goal.parameters()),'high_level_parameters':sum(p.numel() for p in high.parameters())});return out/'checkpoint.pt'
|
||||||
|
|
||||||
|
def main():
|
||||||
|
p=argparse.ArgumentParser();p.add_argument('--config',default='configs/overfit.yaml');train(p.parse_args().config)
|
||||||
|
if __name__=='__main__':main()
|
||||||
@@ -0,0 +1,76 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from pi07mini.context import (
|
||||||
|
PolicyMetadata,
|
||||||
|
build_ambiguous_history_pair,
|
||||||
|
drop_prompt_components,
|
||||||
|
endpoint_subgoal_heatmap,
|
||||||
|
format_policy_prompt,
|
||||||
|
make_history_proxy_frames,
|
||||||
|
prompt_components,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _history() -> torch.Tensor:
|
||||||
|
x = torch.linspace(-8.0, 0.0, 8)
|
||||||
|
return torch.stack((x, 0.04 * x.square()), -1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_subgoal_heatmap_depends_on_endpoint_only() -> None:
|
||||||
|
endpoint = torch.tensor([20.0, 4.0])
|
||||||
|
first = endpoint_subgoal_heatmap(endpoint, size=24)
|
||||||
|
# Two future paths with radically different interiors produce exactly the
|
||||||
|
# same condition because only their shared endpoint is passed to the API.
|
||||||
|
path_a = torch.stack((torch.linspace(0.0, 20.0, 16), torch.linspace(0.0, 4.0, 16)), -1)
|
||||||
|
path_b = path_a.clone()
|
||||||
|
path_b[1:-1, 1] += 100.0
|
||||||
|
second = endpoint_subgoal_heatmap(path_b[-1], size=24)
|
||||||
|
torch.testing.assert_close(first, second, atol=0.0, rtol=0.0)
|
||||||
|
assert first.shape == (1, 24, 24)
|
||||||
|
assert float(first.max()) == 1.0
|
||||||
|
assert torch.isfinite(first).all()
|
||||||
|
|
||||||
|
|
||||||
|
def test_history_proxy_is_oldest_to_newest_and_future_free() -> None:
|
||||||
|
image = torch.linspace(0.0, 1.0, 3 * 20 * 20).reshape(3, 20, 20)
|
||||||
|
frames = make_history_proxy_frames(image, _history(), frames=6)
|
||||||
|
assert frames.shape == (6, 3, 20, 20)
|
||||||
|
torch.testing.assert_close(frames[-1], image * 0.98)
|
||||||
|
assert not torch.equal(frames[0], frames[-1])
|
||||||
|
|
||||||
|
|
||||||
|
def test_ambiguous_histories_are_distinct_but_share_current_state() -> None:
|
||||||
|
pair = build_ambiguous_history_pair(_history(), frames=6)
|
||||||
|
assert pair.shape == (2, 6, 2)
|
||||||
|
torch.testing.assert_close(pair[0, -1], pair[1, -1])
|
||||||
|
assert not torch.allclose(pair[0, :-1], pair[1, :-1])
|
||||||
|
|
||||||
|
|
||||||
|
def test_prompt_component_dropout_is_independent_and_deterministic() -> None:
|
||||||
|
metadata = PolicyMetadata("fast", "good", "none", "autonomous")
|
||||||
|
components = prompt_components(metadata, task="Drive safely.")
|
||||||
|
generator_a = torch.Generator().manual_seed(42)
|
||||||
|
generator_b = torch.Generator().manual_seed(42)
|
||||||
|
visible_a, mask_a = drop_prompt_components(
|
||||||
|
components, drop_probability=0.5, generator=generator_a
|
||||||
|
)
|
||||||
|
visible_b, mask_b = drop_prompt_components(
|
||||||
|
components, drop_probability=0.5, generator=generator_b
|
||||||
|
)
|
||||||
|
assert visible_a == visible_b
|
||||||
|
torch.testing.assert_close(mask_a, mask_b)
|
||||||
|
assert visible_a["task"] == "Drive safely."
|
||||||
|
assert mask_a.shape == (5,)
|
||||||
|
text = format_policy_prompt(visible_a)
|
||||||
|
assert "Drive safely." in text
|
||||||
|
for dropped in set(components) - set(visible_a):
|
||||||
|
assert f"{dropped}=" not in text
|
||||||
|
|
||||||
|
|
||||||
|
def test_full_metadata_dropout_keeps_only_task() -> None:
|
||||||
|
components = prompt_components(PolicyMetadata("slow", "poor", "stop", "autonomous"))
|
||||||
|
visible, mask = drop_prompt_components(components, drop_probability=1.0)
|
||||||
|
assert set(visible) == {"task"}
|
||||||
|
assert mask.tolist() == [True, False, False, False, False]
|
||||||
@@ -0,0 +1,134 @@
|
|||||||
|
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"])
|
||||||
@@ -0,0 +1,63 @@
|
|||||||
|
import torch
|
||||||
|
|
||||||
|
from pi07mini.memory import MEMConfig, MEMVideoEncoder
|
||||||
|
|
||||||
|
|
||||||
|
def tiny_encoder() -> MEMVideoEncoder:
|
||||||
|
return MEMVideoEncoder(
|
||||||
|
MEMConfig(
|
||||||
|
image_size=16,
|
||||||
|
patch_size=8,
|
||||||
|
d_model=16,
|
||||||
|
n_heads=4,
|
||||||
|
n_layers=2,
|
||||||
|
max_history_frames=6,
|
||||||
|
)
|
||||||
|
).eval()
|
||||||
|
|
||||||
|
|
||||||
|
def test_mem_output_token_count_is_independent_of_history_length() -> None:
|
||||||
|
encoder = tiny_encoder()
|
||||||
|
current = torch.randn(2, 3, 16, 16)
|
||||||
|
with torch.no_grad():
|
||||||
|
no_history = encoder(current)
|
||||||
|
long_history = encoder(current, torch.randn(2, 6, 3, 16, 16))
|
||||||
|
assert no_history.shape == (2, 4, 16)
|
||||||
|
assert long_history.shape == no_history.shape
|
||||||
|
|
||||||
|
|
||||||
|
def test_mem_is_sensitive_to_temporal_order() -> None:
|
||||||
|
torch.manual_seed(7)
|
||||||
|
encoder = tiny_encoder()
|
||||||
|
current = torch.randn(1, 3, 16, 16)
|
||||||
|
history = torch.stack(
|
||||||
|
(
|
||||||
|
torch.zeros(3, 16, 16),
|
||||||
|
torch.ones(3, 16, 16),
|
||||||
|
torch.full((3, 16, 16), 2.0),
|
||||||
|
)
|
||||||
|
)[None]
|
||||||
|
with torch.no_grad():
|
||||||
|
chronological = encoder(current, history)
|
||||||
|
reversed_history = encoder(current, history.flip(1))
|
||||||
|
assert not torch.allclose(chronological, reversed_history)
|
||||||
|
|
||||||
|
|
||||||
|
def test_masked_history_frame_cannot_change_memory() -> None:
|
||||||
|
torch.manual_seed(9)
|
||||||
|
encoder = tiny_encoder()
|
||||||
|
current = torch.randn(1, 3, 16, 16)
|
||||||
|
history = torch.randn(1, 2, 3, 16, 16)
|
||||||
|
changed = history.clone()
|
||||||
|
changed[:, 0] = 1000.0
|
||||||
|
mask = torch.tensor([[False, True]])
|
||||||
|
with torch.no_grad():
|
||||||
|
first = encoder(current, history, mask)
|
||||||
|
second = encoder(current, changed, mask)
|
||||||
|
torch.testing.assert_close(first, second)
|
||||||
|
|
||||||
|
|
||||||
|
def test_current_frame_temporal_position_is_exactly_zero() -> None:
|
||||||
|
encoding = MEMVideoEncoder.temporal_position(torch.tensor([3, 2, 1, 0]), 15)
|
||||||
|
assert torch.count_nonzero(encoding[-1]) == 0
|
||||||
|
|
||||||
@@ -0,0 +1,180 @@
|
|||||||
|
import torch
|
||||||
|
|
||||||
|
from pi07mini.model import (
|
||||||
|
HighLevelPolicy,
|
||||||
|
Pi07Config,
|
||||||
|
Pi07Mini,
|
||||||
|
PrefixLayout,
|
||||||
|
TinyGoalGenerator,
|
||||||
|
)
|
||||||
|
from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer
|
||||||
|
|
||||||
|
|
||||||
|
def tiny_model() -> Pi07Mini:
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
return Pi07Mini(
|
||||||
|
Pi07Config(
|
||||||
|
vocab_size=tokenizer.vocab_size,
|
||||||
|
image_size=16,
|
||||||
|
patch_size=16,
|
||||||
|
d_model=32,
|
||||||
|
n_heads=4,
|
||||||
|
n_layers=1,
|
||||||
|
memory_layers=1,
|
||||||
|
max_history_frames=3,
|
||||||
|
max_text_tokens=64,
|
||||||
|
max_action_tokens=4,
|
||||||
|
action_start_token_id=tokenizer.ACTION_START,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def sample_sequence():
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
metadata = EpisodeMetadata(2000, 5, False, "ee", "demonstration")
|
||||||
|
return tokenizer, tokenizer.build_policy_sequence(
|
||||||
|
"drive", "steer_left", torch.zeros(2, 2), metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def common_inputs(sequence):
|
||||||
|
return dict(
|
||||||
|
pixel_values=torch.randn(1, 3, 16, 16),
|
||||||
|
history_pixel_values=torch.randn(1, 2, 3, 16, 16),
|
||||||
|
history_states=torch.randn(1, 2, 2),
|
||||||
|
current_state=torch.randn(1, 2),
|
||||||
|
subgoal_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],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_block_causal_attention_matches_pi07_modal_order() -> None:
|
||||||
|
layout = PrefixLayout(observation_tokens=2, goal_tokens=2, text_tokens=3)
|
||||||
|
valid = torch.ones(1, layout.total_tokens, dtype=torch.bool)
|
||||||
|
edges = Pi07Mini.build_prefix_attention_mask(layout, valid)[0]
|
||||||
|
assert edges[0, :2].all()
|
||||||
|
assert not edges[0, 2:].any() # observation cannot read goal or text
|
||||||
|
assert edges[2, :4].all()
|
||||||
|
assert not edges[2, 4:].any() # goal reads observation + goal, not text
|
||||||
|
assert edges[4, :5].all()
|
||||||
|
assert not edges[4, 5:].any() # first text token is causal
|
||||||
|
assert edges[6].all()
|
||||||
|
|
||||||
|
|
||||||
|
def test_prefix_has_fixed_mem_tokens_and_one_token_per_state() -> None:
|
||||||
|
_, sequence = sample_sequence()
|
||||||
|
model = tiny_model()
|
||||||
|
inputs = common_inputs(sequence)
|
||||||
|
prefix, _, layout = model.encode_prefix(
|
||||||
|
inputs["pixel_values"],
|
||||||
|
inputs["history_pixel_values"],
|
||||||
|
inputs["history_states"],
|
||||||
|
inputs["current_state"],
|
||||||
|
inputs["input_ids"],
|
||||||
|
inputs["subgoal_pixel_values"],
|
||||||
|
)
|
||||||
|
assert layout.observation_tokens == 1 + 2 + 1 # MEM + history states + current
|
||||||
|
assert layout.goal_tokens == 1
|
||||||
|
assert prefix.shape[1] == layout.total_tokens
|
||||||
|
|
||||||
|
|
||||||
|
def test_flow_cannot_read_teacher_forced_fast_tokens() -> None:
|
||||||
|
torch.manual_seed(12)
|
||||||
|
tokenizer, sequence = sample_sequence()
|
||||||
|
model = tiny_model().eval()
|
||||||
|
changed_ids = sequence.input_ids.clone()
|
||||||
|
changed_ids[sequence.action_start_index + 1 : -1] = tokenizer.ACTION_TOKEN_START
|
||||||
|
inputs = common_inputs(sequence)
|
||||||
|
with torch.no_grad():
|
||||||
|
first = model(**inputs).vector_field
|
||||||
|
inputs["input_ids"] = changed_ids[None]
|
||||||
|
second = model(**inputs).vector_field
|
||||||
|
torch.testing.assert_close(first, second)
|
||||||
|
|
||||||
|
|
||||||
|
def test_flow_loss_is_knowledge_insulated_from_all_vlm_parameters() -> None:
|
||||||
|
torch.manual_seed(13)
|
||||||
|
_, sequence = sample_sequence()
|
||||||
|
model = tiny_model()
|
||||||
|
output = model(
|
||||||
|
**common_inputs(sequence), target_velocity=torch.randn(1, 2, 2)
|
||||||
|
)
|
||||||
|
assert output.flow_loss is not None
|
||||||
|
output.flow_loss.backward()
|
||||||
|
assert all(
|
||||||
|
parameter.grad is None or torch.count_nonzero(parameter.grad) == 0
|
||||||
|
for parameter in model.vlm_parameters()
|
||||||
|
)
|
||||||
|
assert any(
|
||||||
|
parameter.grad is not None and torch.count_nonzero(parameter.grad) > 0
|
||||||
|
for parameter in model.action_parameters()
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_fast_cross_entropy_still_trains_vlm() -> None:
|
||||||
|
_, sequence = sample_sequence()
|
||||||
|
model = tiny_model()
|
||||||
|
output = model(
|
||||||
|
**common_inputs(sequence), text_labels=sequence.labels[None]
|
||||||
|
)
|
||||||
|
assert output.text_loss is not None
|
||||||
|
output.text_loss.backward()
|
||||||
|
assert model.token_embedding.weight.grad is not None
|
||||||
|
assert torch.count_nonzero(model.token_embedding.weight.grad) > 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_visual_subgoal_conditions_flow_expert() -> None:
|
||||||
|
torch.manual_seed(14)
|
||||||
|
_, sequence = sample_sequence()
|
||||||
|
model = tiny_model().eval()
|
||||||
|
inputs = common_inputs(sequence)
|
||||||
|
with torch.no_grad():
|
||||||
|
first = model(**inputs).vector_field
|
||||||
|
inputs["subgoal_pixel_values"] = inputs["subgoal_pixel_values"] + 3.0
|
||||||
|
second = model(**inputs).vector_field
|
||||||
|
assert not torch.allclose(first, second)
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_cfg_uses_paper_extrapolation_formula() -> None:
|
||||||
|
conditioned = torch.tensor([1.0, 2.0])
|
||||||
|
unconditional = torch.tensor([0.5, 1.5])
|
||||||
|
torch.testing.assert_close(
|
||||||
|
Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 0.0), conditioned
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
Pi07Mini.combine_metadata_cfg(conditioned, unconditional, 1.0),
|
||||||
|
torch.tensor([1.5, 2.5]),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_high_level_policy_predicts_next_subtask_token() -> None:
|
||||||
|
model = HighLevelPolicy(tiny_model().config)
|
||||||
|
output = model(
|
||||||
|
pixel_values=torch.randn(2, 3, 16, 16),
|
||||||
|
history_pixel_values=torch.randn(2, 2, 3, 16, 16),
|
||||||
|
history_states=torch.randn(2, 2, 2),
|
||||||
|
current_state=torch.randn(2, 2),
|
||||||
|
input_ids=torch.tensor([[256, 257, 264], [256, 257, 264]]),
|
||||||
|
target_token=torch.tensor([271, 272]),
|
||||||
|
)
|
||||||
|
assert output.logits.shape == (2, tiny_model().config.vocab_size)
|
||||||
|
assert output.loss is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_tiny_goal_generator_predicts_image_flow() -> None:
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
generator = TinyGoalGenerator(tokenizer.vocab_size)
|
||||||
|
current = torch.randn(2, 3, 16, 16)
|
||||||
|
noisy = torch.randn_like(current)
|
||||||
|
output = generator(
|
||||||
|
current,
|
||||||
|
torch.tensor([[256, 257, 264], [256, 257, 264]]),
|
||||||
|
noisy,
|
||||||
|
torch.rand(2),
|
||||||
|
target_velocity=torch.randn_like(current),
|
||||||
|
)
|
||||||
|
assert output.vector_field.shape == current.shape
|
||||||
|
assert output.loss is not None
|
||||||
@@ -0,0 +1,50 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from pi07mini.rtc import align_actions_for_rtc, sample_rtc_delay
|
||||||
|
|
||||||
|
|
||||||
|
def _actions() -> torch.Tensor:
|
||||||
|
return torch.arange(32, dtype=torch.float32).reshape(16, 2)
|
||||||
|
|
||||||
|
|
||||||
|
def test_zero_delay_preserves_entire_chunk() -> None:
|
||||||
|
actions = _actions()
|
||||||
|
aligned = align_actions_for_rtc(actions, 0)
|
||||||
|
torch.testing.assert_close(aligned.delayed_target_actions, actions)
|
||||||
|
assert aligned.valid_mask.all()
|
||||||
|
assert aligned.executed_prefix_actions.shape == (0, 2)
|
||||||
|
|
||||||
|
|
||||||
|
def test_delay_shifts_target_and_masks_padded_tail() -> None:
|
||||||
|
actions = _actions()
|
||||||
|
aligned = align_actions_for_rtc(actions, 12)
|
||||||
|
torch.testing.assert_close(aligned.delayed_target_actions[:4], actions[12:])
|
||||||
|
assert torch.count_nonzero(aligned.delayed_target_actions[4:]) == 0
|
||||||
|
assert aligned.valid_mask.tolist() == [True] * 4 + [False] * 12
|
||||||
|
torch.testing.assert_close(aligned.executed_prefix_actions, actions[:12])
|
||||||
|
|
||||||
|
|
||||||
|
def test_rtc_alignment_preserves_batch_dimensions() -> None:
|
||||||
|
actions = torch.stack((_actions(), _actions() + 100.0))
|
||||||
|
aligned = align_actions_for_rtc(actions, 3)
|
||||||
|
assert aligned.delayed_target_actions.shape == (2, 16, 2)
|
||||||
|
assert aligned.valid_mask.shape == (2, 16)
|
||||||
|
torch.testing.assert_close(aligned.delayed_target_actions[:, :13], actions[:, 3:])
|
||||||
|
|
||||||
|
|
||||||
|
def test_delay_sampling_covers_closed_zero_to_twelve_range() -> None:
|
||||||
|
generator = torch.Generator().manual_seed(9)
|
||||||
|
values = {
|
||||||
|
sample_rtc_delay(generator=generator, action_horizon=16) for _ in range(300)
|
||||||
|
}
|
||||||
|
assert min(values) == 0
|
||||||
|
assert max(values) == 12
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("delay", [-1, 13, 16])
|
||||||
|
def test_invalid_delay_is_rejected(delay: int) -> None:
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
align_actions_for_rtc(_actions(), delay)
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
import torch
|
||||||
|
|
||||||
|
from pi07mini.tokenizer import EpisodeMetadata, MiniTokenizer
|
||||||
|
|
||||||
|
|
||||||
|
def metadata() -> EpisodeMetadata:
|
||||||
|
return EpisodeMetadata(
|
||||||
|
speed_steps=2100,
|
||||||
|
quality=5,
|
||||||
|
mistake=False,
|
||||||
|
control_mode="ee",
|
||||||
|
source_type="autonomous_rollout",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_prompt_contains_rich_metadata_and_control_tokens() -> None:
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
prompt, metadata_indices = tokenizer.encode_prompt(
|
||||||
|
"drive on the road", "steer_left", metadata()
|
||||||
|
)
|
||||||
|
assert tokenizer.SPEED_MARKER in prompt
|
||||||
|
assert tokenizer.QUALITY_MARKER in prompt
|
||||||
|
assert tokenizer.MISTAKE_MARKER in prompt
|
||||||
|
assert tokenizer.SOURCE_MARKER in prompt
|
||||||
|
assert tokenizer.CONTROL_MARKER in prompt
|
||||||
|
assert tokenizer.CONTROL_EE in prompt
|
||||||
|
assert len(metadata_indices) == 8
|
||||||
|
|
||||||
|
|
||||||
|
def test_metadata_cfg_pair_drops_metadata_but_retains_control_mode() -> None:
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
actions = torch.zeros(3, 2)
|
||||||
|
conditioned, unconditional = tokenizer.build_metadata_cfg_pair(
|
||||||
|
"drive", "keep_lane", actions, metadata()
|
||||||
|
)
|
||||||
|
assert tokenizer.QUALITY_MARKER in conditioned.input_ids
|
||||||
|
assert tokenizer.QUALITY_MARKER not in unconditional.input_ids
|
||||||
|
assert tokenizer.CONTROL_EE in conditioned.input_ids
|
||||||
|
assert tokenizer.CONTROL_EE in unconditional.input_ids
|
||||||
|
assert conditioned.metadata_indices
|
||||||
|
assert not unconditional.metadata_indices
|
||||||
|
|
||||||
|
|
||||||
|
def test_flow_prefix_ends_before_fast_stream() -> None:
|
||||||
|
tokenizer = MiniTokenizer()
|
||||||
|
sequence = tokenizer.build_policy_sequence(
|
||||||
|
"drive", "keep_lane", torch.zeros(4, 2), metadata()
|
||||||
|
)
|
||||||
|
assert sequence.input_ids[sequence.action_start_index] == tokenizer.ACTION_START
|
||||||
|
assert sequence.action_prefix_mask[: sequence.action_start_index].all()
|
||||||
|
assert not sequence.action_prefix_mask[sequence.action_start_index :].any()
|
||||||
|
assert sequence.labels[sequence.action_start_index - 1] == 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
|
||||||
|
|
||||||
Reference in New Issue
Block a user