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 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) 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) predicted_subtask = predict_subtask_token( recap, tokenizer, image, demo["history_xy"], dataset.action_horizon, device ) 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) 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] 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) 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 def main(): parser=argparse.ArgumentParser();parser.add_argument("--checkpoint",default="outputs/overfit/checkpoint.pt") evaluate(parser.parse_args().checkpoint) if __name__=="__main__":main()