from __future__ import annotations import argparse from pathlib import Path import matplotlib.pyplot as plt import numpy as np import torch from .data import Pi07SteerableDataset from .flow import FlowMatching from .model import HighLevelPolicy, Pi07Config, Pi07Mini, TinyGoalGenerator from .rtc import build_rtc_conditioning from .runtime import choose_device, write_json from .tokenizer import EpisodeMetadata, MiniTokenizer @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["device"]) tokenizer = MiniTokenizer() dataset = 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=checkpoint["action_statistics"], max_rtc_delay_steps=cfg["rtc"]["max_delay_steps"], ) policy_config = Pi07Config(**checkpoint["policy_config"]) policy = Pi07Mini(policy_config) policy.load_state_dict(checkpoint["policy"]) policy.to(device).eval() high_level = HighLevelPolicy(policy_config) high_level.load_state_dict(checkpoint["high_level"]) high_level.to(device).eval() goal_generator = TinyGoalGenerator(tokenizer.vocab_size) goal_generator.load_state_dict(checkpoint["goal"]) goal_generator.to(device).eval() flow = FlowMatching(cfg["flow"]["integration_steps"]) sample = next( dataset[index] for index in range(len(dataset)) if dataset[index]["policy_kind"] == "expert" ) metadata = EpisodeMetadata(**sample["token_metadata"]) image = sample["image"][None].to(device) high_level_prompt = tokenizer.encode_high_level_prompt( "drive safely", sample["past_subtasks"], metadata ) high_level_output = high_level( image, high_level_prompt[None].to(device), sample["history_images"][None].to(device), sample["history_states"][None].to(device), sample["current_state"][None].to(device), ) predicted_subtask = tokenizer.subtasks.decode( int(high_level_output.logits.argmax(-1).cpu()) ) goal_prompt, _ = tokenizer.encode_prompt( "drive safely", predicted_subtask, metadata ) goal_noise = torch.randn( sample["future_observation"][None].shape, generator=torch.Generator().manual_seed(cfg["evaluation"]["fixed_noise_seed"]), ).to(device) goal_grid = torch.linspace(0.0, 1.0, cfg["evaluation"]["goal_flow_steps"] + 1, device=device) generated_goal = goal_noise for index in range(cfg["evaluation"]["goal_flow_steps"]): time = goal_grid[index].expand(1) generated_goal = generated_goal + (goal_grid[index + 1] - goal_grid[index]) * goal_generator( image, goal_prompt[None].to(device), generated_goal, time ).vector_field generated_goal_cpu = generated_goal[0].cpu().clamp(0.0, 1.0) def sample_policy( *, include_metadata: bool = True, use_goal: bool = True, use_history: bool = True, guidance_beta: float | None = None, rtc_delay: int = 0, oracle_goal: bool = False, ) -> tuple[np.ndarray, torch.Tensor]: conditioned = tokenizer.build_policy_sequence( "drive safely", predicted_subtask, sample["actions"], metadata, include_episode_metadata=include_metadata, ) conditioned_ids = conditioned.input_ids[: conditioned.action_start_index][None].to(device) common = { "pixel_values": image, "history_pixel_values": sample["history_images"][None].to(device) if use_history else None, "history_states": sample["history_states"][None].to(device) if use_history else None, "current_state": sample["current_state"][None].to(device), "subgoal_pixel_values": ( sample["future_observation"][None].to(device) if oracle_goal else generated_goal ) if use_goal else None, } if guidance_beta is None: def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor: return policy( input_ids=conditioned_ids, noisy_actions=actions, timesteps=times, **common, ).vector_field else: unconditional = tokenizer.build_policy_sequence( "drive safely", predicted_subtask, sample["actions"], metadata, include_episode_metadata=False, ) unconditional_ids = unconditional.input_ids[ : unconditional.action_start_index ][None].to(device) def field(actions: torch.Tensor, times: torch.Tensor) -> torch.Tensor: return policy.metadata_cfg( conditioned_ids, unconditional_ids, guidance_scale=guidance_beta, noisy_actions=actions, timesteps=times, **common, ) rtc = build_rtc_conditioning( sample["actions"], rtc_delay, max_delay_steps=cfg["rtc"]["max_delay_steps"], ) torch.manual_seed(cfg["evaluation"]["fixed_noise_seed"]) normalized = flow.sample( field, (1, dataset.action_horizon, 2), device, cfg["flow"]["integration_steps"], action_prefix=rtc.action_prefix[None].to(device), delays=torch.tensor([rtc_delay], device=device), )[0].cpu() trajectory = dataset.action_normalizer.decode( normalized, sample["history_xy"] ).numpy() return trajectory, normalized curves: dict[str, np.ndarray] = {} normalized_predictions: dict[str, torch.Tensor] = {} configurations = { "full_generated_goal": {}, "full_oracle_goal": {"oracle_goal": True}, "no_metadata": {"include_metadata": False}, "no_goal": {"use_goal": False}, "no_history": {"use_history": False}, "metadata_cfg": {"guidance_beta": cfg["steering"]["cfg_beta"]}, } for name, arguments in configurations.items(): curves[name], normalized_predictions[name] = sample_policy(**arguments) target = sample["future_xy"].numpy() target_actions = sample["actions"] trajectory_metrics = { name: float(np.mean((trajectory - target) ** 2)) for name, trajectory in curves.items() } action_metrics = { name: float(torch.mean((actions - target_actions) ** 2)) for name, actions in normalized_predictions.items() } rtc_metrics: dict[str, dict[str, float]] = {} for delay in cfg["evaluation"]["compare_rtc_delays"]: trajectory, actions = sample_policy(rtc_delay=delay) postfix = torch.arange(dataset.action_horizon) >= delay rtc_metrics[str(delay)] = { "postfix_action_mse": float( torch.mean((actions[postfix] - target_actions[postfix]) ** 2) ), "trajectory_mse": float(np.mean((trajectory - target) ** 2)), } output_dir = Path(checkpoint_path).parent figure, axis = plt.subplots(figsize=(10, 7)) axis.plot(target[:, 0], target[:, 1], "o-", label="expert target", linewidth=2) for name, trajectory in curves.items(): axis.plot( trajectory[:, 0], trajectory[:, 1], "x--", label=f"{name} ({trajectory_metrics[name]:.3f})", ) axis.axis("equal") axis.grid(True, alpha=0.3) axis.legend(fontsize=8) axis.set_xlabel("forward x (m)") axis.set_ylabel("left y (m)") axis.set_title("pi0.7 predicted-subtask context and RTC evaluation") figure.tight_layout() trajectory_path = output_dir / "context_ablation.png" figure.savefig(trajectory_path, dpi=170) plt.close(figure) figure, 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(sample["future_observation"].permute(1, 2, 0)) axes[1].set_title("future observation target") axes[2].imshow(generated_goal_cpu.permute(1, 2, 0)) axes[2].set_title("generated visual subgoal") for axis in axes: axis.axis("off") figure.tight_layout() goal_path = output_dir / "goal_generation.png" figure.savefig(goal_path, dpi=170) plt.close(figure) results = { "target_subtask": tokenizer._classify_subtask(sample["reasoning"]), "predicted_subtask": predicted_subtask, "goal_image_mse": float( torch.mean((generated_goal_cpu - sample["future_observation"]) ** 2) ), "trajectory_metrics": trajectory_metrics, "normalized_action_metrics": action_metrics, "rtc_metrics": rtc_metrics, } write_json(output_dir / "evaluation.json", results) print(results) return trajectory_path, goal_path def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--checkpoint", default="outputs/overfit/checkpoint.pt") evaluate(parser.parse_args().checkpoint) if __name__ == "__main__": main()