Initial pi0.6 minimal validation

src/pi06mini/infer.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from .data import Pi06ExperienceDataset
+from .flow import FlowMatching
+from .model import Pi06Config, Pi06Mini
+from .recap import classifier_free_guidance
+from .runtime import choose_device, write_json
+from .tokenizer import MiniTokenizer
+from .value import ValueConfig, ValueModel
+
+
+@torch.no_grad()
+def evaluate(checkpoint_path: str) -> tuple[Path, Path]:
+    checkpoint=torch.load(checkpoint_path,map_location="cpu",weights_only=False)
+    cfg=checkpoint["config"]; root=Path(__file__).resolve().parents[2]
+    device=choose_device(cfg.get("device","auto")); tokenizer=MiniTokenizer()
+    dataset=Pi06ExperienceDataset(
+        root/"data/overfit/overfit_manifest.jsonl",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"],
+        action_statistics=checkpoint["action_statistics"],
+    )
+    mixed=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));mixed.load_state_dict(checkpoint["mixed"]);mixed.to(device).eval()
+    recap=Pi06Mini(Pi06Config(**checkpoint["policy_config"]));recap.load_state_dict(checkpoint["recap"]);recap.to(device).eval()
+    value=ValueModel(ValueConfig(**checkpoint["value_config"]));value.load_state_dict(checkpoint["value"]);value.to(device).eval()
+    flow=FlowMatching(cfg["inference"]["flow_steps"])
+    demo=next(dataset[i] for i in range(len(dataset)) if dataset[i]["source_type"]=="demo_success")
+    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,
+        )
+        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 sample_field(fn,seed:int=2):
+        torch.manual_seed(seed)
+        return flow.sample(fn,(1,dataset.action_horizon,2),device,cfg["inference"]["flow_steps"])[0].cpu()
+
+    mixed_actions=sample_field(lambda x,t:condition(mixed,None,x,t))
+    uncond_actions=sample_field(lambda x,t:condition(recap,None,x,t))
+    positive_actions=sample_field(lambda x,t:condition(recap,"positive",x,t))
+    negative_actions=sample_field(lambda x,t:condition(recap,"negative",x,t))
+    def cfg_field(x,t):
+        return classifier_free_guidance(
+            condition(recap,"positive",x,t),condition(recap,None,x,t),
+            guidance_scale=cfg["inference"]["cfg_beta"],
+        )
+    cfg_actions=sample_field(cfg_field)
+    decoded={
+        "mixed":dataset.action_normalizer.decode(mixed_actions,demo["history_xy"]).numpy(),
+        "unconditional":dataset.action_normalizer.decode(uncond_actions,demo["history_xy"]).numpy(),
+        "positive":dataset.action_normalizer.decode(positive_actions,demo["history_xy"]).numpy(),
+        "negative":dataset.action_normalizer.decode(negative_actions,demo["history_xy"]).numpy(),
+        "cfg":dataset.action_normalizer.decode(cfg_actions,demo["history_xy"]).numpy(),
+    }
+    target=demo["future_xy"].numpy(); failure_xy=failure["future_xy"].numpy()
+    metrics={name:float(np.mean((xy-target)**2)) for name,xy in decoded.items()}
+    metrics["negative_to_failure_mse"]=float(np.mean((decoded["negative"]-failure_xy)**2))
+    fig,ax=plt.subplots(figsize=(9,6));ax.plot(target[:,0],target[:,1],"o-",label="success target",linewidth=2)
+    ax.plot(failure_xy[:,0],failure_xy[:,1],":",label="failure experience")
+    for name,color in [("mixed","gray"),("unconditional","purple"),("positive","green"),("negative","red"),("cfg","orange")]:
+        xy=decoded[name];ax.plot(xy[:,0],xy[:,1],"x--",label=f"{name} ({metrics[name]:.3f})",color=color)
+    ax.scatter([0],[0],c="black",marker="s",label="ego");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.6 Recap advantage-conditioned trajectories")
+    outdir=Path(checkpoint_path).parent;trajectory_path=outdir/"recap_trajectory_comparison.png";fig.tight_layout();fig.savefig(trajectory_path,dpi=170);plt.close(fig)
+
+    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")
+    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})
+    return trajectory_path,value_path
+
+
+def main():
+    parser=argparse.ArgumentParser();parser.add_argument("--checkpoint",default="outputs/overfit/checkpoint.pt")
+    evaluate(parser.parse_args().checkpoint)
+
+
+if __name__=="__main__":main()