Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 seeker
|
bc9d352a3a | ||
|
seeker
|
e951035d9b | ||
|
seeker
|
bd88a46b23 |
@@ -6,4 +6,5 @@ __pycache__/
|
||||
.venv/
|
||||
data/
|
||||
outputs/**/*.pt
|
||||
notebook_runs/
|
||||
.DS_Store
|
||||
|
||||
@@ -1,19 +1,12 @@
|
||||
# π0.6 / RECAP Minimal Validation
|
||||
# π0.6 / RECAP 最小复现
|
||||
|
||||
这是一个用于**小样本过拟合验证**的 π0.6 / `π*0.6` 风格项目。它把 π0.5
|
||||
已有的“高层 subtask + FAST 离散动作 + flow action expert”扩展为一个最小
|
||||
RECAP 闭环：收集示范、失败与人工纠正经验，训练分布式 value function，估计
|
||||
advantage，再用正/负 advantage 条件训练策略，并在推理时比较普通策略与
|
||||
classifier-free guidance（CFG）策略。
|
||||
本项目是 Physical Intelligence π0.6 / `π*0.6` RECAP 的最小复现。它在 π0.5 的高层 subtask、FAST 离散动作和 flow action expert 之上，加入示范、失败与人工纠正经验，完成 value function、advantage relabeling、条件策略训练和 CFG 推理的最小闭环。
|
||||
|
||||
它不是 Physical Intelligence 官方实现，也不能用于真实车辆或机器人控制。
|
||||
|
||||
## 公开资料边界
|
||||
## 复现范围
|
||||
|
||||
本项目以公开的 **“π*0.6: a VLA That Learns From Experience”** RECAP
|
||||
论文和 π0.6 model card 为主要依据。π0.6 没有一篇单独公开、包含全部训练细节的
|
||||
基础模型论文；截至本项目实现时，公开 `openpi` 仓库中也没有对应的官方
|
||||
π0.6/`π*0.6` 训练源码、权重或完整 RECAP 数据流水线。因此：
|
||||
论文和 π0.6 model card 为主要依据。公开资料没有提供完整的
|
||||
π0.6/`π*0.6` 训练源码、权重和 RECAP 数据流水线，未公开的部分在本项目中使用小型可运行结构替代：
|
||||
|
||||
- 论文明确给出的 RECAP 顺序、201-bin value distribution、二值 advantage
|
||||
条件、30% 条件 dropout 和 CFG 被保留；
|
||||
@@ -27,8 +20,8 @@ classifier-free guidance（CFG）策略。
|
||||
| 1. Base policy | `64×64` 图像、离散历史 XY、任务提示 | 小型 VLM policy + stop-gradient flow expert | subtask、FAST-lite tokens、连续 `16×2` action chunk | token CE + flow matching |
|
||||
| 2. Experience | 一条真实示范轨迹 | 确定性经验构造 | 成功示范、understeer、提前停止、人工纠正 4 类 episode | 无梯度，仅构造 buffer |
|
||||
| 3. Value | 图像、离散状态、任务、episode return | 小型独立 value VLM | `201` 个 return bins，期望 value 位于 `[-1,0]` | distributional CE |
|
||||
| 4. Advantage | reward、当前/未来 value、intervention mask | n-step estimator + quantile threshold | positive/negative indicator；人工纠正强制 positive | 无梯度，生成策略条件 |
|
||||
| 5. Mixed warm-up | 示范与合成经验、advantage 条件 | advantage-conditioned policy | 条件化离散动作和连续 flow | token CE + flow matching |
|
||||
| 4. Advantage | reward、学习到的当前/未来 value、intervention mask | critic inference + n-step estimator + quantile threshold | positive/negative indicator；人工纠正强制 positive | value 更新后重新标注策略数据 |
|
||||
| 5. Mixed SFT baseline | 示范与合成经验，不提供 advantage | unconditioned mixed policy | 离散动作和连续 flow | token CE + flow matching |
|
||||
| 6. RECAP extraction | 更新后的 value、重新标注的全部经验 | policy 从 base checkpoint 再训练 | 改进后的 `π*0.6`-style policy | 条件化 token CE + flow matching |
|
||||
| 7. CFG inference | 同一 observation 分别使用 positive 和 omitted 条件 | 两次 flow 前向 + CFG，默认 `β=2` | guided 连续 action chunk | 与真值及 unguided 输出比较 |
|
||||
|
||||
@@ -43,6 +36,11 @@ advantage token 位于 subtask 之后、动作之前，因此它只改变动作
|
||||
读取 teacher-forced FAST-lite 真值动作。flow loss 对 policy/VLM prefix 使用
|
||||
stop-gradient，以最小形式模拟论文的 Knowledge Insulation（KI）训练原则。
|
||||
|
||||
critic 只读取 observation/state（本单任务实验隐含固定任务提示），不读取
|
||||
`demo_success`、失败类型或 intervention source 等结果标签。value 训练完成后，
|
||||
程序会对每条 episode 的所有状态重新运行 critic，并用预测值重新计算 n-step
|
||||
advantage；最终 RECAP policy 使用该重标注结果，而不是手工 baseline。
|
||||
|
||||
## 与 π0.5 最小项目的重点差异
|
||||
|
||||
| π0.5 minimal | π0.6 / `π*0.6` minimal |
|
||||
@@ -66,7 +64,7 @@ width 为 96，value width 为 64；没有复现这些规模和数据改进。
|
||||
- `autonomous_stop`：中途停止；
|
||||
- `intervention_correction`：先偏离，再由“人工”轨迹拉回。
|
||||
|
||||
默认用 4-step bootstrapped advantage，`gamma=0.99`。训练集的 advantage
|
||||
默认用 4-step bootstrapped advantage，按论文使用无折扣 return（`gamma=1.0`）。训练集的 advantage
|
||||
第 70 百分位作为 positive 阈值，即约 30% 普通样本为 positive；intervention
|
||||
动作无条件标为 positive。这个合成 buffer 只用于检验 RECAP 接口，不能替代
|
||||
真实机器人 rollout、任务成功标签或在线人工干预。
|
||||
@@ -83,6 +81,16 @@ python -m pip install -U pip
|
||||
python -m pip install -e '.[dev]'
|
||||
```
|
||||
|
||||
完整的交互式复现入口见
|
||||
[`notebooks/pi06_recap_minimal_reproduction.ipynb`](notebooks/pi06_recap_minimal_reproduction.ipynb)。
|
||||
它包含论文链接、结构说明、环境安装、数据/checkpoint 本地导入或 URL 下载、
|
||||
测试、训练、评估与结果图展示。启动方式：
|
||||
|
||||
```bash
|
||||
python -m pip install -e '.[dev,notebook]'
|
||||
jupyter lab notebooks/pi06_recap_minimal_reproduction.ipynb
|
||||
```
|
||||
|
||||
默认配置读取：
|
||||
|
||||
```text
|
||||
@@ -120,7 +128,7 @@ python -m pi06mini.train --config configs/overfit.yaml
|
||||
|---|---:|---|
|
||||
| base | 300 | 先记住示范的 subtask、FAST-lite 和连续动作 |
|
||||
| value | 400 | 记住成功、失败和纠正 episode 的 201-bin return |
|
||||
| mixed | 300 | 让策略学会 positive/negative/omitted 三种条件 |
|
||||
| mixed | 300 | 训练不带 advantage 的全经验 SFT 对照组 |
|
||||
| recap | 700 | 使用合成经验中的 n-step advantage 标签与 30% 条件 dropout 提取最终策略 |
|
||||
|
||||
训练应保存 checkpoint、每阶段 loss、value/advantage 诊断和最终评估。CPU 可以
|
||||
@@ -148,6 +156,8 @@ python -m pi06mini.infer --checkpoint outputs/overfit/checkpoint.pt
|
||||
|
||||
在极小数据过拟合场景中，目标是看到 value 能区分合成成功/失败，positive 或
|
||||
CFG 策略比 negative/unguided 更接近专家轨迹。它不代表模型具备泛化能力。
|
||||
推理链会先由 VLM 预测 subtask token，再将该预测与 advantage indicator 一起
|
||||
交给连续 action expert，不使用数据集中的真值 subtask 作为低层条件。
|
||||
|
||||
## 关键配置
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ data:
|
||||
|
||||
experience:
|
||||
augmentation_factor: 1
|
||||
gamma: 0.99
|
||||
gamma: 1.0
|
||||
n_step: 4
|
||||
positive_quantile: 0.7
|
||||
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,88 @@
|
||||
{
|
||||
"predicted_values": [
|
||||
[
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.22784602642059326,
|
||||
-0.2076655775308609,
|
||||
-0.1778903752565384,
|
||||
-0.06647361069917679,
|
||||
-0.028201263397932053,
|
||||
-0.02679925411939621,
|
||||
-0.018520547077059746,
|
||||
-0.018748600035905838,
|
||||
-0.01845482364296913,
|
||||
-0.017860613763332367,
|
||||
-0.01756754145026207,
|
||||
-0.017646076157689095,
|
||||
-0.017747100442647934
|
||||
],
|
||||
[
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.1956789344549179,
|
||||
-0.1798555701971054,
|
||||
-0.18803948163986206,
|
||||
-0.1962352693080902,
|
||||
-0.177969291806221,
|
||||
-0.21014070510864258,
|
||||
-0.21469540894031525,
|
||||
-0.218982994556427,
|
||||
-0.170505091547966,
|
||||
-0.16970007121562958,
|
||||
-0.1965235322713852,
|
||||
-0.06712071597576141,
|
||||
-0.026878785341978073
|
||||
],
|
||||
[
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.22784602642059326,
|
||||
-0.2076655775308609,
|
||||
-0.18868279457092285,
|
||||
-0.3445330560207367,
|
||||
-0.33078309893608093,
|
||||
-0.3386183977127075,
|
||||
-0.3102336525917053,
|
||||
-0.29943764209747314,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745
|
||||
],
|
||||
[
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.1956789344549179,
|
||||
-0.1798555701971054,
|
||||
-0.16467693448066711,
|
||||
-0.09015780687332153,
|
||||
-0.07624924182891846,
|
||||
-0.07263262569904327,
|
||||
-0.016775300726294518,
|
||||
-0.01684822142124176,
|
||||
-0.017563480883836746,
|
||||
-0.01873198337852955,
|
||||
-0.019441569223999977,
|
||||
-0.019780147820711136,
|
||||
-0.0206752996891737
|
||||
]
|
||||
],
|
||||
"advantage_threshold": 0.02545459009706974,
|
||||
"episode_advantage_threshold": 0.046584345400333405,
|
||||
"episode_labels": [
|
||||
"positive",
|
||||
"negative",
|
||||
"negative",
|
||||
"positive"
|
||||
]
|
||||
}
|
||||
+164
-14
@@ -1,32 +1,182 @@
|
||||
{
|
||||
"predicted_subtask": "steer_left",
|
||||
"target_subtask": "steer_left",
|
||||
"trajectory_metrics": {
|
||||
"mixed": 100.86174011230469,
|
||||
"unconditional": 21.51112937927246,
|
||||
"positive": 0.061239250004291534,
|
||||
"negative": 55.823570251464844,
|
||||
"cfg": 0.24478793144226074,
|
||||
"negative_to_failure_mse": 48.71984100341797
|
||||
"mixed": 5.046791076660156,
|
||||
"unconditional": 330.43865966796875,
|
||||
"positive": 1.3081321716308594,
|
||||
"negative": 351.36962890625,
|
||||
"cfg": 0.06475023925304413,
|
||||
"negative_to_failure_mse": 20.35787582397461
|
||||
},
|
||||
"values": [
|
||||
{
|
||||
"source_type": "demo_success",
|
||||
"target": 0.0,
|
||||
"prediction": -0.007052943576127291
|
||||
"targets": [
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0
|
||||
],
|
||||
"predictions": [
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.22784602642059326,
|
||||
-0.2076655775308609,
|
||||
-0.1778903752565384,
|
||||
-0.06647361069917679,
|
||||
-0.028201263397932053,
|
||||
-0.02679925411939621,
|
||||
-0.018520547077059746,
|
||||
-0.018748600035905838,
|
||||
-0.01845482364296913,
|
||||
-0.017860613763332367,
|
||||
-0.01756754145026207,
|
||||
-0.017646076157689095,
|
||||
-0.017747100442647934
|
||||
],
|
||||
"curve_mse": 0.018043329939246178
|
||||
},
|
||||
{
|
||||
"source_type": "autonomous_understeer",
|
||||
"target": -0.03799404203891754,
|
||||
"prediction": -0.04714635759592056
|
||||
"targets": [
|
||||
-0.03993216156959534,
|
||||
-0.03993084654211998,
|
||||
-0.03992237150669098,
|
||||
-0.03984469920396805,
|
||||
-0.0396307148039341,
|
||||
-0.0392293743789196,
|
||||
-0.03840578347444534,
|
||||
-0.037170957773923874,
|
||||
-0.035448621958494186,
|
||||
-0.03320406377315521,
|
||||
-0.030476942658424377,
|
||||
-0.027057264000177383,
|
||||
-0.022992689162492752,
|
||||
-0.018253253772854805,
|
||||
-0.01285458542406559,
|
||||
-0.006825999356806278,
|
||||
0.0
|
||||
],
|
||||
"predictions": [
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.1956789344549179,
|
||||
-0.1798555701971054,
|
||||
-0.18803948163986206,
|
||||
-0.1962352693080902,
|
||||
-0.177969291806221,
|
||||
-0.21014070510864258,
|
||||
-0.21469540894031525,
|
||||
-0.218982994556427,
|
||||
-0.170505091547966,
|
||||
-0.16970007121562958,
|
||||
-0.1965235322713852,
|
||||
-0.06712071597576141,
|
||||
-0.026878785341978073
|
||||
],
|
||||
"curve_mse": 0.024060551077127457
|
||||
},
|
||||
{
|
||||
"source_type": "autonomous_stop",
|
||||
"target": -0.9509902000427246,
|
||||
"prediction": -0.9405871629714966
|
||||
"targets": [
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
-0.9818642139434814,
|
||||
-0.949521541595459,
|
||||
-0.9027397632598877,
|
||||
-0.8413366675376892,
|
||||
-0.7655149698257446,
|
||||
-0.6720305681228638,
|
||||
-0.5647895336151123,
|
||||
-0.444036602973938,
|
||||
-0.30988287925720215,
|
||||
-0.16282053291797638,
|
||||
0.0
|
||||
],
|
||||
"predictions": [
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.22784602642059326,
|
||||
-0.2076655775308609,
|
||||
-0.18868279457092285,
|
||||
-0.3445330560207367,
|
||||
-0.33078309893608093,
|
||||
-0.3386183977127075,
|
||||
-0.3102336525917053,
|
||||
-0.29943764209747314,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745,
|
||||
-0.2533053159713745
|
||||
],
|
||||
"curve_mse": 0.34516632556915283
|
||||
},
|
||||
{
|
||||
"source_type": "intervention_correction",
|
||||
"target": -0.0051981499418616295,
|
||||
"prediction": -0.0118478462100029
|
||||
"targets": [
|
||||
-0.005220414604991674,
|
||||
-0.005219103302806616,
|
||||
-0.005210627801716328,
|
||||
-0.0051329550333321095,
|
||||
-0.004918968304991722,
|
||||
-0.00456627644598484,
|
||||
-0.003942344803363085,
|
||||
-0.003156546037644148,
|
||||
-0.0022692831698805094,
|
||||
-0.0013850623508915305,
|
||||
-0.0006413022056221962,
|
||||
-0.00012316892389208078,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0,
|
||||
0.0
|
||||
],
|
||||
"predictions": [
|
||||
-0.20591823756694794,
|
||||
-0.20782312750816345,
|
||||
-0.20701773464679718,
|
||||
-0.2082844078540802,
|
||||
-0.1956789344549179,
|
||||
-0.1798555701971054,
|
||||
-0.16467693448066711,
|
||||
-0.09015780687332153,
|
||||
-0.07624924182891846,
|
||||
-0.07263262569904327,
|
||||
-0.016775300726294518,
|
||||
-0.01684822142124176,
|
||||
-0.017563480883836746,
|
||||
-0.01873198337852955,
|
||||
-0.019441569223999977,
|
||||
-0.019780147820711136,
|
||||
-0.0206752996891737
|
||||
],
|
||||
"curve_mse": 0.016281794756650925
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 87 KiB After Width: | Height: | Size: 84 KiB |
@@ -2,624 +2,624 @@
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 1,
|
||||
"loss": 72.41638946533203,
|
||||
"text_loss": 68.23114013671875,
|
||||
"flow_loss": 2.0926246643066406
|
||||
"loss": 70.63429260253906,
|
||||
"text_loss": 68.23114776611328,
|
||||
"flow_loss": 1.2015743255615234
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 20,
|
||||
"loss": 4.305421829223633,
|
||||
"text_loss": 1.9023208618164062,
|
||||
"flow_loss": 1.2015506029129028
|
||||
"loss": 3.8942408561706543,
|
||||
"text_loss": 1.7884641885757446,
|
||||
"flow_loss": 1.0528883934020996
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 40,
|
||||
"loss": 2.912196159362793,
|
||||
"text_loss": 0.8887173533439636,
|
||||
"flow_loss": 1.0117393732070923
|
||||
"loss": 3.1883139610290527,
|
||||
"text_loss": 0.8112223148345947,
|
||||
"flow_loss": 1.188545823097229
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 60,
|
||||
"loss": 2.184612512588501,
|
||||
"text_loss": 0.5364341735839844,
|
||||
"flow_loss": 0.8240891695022583
|
||||
"loss": 2.8662467002868652,
|
||||
"text_loss": 0.5043408274650574,
|
||||
"flow_loss": 1.1809529066085815
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 80,
|
||||
"loss": 1.578565239906311,
|
||||
"text_loss": 0.23946531116962433,
|
||||
"flow_loss": 0.6695499420166016
|
||||
"loss": 2.9415433406829834,
|
||||
"text_loss": 0.2767218053340912,
|
||||
"flow_loss": 1.3324108123779297
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 100,
|
||||
"loss": 1.5856558084487915,
|
||||
"text_loss": 0.04227600246667862,
|
||||
"flow_loss": 0.7716898918151855
|
||||
"loss": 1.23435640335083,
|
||||
"text_loss": 0.013344191014766693,
|
||||
"flow_loss": 0.6105061173439026
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 120,
|
||||
"loss": 0.3423895239830017,
|
||||
"text_loss": 0.0004414900322444737,
|
||||
"flow_loss": 0.1709740161895752
|
||||
"loss": 0.6844875812530518,
|
||||
"text_loss": 0.1804233342409134,
|
||||
"flow_loss": 0.2520321309566498
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 140,
|
||||
"loss": 0.24176472425460815,
|
||||
"text_loss": 1.8706961782299913e-05,
|
||||
"flow_loss": 0.12087301164865494
|
||||
"loss": 1.1563541889190674,
|
||||
"text_loss": 6.113824201747775e-05,
|
||||
"flow_loss": 0.5781465172767639
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 160,
|
||||
"loss": 0.3000774383544922,
|
||||
"text_loss": 7.637506496394053e-05,
|
||||
"flow_loss": 0.15000052750110626
|
||||
"loss": 0.49136143922805786,
|
||||
"text_loss": 0.00014380548964254558,
|
||||
"flow_loss": 0.24560882151126862
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 180,
|
||||
"loss": 0.6410971879959106,
|
||||
"text_loss": 1.6521895304322243e-05,
|
||||
"flow_loss": 0.32054033875465393
|
||||
"loss": 0.5463516712188721,
|
||||
"text_loss": 1.483415599068394e-05,
|
||||
"flow_loss": 0.2731684148311615
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 200,
|
||||
"loss": 0.6132829189300537,
|
||||
"text_loss": 9.676185072748922e-06,
|
||||
"flow_loss": 0.30663663148880005
|
||||
"loss": 0.21821065247058868,
|
||||
"text_loss": 4.533280844043475e-06,
|
||||
"flow_loss": 0.10910306125879288
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 220,
|
||||
"loss": 0.3431224524974823,
|
||||
"text_loss": 7.356796231761109e-06,
|
||||
"flow_loss": 0.17155754566192627
|
||||
"loss": 0.2135496884584427,
|
||||
"text_loss": 3.412749947528937e-06,
|
||||
"flow_loss": 0.10677313804626465
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 240,
|
||||
"loss": 0.28047603368759155,
|
||||
"text_loss": 5.987634267512476e-06,
|
||||
"flow_loss": 0.1402350217103958
|
||||
"loss": 0.15570451319217682,
|
||||
"text_loss": 3.027882712558494e-06,
|
||||
"flow_loss": 0.07785074412822723
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 260,
|
||||
"loss": 0.13324393332004547,
|
||||
"text_loss": 4.9011541705112904e-06,
|
||||
"flow_loss": 0.06661951541900635
|
||||
"loss": 0.17475822567939758,
|
||||
"text_loss": 2.816716460074531e-06,
|
||||
"flow_loss": 0.08737770467996597
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 280,
|
||||
"loss": 0.1635851263999939,
|
||||
"text_loss": 4.0462705328536686e-06,
|
||||
"flow_loss": 0.08179053664207458
|
||||
"loss": 0.12843455374240875,
|
||||
"text_loss": 2.6055499802168924e-06,
|
||||
"flow_loss": 0.06421597301959991
|
||||
},
|
||||
{
|
||||
"stage": "base",
|
||||
"step": 300,
|
||||
"loss": 0.15081387758255005,
|
||||
"text_loss": 3.4570459774840856e-06,
|
||||
"flow_loss": 0.07540521025657654
|
||||
"loss": 0.3167646825313568,
|
||||
"text_loss": 2.387570930295624e-06,
|
||||
"flow_loss": 0.1583811491727829
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 1,
|
||||
"loss": 5.427414894104004,
|
||||
"prediction": -0.5049723982810974,
|
||||
"loss": 5.466859817504883,
|
||||
"prediction": -0.5042269229888916,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 20,
|
||||
"loss": 2.24937105178833,
|
||||
"prediction": -0.3925527334213257,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 3.274851083755493,
|
||||
"prediction": -0.3775767683982849,
|
||||
"target": -0.005210627801716328
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 40,
|
||||
"loss": 1.728512167930603,
|
||||
"prediction": -0.32079368829727173,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 2.3394997119903564,
|
||||
"prediction": -0.2888708710670471,
|
||||
"target": -0.00456627644598484
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 60,
|
||||
"loss": 1.4823004007339478,
|
||||
"prediction": -0.3057250678539276,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.9264469146728516,
|
||||
"prediction": -0.2154604196548462,
|
||||
"target": -0.0022692831698805094
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 80,
|
||||
"loss": 1.1117331981658936,
|
||||
"prediction": -0.3553512394428253,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.756015956401825,
|
||||
"prediction": -0.19315782189369202,
|
||||
"target": -0.00012316892389208078
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 100,
|
||||
"loss": 0.2580283582210541,
|
||||
"prediction": -0.11316870898008347,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.6257156729698181,
|
||||
"prediction": -0.18343542516231537,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 120,
|
||||
"loss": 0.12432077527046204,
|
||||
"prediction": -0.054820768535137177,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.3138372898101807,
|
||||
"prediction": -0.37475356459617615,
|
||||
"target": -0.005220414604991674
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 140,
|
||||
"loss": 0.08916354924440384,
|
||||
"prediction": -0.04263461381196976,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.4211702346801758,
|
||||
"prediction": -0.1804044544696808,
|
||||
"target": -0.0051329550333321095
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 160,
|
||||
"loss": 0.07165656983852386,
|
||||
"prediction": -0.034443676471710205,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.448256254196167,
|
||||
"prediction": -0.2506459355354309,
|
||||
"target": -0.003942344803363085
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 180,
|
||||
"loss": 0.05974685773253441,
|
||||
"prediction": -0.029927225783467293,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.33143094182014465,
|
||||
"prediction": -0.08220178633928299,
|
||||
"target": -0.0013850623508915305
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 200,
|
||||
"loss": 0.05075548216700554,
|
||||
"prediction": -0.026268303394317627,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.2722547948360443,
|
||||
"prediction": -0.06349469721317291,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 220,
|
||||
"loss": 0.043695785105228424,
|
||||
"prediction": -0.023351771757006645,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.22779540717601776,
|
||||
"prediction": -0.05417971312999725,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 240,
|
||||
"loss": 0.03802856057882309,
|
||||
"prediction": -0.021020548418164253,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.5308982133865356,
|
||||
"prediction": -0.3449951708316803,
|
||||
"target": -0.005219103302806616
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 260,
|
||||
"loss": 0.03340449184179306,
|
||||
"prediction": -0.019105376675724983,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.1563409566879272,
|
||||
"prediction": -0.2026708871126175,
|
||||
"target": -0.004918968304991722
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 280,
|
||||
"loss": 0.02958041802048683,
|
||||
"prediction": -0.017514187842607498,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.2716314792633057,
|
||||
"prediction": -0.07276482880115509,
|
||||
"target": -0.003156546037644148
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 300,
|
||||
"loss": 0.02638067863881588,
|
||||
"prediction": -0.016178548336029053,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.3126322329044342,
|
||||
"prediction": -0.07691694051027298,
|
||||
"target": -0.0006413022056221962
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 320,
|
||||
"loss": 0.023675711825489998,
|
||||
"prediction": -0.01504577323794365,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.1044526994228363,
|
||||
"prediction": -0.03439018875360489,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 340,
|
||||
"loss": 0.021368399262428284,
|
||||
"prediction": -0.014076616615056992,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.1302744299173355,
|
||||
"prediction": -0.03440919518470764,
|
||||
"target": 0.0
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 360,
|
||||
"loss": 0.019383536651730537,
|
||||
"prediction": -0.013240991160273552,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.5437527894973755,
|
||||
"prediction": -0.24599501490592957,
|
||||
"target": -0.005210627801716328
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 380,
|
||||
"loss": 0.017663555219769478,
|
||||
"prediction": -0.01251534465700388,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 1.4080818891525269,
|
||||
"prediction": -0.3180582821369171,
|
||||
"target": -0.00456627644598484
|
||||
},
|
||||
{
|
||||
"stage": "value",
|
||||
"step": 400,
|
||||
"loss": 0.016163267195224762,
|
||||
"prediction": -0.011881118640303612,
|
||||
"target": -0.0051981499418616295
|
||||
"loss": 0.8151843547821045,
|
||||
"prediction": -0.07982967793941498,
|
||||
"target": -0.0022692831698805094
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 1,
|
||||
"loss": 0.11788154393434525,
|
||||
"text_loss": 3.4366105410299497e-06,
|
||||
"flow_loss": 0.05893905460834503
|
||||
"loss": 0.05100901424884796,
|
||||
"text_loss": 2.3773534394422313e-06,
|
||||
"flow_loss": 0.02550331875681877
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 20,
|
||||
"loss": 4.7238688468933105,
|
||||
"text_loss": 2.6195578575134277,
|
||||
"flow_loss": 1.0521554946899414
|
||||
"loss": 3.583022117614746,
|
||||
"text_loss": 2.2782084941864014,
|
||||
"flow_loss": 0.6524068117141724
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 40,
|
||||
"loss": 5.051268577575684,
|
||||
"text_loss": 1.2949998378753662,
|
||||
"flow_loss": 1.8781343698501587
|
||||
"loss": 2.2450661659240723,
|
||||
"text_loss": 1.1263315677642822,
|
||||
"flow_loss": 0.5593672394752502
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 60,
|
||||
"loss": 1.6957576274871826,
|
||||
"text_loss": 0.8319522142410278,
|
||||
"flow_loss": 0.4319027364253998
|
||||
"loss": 2.665435552597046,
|
||||
"text_loss": 0.9095665812492371,
|
||||
"flow_loss": 0.8779345154762268
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 80,
|
||||
"loss": 2.1900432109832764,
|
||||
"text_loss": 0.8061568140983582,
|
||||
"flow_loss": 0.6919432282447815
|
||||
"loss": 0.9744815826416016,
|
||||
"text_loss": 0.7057558298110962,
|
||||
"flow_loss": 0.1343628615140915
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 100,
|
||||
"loss": 2.2935876846313477,
|
||||
"text_loss": 0.5968648791313171,
|
||||
"flow_loss": 0.8483614325523376
|
||||
"loss": 2.257481575012207,
|
||||
"text_loss": 0.7136633992195129,
|
||||
"flow_loss": 0.7719091176986694
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 120,
|
||||
"loss": 1.4990849494934082,
|
||||
"text_loss": 0.6506502032279968,
|
||||
"flow_loss": 0.4242174029350281
|
||||
"loss": 1.557675838470459,
|
||||
"text_loss": 0.6338224411010742,
|
||||
"flow_loss": 0.46192666888237
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 140,
|
||||
"loss": 1.6056029796600342,
|
||||
"text_loss": 0.5977785587310791,
|
||||
"flow_loss": 0.5039122104644775
|
||||
"loss": 1.2216718196868896,
|
||||
"text_loss": 0.7448893785476685,
|
||||
"flow_loss": 0.2383912056684494
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 160,
|
||||
"loss": 1.3810169696807861,
|
||||
"text_loss": 0.48245665431022644,
|
||||
"flow_loss": 0.44928014278411865
|
||||
"loss": 0.5911489725112915,
|
||||
"text_loss": 0.27954256534576416,
|
||||
"flow_loss": 0.15580320358276367
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 180,
|
||||
"loss": 1.1724852323532104,
|
||||
"text_loss": 0.3958587944507599,
|
||||
"flow_loss": 0.3883132338523865
|
||||
"loss": 2.7055318355560303,
|
||||
"text_loss": 0.25718608498573303,
|
||||
"flow_loss": 1.224172830581665
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 200,
|
||||
"loss": 1.2829020023345947,
|
||||
"text_loss": 0.18238645792007446,
|
||||
"flow_loss": 0.5502578020095825
|
||||
"loss": 1.3802378177642822,
|
||||
"text_loss": 0.32982689142227173,
|
||||
"flow_loss": 0.5252054929733276
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 220,
|
||||
"loss": 0.7328482866287231,
|
||||
"text_loss": 0.1418272703886032,
|
||||
"flow_loss": 0.29551050066947937
|
||||
"loss": 1.1629283428192139,
|
||||
"text_loss": 0.26602569222450256,
|
||||
"flow_loss": 0.44845134019851685
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 240,
|
||||
"loss": 1.196341872215271,
|
||||
"text_loss": 0.28424081206321716,
|
||||
"flow_loss": 0.4560505449771881
|
||||
"loss": 0.5630282759666443,
|
||||
"text_loss": 0.2572731077671051,
|
||||
"flow_loss": 0.1528775840997696
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 260,
|
||||
"loss": 4.298240661621094,
|
||||
"text_loss": 0.257877379655838,
|
||||
"flow_loss": 2.020181655883789
|
||||
"loss": 1.1026716232299805,
|
||||
"text_loss": 0.15438926219940186,
|
||||
"flow_loss": 0.4741412103176117
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 280,
|
||||
"loss": 1.3184027671813965,
|
||||
"text_loss": 0.07380203157663345,
|
||||
"flow_loss": 0.622300386428833
|
||||
"loss": 1.9438565969467163,
|
||||
"text_loss": 0.18404552340507507,
|
||||
"flow_loss": 0.8799055218696594
|
||||
},
|
||||
{
|
||||
"stage": "mixed",
|
||||
"step": 300,
|
||||
"loss": 2.4886393547058105,
|
||||
"text_loss": 0.30716001987457275,
|
||||
"flow_loss": 1.0907396078109741
|
||||
"loss": 0.8524055480957031,
|
||||
"text_loss": 0.1450427919626236,
|
||||
"flow_loss": 0.35368138551712036
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 1,
|
||||
"loss": 10.11633586883545,
|
||||
"text_loss": 9.866013526916504,
|
||||
"flow_loss": 0.12516121566295624
|
||||
"loss": 10.587621688842773,
|
||||
"text_loss": 9.856688499450684,
|
||||
"flow_loss": 0.3654666244983673
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 20,
|
||||
"loss": 3.1665196418762207,
|
||||
"text_loss": 2.23919415473938,
|
||||
"flow_loss": 0.4636627733707428
|
||||
"loss": 3.289060354232788,
|
||||
"text_loss": 2.240044355392456,
|
||||
"flow_loss": 0.524507999420166
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 40,
|
||||
"loss": 3.1297008991241455,
|
||||
"text_loss": 1.4351346492767334,
|
||||
"flow_loss": 0.847283124923706
|
||||
"loss": 2.791560173034668,
|
||||
"text_loss": 1.4383518695831299,
|
||||
"flow_loss": 0.676604151725769
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 60,
|
||||
"loss": 2.608584403991699,
|
||||
"text_loss": 1.171854853630066,
|
||||
"flow_loss": 0.7183647751808167
|
||||
"loss": 1.8424029350280762,
|
||||
"text_loss": 1.1395618915557861,
|
||||
"flow_loss": 0.351420521736145
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 80,
|
||||
"loss": 3.5118634700775146,
|
||||
"text_loss": 0.9280703067779541,
|
||||
"flow_loss": 1.2918965816497803
|
||||
"loss": 1.8679100275039673,
|
||||
"text_loss": 0.9174721837043762,
|
||||
"flow_loss": 0.47521892189979553
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 100,
|
||||
"loss": 2.7196712493896484,
|
||||
"text_loss": 0.8991937041282654,
|
||||
"flow_loss": 0.9102387428283691
|
||||
"loss": 2.929778814315796,
|
||||
"text_loss": 0.9240415096282959,
|
||||
"flow_loss": 1.00286865234375
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 120,
|
||||
"loss": 1.1335762739181519,
|
||||
"text_loss": 0.7848832011222839,
|
||||
"flow_loss": 0.17434655129909515
|
||||
"loss": 1.4325169324874878,
|
||||
"text_loss": 0.7645139694213867,
|
||||
"flow_loss": 0.33400148153305054
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 140,
|
||||
"loss": 2.5729737281799316,
|
||||
"text_loss": 0.5885628461837769,
|
||||
"flow_loss": 0.9922055006027222
|
||||
"loss": 1.4750593900680542,
|
||||
"text_loss": 0.7536309957504272,
|
||||
"flow_loss": 0.3607141971588135
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 160,
|
||||
"loss": 1.9800522327423096,
|
||||
"text_loss": 0.6485287547111511,
|
||||
"flow_loss": 0.6657617688179016
|
||||
"loss": 1.5970324277877808,
|
||||
"text_loss": 1.017958164215088,
|
||||
"flow_loss": 0.28953713178634644
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 180,
|
||||
"loss": 1.5937684774398804,
|
||||
"text_loss": 0.6764931678771973,
|
||||
"flow_loss": 0.45863765478134155
|
||||
"loss": 1.6522973775863647,
|
||||
"text_loss": 0.922364354133606,
|
||||
"flow_loss": 0.3649665117263794
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 200,
|
||||
"loss": 1.588905692100525,
|
||||
"text_loss": 1.1673603057861328,
|
||||
"flow_loss": 0.21077267825603485
|
||||
"loss": 1.3916192054748535,
|
||||
"text_loss": 0.9831637740135193,
|
||||
"flow_loss": 0.2042277306318283
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 220,
|
||||
"loss": 1.1653995513916016,
|
||||
"text_loss": 0.5840165019035339,
|
||||
"flow_loss": 0.2906915545463562
|
||||
"loss": 1.138975739479065,
|
||||
"text_loss": 0.8730899691581726,
|
||||
"flow_loss": 0.13294290006160736
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 240,
|
||||
"loss": 2.6636130809783936,
|
||||
"text_loss": 1.3158458471298218,
|
||||
"flow_loss": 0.6738836169242859
|
||||
"loss": 1.502185344696045,
|
||||
"text_loss": 0.8015594482421875,
|
||||
"flow_loss": 0.3503129780292511
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 260,
|
||||
"loss": 0.5980206727981567,
|
||||
"text_loss": 0.42834123969078064,
|
||||
"flow_loss": 0.08483970910310745
|
||||
"loss": 0.8756294846534729,
|
||||
"text_loss": 0.31519001722335815,
|
||||
"flow_loss": 0.2802197337150574
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 280,
|
||||
"loss": 1.0666983127593994,
|
||||
"text_loss": 0.47541406750679016,
|
||||
"flow_loss": 0.2956421375274658
|
||||
"loss": 0.779049277305603,
|
||||
"text_loss": 0.3690524101257324,
|
||||
"flow_loss": 0.2049984186887741
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 300,
|
||||
"loss": 0.8535268306732178,
|
||||
"text_loss": 0.4074874520301819,
|
||||
"flow_loss": 0.22301967442035675
|
||||
"loss": 0.7172976732254028,
|
||||
"text_loss": 0.23720531165599823,
|
||||
"flow_loss": 0.2400461882352829
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 320,
|
||||
"loss": 5.488016605377197,
|
||||
"text_loss": 0.851451575756073,
|
||||
"flow_loss": 2.3182826042175293
|
||||
"loss": 2.9586782455444336,
|
||||
"text_loss": 0.7825477719306946,
|
||||
"flow_loss": 1.088065266609192
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 340,
|
||||
"loss": 0.7862600088119507,
|
||||
"text_loss": 0.37299302220344543,
|
||||
"flow_loss": 0.20663349330425262
|
||||
"loss": 1.745682954788208,
|
||||
"text_loss": 0.43472984433174133,
|
||||
"flow_loss": 0.6554765701293945
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 360,
|
||||
"loss": 1.3418667316436768,
|
||||
"text_loss": 0.7838777899742126,
|
||||
"flow_loss": 0.27899447083473206
|
||||
"loss": 1.5331687927246094,
|
||||
"text_loss": 0.7225565314292908,
|
||||
"flow_loss": 0.4053061604499817
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 380,
|
||||
"loss": 0.7067349553108215,
|
||||
"text_loss": 0.15938280522823334,
|
||||
"flow_loss": 0.2736760675907135
|
||||
"loss": 1.6411525011062622,
|
||||
"text_loss": 0.14736497402191162,
|
||||
"flow_loss": 0.7468937635421753
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 400,
|
||||
"loss": 0.4985463619232178,
|
||||
"text_loss": 0.2566893398761749,
|
||||
"flow_loss": 0.12092851102352142
|
||||
"loss": 0.5169777870178223,
|
||||
"text_loss": 0.24374954402446747,
|
||||
"flow_loss": 0.1366141140460968
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 420,
|
||||
"loss": 0.5810034871101379,
|
||||
"text_loss": 0.24579037725925446,
|
||||
"flow_loss": 0.16760656237602234
|
||||
"loss": 0.6074326038360596,
|
||||
"text_loss": 0.11582576483488083,
|
||||
"flow_loss": 0.24580343067646027
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 440,
|
||||
"loss": 0.40272825956344604,
|
||||
"text_loss": 0.19410930573940277,
|
||||
"flow_loss": 0.10430946946144104
|
||||
"loss": 0.7017235159873962,
|
||||
"text_loss": 0.3753448724746704,
|
||||
"flow_loss": 0.16318932175636292
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 460,
|
||||
"loss": 0.7170483469963074,
|
||||
"text_loss": 0.32609468698501587,
|
||||
"flow_loss": 0.19547683000564575
|
||||
"loss": 1.029306411743164,
|
||||
"text_loss": 0.2204148769378662,
|
||||
"flow_loss": 0.40444573760032654
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 480,
|
||||
"loss": 0.3233756422996521,
|
||||
"text_loss": 0.08512496948242188,
|
||||
"flow_loss": 0.11912534385919571
|
||||
"loss": 0.9227437376976013,
|
||||
"text_loss": 0.21806900203227997,
|
||||
"flow_loss": 0.3523373603820801
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 500,
|
||||
"loss": 0.625529408454895,
|
||||
"text_loss": 0.23650097846984863,
|
||||
"flow_loss": 0.1945142149925232
|
||||
"loss": 0.7456694841384888,
|
||||
"text_loss": 0.3062073886394501,
|
||||
"flow_loss": 0.21973104774951935
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 520,
|
||||
"loss": 1.645598292350769,
|
||||
"text_loss": 0.874957263469696,
|
||||
"flow_loss": 0.3853205144405365
|
||||
"loss": 1.2862324714660645,
|
||||
"text_loss": 0.6478312015533447,
|
||||
"flow_loss": 0.3192006051540375
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 540,
|
||||
"loss": 0.8086838722229004,
|
||||
"text_loss": 0.31192541122436523,
|
||||
"flow_loss": 0.24837924540042877
|
||||
"loss": 0.7051770687103271,
|
||||
"text_loss": 0.23054330050945282,
|
||||
"flow_loss": 0.23731689155101776
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 560,
|
||||
"loss": 1.9226100444793701,
|
||||
"text_loss": 0.8500325083732605,
|
||||
"flow_loss": 0.5362887382507324
|
||||
"loss": 1.9938247203826904,
|
||||
"text_loss": 1.1214075088500977,
|
||||
"flow_loss": 0.4362086057662964
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 580,
|
||||
"loss": 2.138998031616211,
|
||||
"text_loss": 0.7942134141921997,
|
||||
"flow_loss": 0.6723923683166504
|
||||
"loss": 1.6423680782318115,
|
||||
"text_loss": 1.1029378175735474,
|
||||
"flow_loss": 0.2697151303291321
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 600,
|
||||
"loss": 0.504305362701416,
|
||||
"text_loss": 0.13201656937599182,
|
||||
"flow_loss": 0.1861443817615509
|
||||
"loss": 1.3713027238845825,
|
||||
"text_loss": 0.1559465080499649,
|
||||
"flow_loss": 0.6076781153678894
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 620,
|
||||
"loss": 0.5794067978858948,
|
||||
"text_loss": 0.26564836502075195,
|
||||
"flow_loss": 0.1568792164325714
|
||||
"loss": 0.3619869351387024,
|
||||
"text_loss": 0.06553713977336884,
|
||||
"flow_loss": 0.14822489023208618
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 640,
|
||||
"loss": 0.8541183471679688,
|
||||
"text_loss": 0.154631569981575,
|
||||
"flow_loss": 0.34974339604377747
|
||||
"loss": 0.5116170048713684,
|
||||
"text_loss": 0.07361899316310883,
|
||||
"flow_loss": 0.2189989984035492
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 660,
|
||||
"loss": 0.16945119202136993,
|
||||
"text_loss": 0.056921038776636124,
|
||||
"flow_loss": 0.056265074759721756
|
||||
"loss": 0.8094859719276428,
|
||||
"text_loss": 0.07687854766845703,
|
||||
"flow_loss": 0.3663037121295929
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 680,
|
||||
"loss": 0.47458726167678833,
|
||||
"text_loss": 0.04364287480711937,
|
||||
"flow_loss": 0.21547219157218933
|
||||
"loss": 1.4220974445343018,
|
||||
"text_loss": 0.13382284343242645,
|
||||
"flow_loss": 0.6441373229026794
|
||||
},
|
||||
{
|
||||
"stage": "recap",
|
||||
"step": 700,
|
||||
"loss": 0.18534913659095764,
|
||||
"text_loss": 0.05316398665308952,
|
||||
"flow_loss": 0.06609257310628891
|
||||
"loss": 0.40469539165496826,
|
||||
"text_loss": 0.04384676367044449,
|
||||
"flow_loss": 0.18042431771755219
|
||||
}
|
||||
]
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 46 KiB After Width: | Height: | Size: 112 KiB |
@@ -22,6 +22,11 @@ dependencies = [
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
]
|
||||
notebook = [
|
||||
"ipykernel>=6.29",
|
||||
"jupyterlab>=4.2",
|
||||
"nbclient>=0.10",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
pi06mini-train = "pi06mini.train:main"
|
||||
|
||||
+66
-4
@@ -214,7 +214,7 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
history_steps: int = 8,
|
||||
action_horizon: int = 16,
|
||||
n_step: int = 4,
|
||||
gamma: float = 0.99,
|
||||
gamma: float = 1.0,
|
||||
positive_quantile: float = 0.7,
|
||||
seed: int = 7,
|
||||
action_statistics: ActionStatistics | dict[str, Sequence[float]] | None = None,
|
||||
@@ -229,6 +229,9 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
self.seed = int(seed)
|
||||
self.state_dim = self.history_steps * 2
|
||||
self.action_dim = 2
|
||||
self.n_step = int(n_step)
|
||||
self.gamma = float(gamma)
|
||||
self.positive_quantile = float(positive_quantile)
|
||||
|
||||
episodes: list[ExperienceEpisode] = []
|
||||
episode_record_indices: list[int] = []
|
||||
@@ -275,6 +278,12 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
)
|
||||
)
|
||||
|
||||
self.value_histories = torch.stack(
|
||||
[self._episode_value_histories(episode) for episode in self.episodes]
|
||||
)
|
||||
terminal = torch.zeros(len(self.episodes), 1, dtype=self.normalized_returns.dtype)
|
||||
self.value_targets = torch.cat((self.normalized_returns, terminal), dim=1)
|
||||
|
||||
if action_statistics is None:
|
||||
self.action_normalizer = ActionNormalizer.fit(
|
||||
[episode.history_xy for episode in episodes],
|
||||
@@ -287,6 +296,58 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
else:
|
||||
self.action_normalizer = ActionNormalizer.from_dict(action_statistics)
|
||||
|
||||
def _episode_value_histories(self, episode: ExperienceEpisode) -> torch.Tensor:
|
||||
"""Create observed state windows for every step and the terminal state."""
|
||||
|
||||
path = torch.cat((episode.history_xy, episode.future_xy), dim=0)
|
||||
windows: list[torch.Tensor] = []
|
||||
history_end = len(episode.history_xy)
|
||||
for step in range(len(episode.future_xy) + 1):
|
||||
end = history_end + step
|
||||
start = max(0, end - self.history_steps)
|
||||
window = path[start:end]
|
||||
if len(window) < self.history_steps:
|
||||
pad = window[:1].expand(self.history_steps - len(window), -1)
|
||||
window = torch.cat((pad, window), dim=0)
|
||||
windows.append(window[-self.history_steps :])
|
||||
return torch.stack(windows)
|
||||
|
||||
def relabel_from_value_predictions(
|
||||
self,
|
||||
predicted_values: torch.Tensor,
|
||||
*,
|
||||
force_interventions_positive: bool = True,
|
||||
) -> None:
|
||||
"""Recompute RECAP indicators from the learned critic predictions."""
|
||||
|
||||
predictions = torch.as_tensor(predicted_values, dtype=torch.float32)
|
||||
expected = (len(self.episodes), self.action_horizon + 1)
|
||||
if predictions.shape != expected:
|
||||
raise ValueError(f"predicted_values must have shape {expected}")
|
||||
rewards = stack_episode_field(self.episodes, "rewards")
|
||||
self.advantages = n_step_advantages(
|
||||
rewards,
|
||||
predictions,
|
||||
n_step=self.n_step,
|
||||
gamma=self.gamma,
|
||||
)
|
||||
intervention_masks = stack_episode_field(self.episodes, "intervention_mask").bool()
|
||||
step_interventions = intervention_masks if force_interventions_positive else None
|
||||
self.positive_steps, self.advantage_threshold = quantile_advantage_labels(
|
||||
self.advantages,
|
||||
quantile=self.positive_quantile,
|
||||
intervention_mask=step_interventions,
|
||||
)
|
||||
episode_advantages = self.advantages.mean(dim=-1)
|
||||
episode_interventions = (
|
||||
intervention_masks.any(dim=-1) if force_interventions_positive else None
|
||||
)
|
||||
self.positive_episodes, self.episode_advantage_threshold = quantile_advantage_labels(
|
||||
episode_advantages,
|
||||
quantile=self.positive_quantile,
|
||||
intervention_mask=episode_interventions,
|
||||
)
|
||||
|
||||
@property
|
||||
def action_statistics(self) -> ActionStatistics:
|
||||
return self.action_normalizer.statistics
|
||||
@@ -334,9 +395,10 @@ class Pi06ExperienceDataset(Dataset[dict[str, Any]]):
|
||||
"future_xy": episode.future_xy.clone(),
|
||||
"rewards": episode.rewards.clone(),
|
||||
"returns": self.normalized_returns[episode_index].clone(),
|
||||
# The critic predicts the beginning-of-episode return for this
|
||||
# minimal trajectory-level buffer; the full vector remains
|
||||
# available for diagnostics and advantage recomputation.
|
||||
"value_histories": self.value_histories[episode_index].clone(),
|
||||
"value_targets": self.value_targets[episode_index].clone(),
|
||||
# Scalar aliases are retained for simple diagnostics; value training
|
||||
# and RECAP relabeling use the complete per-state vectors above.
|
||||
"return_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"value_target": self.normalized_returns[episode_index, 0].clone(),
|
||||
"advantages": self.advantages[episode_index].clone(),
|
||||
|
||||
+65
-15
@@ -15,6 +15,29 @@ from .tokenizer import MiniTokenizer
|
||||
from .value import ValueConfig, ValueModel
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_subtask_token(
|
||||
model: Pi06Mini,
|
||||
tokenizer: MiniTokenizer,
|
||||
image: torch.Tensor,
|
||||
history_xy: torch.Tensor,
|
||||
action_horizon: int,
|
||||
device: torch.device,
|
||||
) -> int:
|
||||
"""Autoregressively produce the semantic token used by low-level control."""
|
||||
|
||||
query = tokenizer.build_subtask_query(history_xy)[None].to(device)
|
||||
zero_actions = torch.zeros(1, action_horizon, model.config.action_dim, device=device)
|
||||
output = model(
|
||||
image,
|
||||
query,
|
||||
zero_actions,
|
||||
torch.zeros(1, device=device),
|
||||
generated_token_start=query.shape[1],
|
||||
)
|
||||
return int(output.text_logits[0, -1].argmax().cpu())
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
checkpoint=torch.load(checkpoint_path,map_location="cpu",weights_only=False)
|
||||
@@ -35,13 +58,18 @@ def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
failure=next(dataset[i] for i in range(len(dataset)) if "stop" in dataset[i]["source_type"])
|
||||
image=demo["image"][None].to(device)
|
||||
|
||||
def condition(model:Pi06Mini, advantage:str|None, x:torch.Tensor,t:torch.Tensor)->torch.Tensor:
|
||||
seq=tokenizer.build_policy_sequence(
|
||||
demo["history_xy"],demo["reasoning"],demo["actions"],advantage,
|
||||
metadata=None,
|
||||
predicted_subtask = predict_subtask_token(
|
||||
recap, tokenizer, image, demo["history_xy"], dataset.action_horizon, device
|
||||
)
|
||||
ids=seq.input_ids[:seq.action_start_index][None].to(device)
|
||||
return model(image,ids,x,t,generated_token_start=seq.generated_token_start).vector_field
|
||||
|
||||
def condition(model:Pi06Mini, advantage:str|None, x:torch.Tensor,t:torch.Tensor)->torch.Tensor:
|
||||
ids = tokenizer.build_action_condition(
|
||||
demo["history_xy"], predicted_subtask, advantage
|
||||
)[None].to(device)
|
||||
return model(
|
||||
image, ids, x, t,
|
||||
generated_token_start=tokenizer.build_subtask_query(demo["history_xy"]).numel(),
|
||||
).vector_field
|
||||
|
||||
def sample_field(fn,seed:int=2):
|
||||
torch.manual_seed(seed)
|
||||
@@ -77,16 +105,38 @@ def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
|
||||
|
||||
value_rows=[]
|
||||
for i in range(len(dataset)):
|
||||
s=dataset[i];ids=tokenizer.build_value_input(s["history_xy"],s["source_type"])[None].to(device)
|
||||
prediction=float(value(s["image"][None].to(device),ids).expectation.cpu())
|
||||
value_rows.append({"source_type":s["source_type"],"target":float(s["value_target"]),"prediction":prediction})
|
||||
fig,ax=plt.subplots(figsize=(8,4));x=np.arange(len(value_rows));w=.35
|
||||
ax.bar(x-w/2,[r["target"] for r in value_rows],w,label="return target")
|
||||
ax.bar(x+w/2,[r["prediction"] for r in value_rows],w,label="value prediction")
|
||||
ax.set_xticks(x,[r["source_type"] for r in value_rows],rotation=20,ha="right");ax.set_ylim(-1.05,.05);ax.grid(True,axis="y",alpha=.3);ax.legend();ax.set_title("201-bin distributional value overfit")
|
||||
s=dataset[i]
|
||||
histories=s["value_histories"]
|
||||
ids=torch.stack([tokenizer.build_value_input(history) for history in histories]).to(device)
|
||||
images=s["image"][None].expand(len(histories),-1,-1,-1).to(device)
|
||||
predictions=value(images,ids).expectation.cpu()
|
||||
targets=s["value_targets"]
|
||||
value_rows.append({
|
||||
"source_type":s["source_type"],
|
||||
"targets":targets.tolist(),
|
||||
"predictions":predictions.tolist(),
|
||||
"curve_mse":float(torch.mean((predictions-targets).square())),
|
||||
})
|
||||
fig,axes=plt.subplots(2,2,figsize=(10,7),sharex=True,sharey=True)
|
||||
for ax,row in zip(axes.flat,value_rows):
|
||||
ax.plot(row["targets"],"o-",label="return target")
|
||||
ax.plot(row["predictions"],"x--",label="value prediction")
|
||||
ax.set_title(f'{row["source_type"]} MSE={row["curve_mse"]:.3f}')
|
||||
ax.set_ylim(-1.05,.05);ax.grid(True,alpha=.3)
|
||||
axes[0,0].legend(fontsize=8);fig.supxlabel("episode step");fig.supylabel("normalized value")
|
||||
value_path=outdir/"value_comparison.png";fig.tight_layout();fig.savefig(value_path,dpi=170);plt.close(fig)
|
||||
write_json(outdir/"evaluation.json",{"trajectory_metrics":metrics,"values":value_rows})
|
||||
print({"trajectory_metrics":metrics,"values":value_rows})
|
||||
try:
|
||||
predicted_subtask_label = tokenizer.subtasks.decode(predicted_subtask)
|
||||
except ValueError:
|
||||
predicted_subtask_label = f"token_{predicted_subtask}"
|
||||
evaluation = {
|
||||
"predicted_subtask": predicted_subtask_label,
|
||||
"target_subtask": tokenizer.subtasks.classify(demo["reasoning"]),
|
||||
"trajectory_metrics": metrics,
|
||||
"values": value_rows,
|
||||
}
|
||||
write_json(outdir/"evaluation.json",evaluation)
|
||||
print(evaluation)
|
||||
return trajectory_path,value_path
|
||||
|
||||
|
||||
|
||||
@@ -191,15 +191,18 @@ class MiniTokenizer:
|
||||
parts.append(query)
|
||||
return torch.cat(parts)
|
||||
|
||||
def build_value_input(self, history_xy: Tensor, metadata: str | None = None) -> Tensor:
|
||||
"""Return image marker plus discretized state, without policy query tokens."""
|
||||
def build_value_input(self, history_xy: Tensor) -> Tensor:
|
||||
"""Return observation tokens for the critic without outcome metadata.
|
||||
|
||||
RECAP's value function estimates ``V(o, task)``. Experience provenance
|
||||
and success/failure labels are training targets, not observable inputs;
|
||||
including them would let the critic solve the toy buffer by leakage.
|
||||
The single-task minimal experiment keeps the task command implicit.
|
||||
"""
|
||||
|
||||
coordinate_tokens = self.coordinates.encode(history_xy).reshape(-1).cpu()
|
||||
specials = torch.tensor([self.BOS, self.IMAGE, self.STATE], dtype=torch.long)
|
||||
parts = [specials, coordinate_tokens]
|
||||
metadata_token = self.encode_metadata(metadata)
|
||||
if metadata_token is not None: parts.append(metadata_token.reshape(1))
|
||||
return torch.cat(parts)
|
||||
return torch.cat((specials, coordinate_tokens))
|
||||
|
||||
# Backward-compatible name used by the pi0.5 data utilities.
|
||||
build_prefix = build_observation
|
||||
@@ -216,6 +219,28 @@ class MiniTokenizer:
|
||||
return None
|
||||
raise ValueError("advantage must be positive, negative, omitted, bool, or None")
|
||||
|
||||
def build_subtask_query(self, history_xy: Tensor) -> Tensor:
|
||||
"""Observation prefix whose final token predicts the next subtask."""
|
||||
|
||||
return self.build_observation(history_xy, metadata=None)
|
||||
|
||||
def build_action_condition(
|
||||
self,
|
||||
history_xy: Tensor,
|
||||
subtask_token: int | Tensor,
|
||||
advantage: Advantage | bool | None,
|
||||
) -> Tensor:
|
||||
"""Build the inference prefix ``observation -> subtask -> advantage``."""
|
||||
|
||||
pieces = [
|
||||
self.build_observation(history_xy, metadata=None),
|
||||
torch.as_tensor(subtask_token, dtype=torch.long).reshape(1).cpu(),
|
||||
]
|
||||
advantage_token = self.encode_advantage(advantage)
|
||||
if advantage_token is not None:
|
||||
pieces.append(advantage_token.reshape(1))
|
||||
return torch.cat(pieces)
|
||||
|
||||
@staticmethod
|
||||
def _orthonormal_dct_matrix(
|
||||
horizon: int, *, device: torch.device, dtype: torch.dtype
|
||||
|
||||
+46
-4
@@ -63,6 +63,27 @@ def policy_step(
|
||||
}
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_dataset_values(
|
||||
value: ValueModel,
|
||||
dataset: Pi06ExperienceDataset,
|
||||
tokenizer: MiniTokenizer,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
"""Run the learned critic over every state used for RECAP relabeling."""
|
||||
|
||||
value.eval()
|
||||
predictions: list[torch.Tensor] = []
|
||||
for episode_index in range(len(dataset.episodes)):
|
||||
sample = dataset[episode_index * dataset.augmentation_factor]
|
||||
histories = sample["value_histories"]
|
||||
ids = torch.stack([tokenizer.build_value_input(history) for history in histories])
|
||||
images = sample["image"][None].expand(len(histories), -1, -1, -1)
|
||||
output = value(images.to(device), ids.to(device))
|
||||
predictions.append(output.expectation.cpu())
|
||||
return torch.stack(predictions)
|
||||
|
||||
|
||||
def train(config_path: str) -> Path:
|
||||
cfg = load_config(config_path)
|
||||
root = Path(__file__).resolve().parents[2]
|
||||
@@ -73,7 +94,8 @@ def train(config_path: str) -> Path:
|
||||
manifest, image_size=cfg["data"]["image_size"],
|
||||
augmentation_factor=cfg["experience"]["augmentation_factor"],
|
||||
action_horizon=cfg["data"]["action_horizon"], n_step=cfg["experience"]["n_step"],
|
||||
gamma=cfg["experience"]["gamma"], positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
gamma=cfg["experience"].get("gamma", 1.0),
|
||||
positive_quantile=cfg["experience"]["positive_quantile"],
|
||||
)
|
||||
tokenizer = MiniTokenizer()
|
||||
base, value = build_models(cfg, tokenizer)
|
||||
@@ -99,11 +121,31 @@ def train(config_path: str) -> Path:
|
||||
value.train()
|
||||
for step in range(1, cfg["train"]["value_steps"] + 1):
|
||||
sample = dataset[(step - 1) % len(dataset)]
|
||||
ids = tokenizer.build_value_input(sample["history_xy"], sample["source_type"])[None].to(device)
|
||||
output = value(sample["image"][None].to(device), ids, target_values=sample["value_target"][None].to(device))
|
||||
time_index = (step - 1) % (dataset.action_horizon + 1)
|
||||
history_xy = sample["value_histories"][time_index]
|
||||
ids = tokenizer.build_value_input(history_xy)[None].to(device)
|
||||
target = sample["value_targets"][time_index][None].to(device)
|
||||
output = value(sample["image"][None].to(device), ids, target_values=target)
|
||||
value_opt.zero_grad(set_to_none=True); output.loss.backward(); value_opt.step()
|
||||
if step == 1 or step % cfg["train"]["log_every"] == 0:
|
||||
row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(sample["value_target"])};history.append(row);print(json.dumps(row))
|
||||
row={"stage":"value","step":step,"loss":float(output.loss.detach().cpu()),"prediction":float(output.expectation.detach().cpu()),"target":float(target.cpu())};history.append(row);print(json.dumps(row))
|
||||
|
||||
critic_values = predict_dataset_values(value, dataset, tokenizer, device)
|
||||
dataset.relabel_from_value_predictions(
|
||||
critic_values,
|
||||
force_interventions_positive=cfg["recap"].get(
|
||||
"force_interventions_positive", True
|
||||
),
|
||||
)
|
||||
write_json(outdir / "critic_relabeling.json", {
|
||||
"predicted_values": critic_values.tolist(),
|
||||
"advantage_threshold": float(dataset.advantage_threshold),
|
||||
"episode_advantage_threshold": float(dataset.episode_advantage_threshold),
|
||||
"episode_labels": [
|
||||
"positive" if bool(label) else "negative"
|
||||
for label in dataset.positive_episodes
|
||||
],
|
||||
})
|
||||
|
||||
def fit_policy(stage: str, steps: int, conditioned: bool) -> Pi06Mini:
|
||||
model, _ = build_models(cfg, tokenizer); model.load_state_dict(base_state); model.to(device).train()
|
||||
|
||||
@@ -71,6 +71,8 @@ def test_experience_dataset_contract_and_intervention_label(tmp_path: Path) -> N
|
||||
assert intervention["returns"].shape == (16,)
|
||||
assert intervention["return_target"].shape == ()
|
||||
assert intervention["value_target"].shape == ()
|
||||
assert intervention["value_histories"].shape == (17, 8, 2)
|
||||
assert intervention["value_targets"].shape == (17,)
|
||||
assert intervention["advantages"].shape == (16,)
|
||||
assert intervention["advantage_label"] == "positive"
|
||||
assert intervention["reasoning"] == "Steer left."
|
||||
@@ -78,3 +80,21 @@ def test_experience_dataset_contract_and_intervention_label(tmp_path: Path) -> N
|
||||
intervention["actions"], intervention["history_xy"]
|
||||
)
|
||||
torch.testing.assert_close(decoded, intervention["future_xy"], atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def test_learned_value_predictions_drive_advantage_relabeling(tmp_path: Path) -> None:
|
||||
manifest = _make_source_project(tmp_path / "source")
|
||||
dataset = Pi06ExperienceDataset(manifest, image_size=32)
|
||||
before = dataset.advantages.clone()
|
||||
predicted = torch.zeros(len(dataset.episodes), dataset.action_horizon + 1)
|
||||
predicted[:, :-1] = torch.linspace(-0.9, -0.1, dataset.action_horizon)
|
||||
|
||||
dataset.relabel_from_value_predictions(predicted)
|
||||
|
||||
assert not torch.allclose(dataset.advantages, before)
|
||||
intervention_index = next(
|
||||
index
|
||||
for index, episode in enumerate(dataset.episodes)
|
||||
if episode.source_type == "intervention_correction"
|
||||
)
|
||||
assert bool(dataset.positive_episodes[intervention_index])
|
||||
|
||||
@@ -37,6 +37,28 @@ def test_omitted_advantage_is_absent_not_a_third_condition_token() -> None:
|
||||
assert tokenizer.ADVANTAGE_NEGATIVE not in sequence.input_ids
|
||||
|
||||
|
||||
def test_value_input_does_not_encode_experience_outcome() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
history = torch.zeros(8, 2)
|
||||
value_input = tokenizer.build_value_input(history)
|
||||
|
||||
assert tokenizer.METADATA_DEMO not in value_input
|
||||
assert tokenizer.METADATA_UNDERSTEER not in value_input
|
||||
assert tokenizer.METADATA_STOP not in value_input
|
||||
assert tokenizer.METADATA_INTERVENTION not in value_input
|
||||
|
||||
|
||||
def test_predicted_subtask_can_be_connected_to_action_condition() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
history = torch.zeros(8, 2)
|
||||
predicted = tokenizer.subtasks.encode("steer_left")
|
||||
condition = tokenizer.build_action_condition(history, predicted, "positive")
|
||||
|
||||
assert condition[-2].item() == predicted
|
||||
assert condition[-1].item() == tokenizer.ADVANTAGE_POSITIVE
|
||||
assert tokenizer.ACTION_START not in condition
|
||||
|
||||
|
||||
def test_labels_predict_subtask_and_fast_but_not_external_advantage() -> None:
|
||||
tokenizer = MiniTokenizer()
|
||||
positive = tokenizer.build_policy_sequence(
|
||||
|
||||
Reference in New Issue
Block a user