seeker
259 lines
9.5 KiB
Python
259 lines
9.5 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 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()
|