seeker
149 lines
7.1 KiB
Python
149 lines
7.1 KiB
Python
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()
|