seeker
245 lines
7.9 KiB
Python
245 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import statistics
|
|
from functools import partial
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
from alpamayo_lite.action_space import RelativeDeltaActionSpace
|
|
from alpamayo_lite.config import ExperimentConfig, load_config
|
|
from alpamayo_lite.data import OverfitSample, load_manifest
|
|
from alpamayo_lite.diffusion import FlowMatching
|
|
from alpamayo_lite.evaluation import (
|
|
save_prediction_arrays,
|
|
save_trajectory_plot,
|
|
trajectory_metrics,
|
|
)
|
|
from alpamayo_lite.expert import ExpertArchitecture, KVCachedDiffusionExpert, PrefixKV
|
|
from alpamayo_lite.prompts import build_expert_prefix, build_vlm_prefix
|
|
from alpamayo_lite.runtime import (
|
|
empty_device_cache,
|
|
parameter_counts,
|
|
resolve_dtype,
|
|
seed_everything,
|
|
select_device,
|
|
write_json,
|
|
)
|
|
from alpamayo_lite.vlm import (
|
|
extract_prefix_kv,
|
|
generate_reasoning,
|
|
load_vlm,
|
|
vlm_text_config,
|
|
)
|
|
|
|
|
|
def _resolve(cfg: ExperimentConfig, value: str) -> Path:
|
|
path = Path(value).expanduser()
|
|
return path.resolve() if path.is_absolute() else (cfg.root / path).resolve()
|
|
|
|
|
|
def _read_json(path: Path) -> dict[str, Any]:
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def _summary(values: list[float]) -> dict[str, float]:
|
|
return {
|
|
"mean": statistics.fmean(values),
|
|
"std": statistics.pstdev(values),
|
|
"min": min(values),
|
|
"max": max(values),
|
|
}
|
|
|
|
|
|
def _evaluate_expert(
|
|
stage2_dir: Path,
|
|
vlm: Any,
|
|
prefix: PrefixKV,
|
|
sample: OverfitSample,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
num_rollouts: int,
|
|
base_seed: int,
|
|
) -> dict[str, Any]:
|
|
expert_config = _read_json(stage2_dir / "expert_config.json")
|
|
metrics_path = stage2_dir / "stage2_metrics.json"
|
|
saved_metrics = _read_json(metrics_path) if metrics_path.exists() else {}
|
|
architecture = ExpertArchitecture(**expert_config["architecture"])
|
|
expert = KVCachedDiffusionExpert(vlm_text_config(vlm), architecture).to(
|
|
device=device, dtype=dtype
|
|
)
|
|
expert.load_state_dict(load_file(stage2_dir / "expert.safetensors"), strict=True)
|
|
expert.requires_grad_(False)
|
|
expert.eval()
|
|
|
|
action_shape = tuple(int(value) for value in expert_config["action_shape"])
|
|
if len(action_shape) != 2:
|
|
raise ValueError(f"Invalid action shape in {stage2_dir}: {action_shape}")
|
|
flow = FlowMatching(
|
|
action_shape=action_shape,
|
|
inference_steps=int(expert_config["inference_steps"]),
|
|
)
|
|
action_space = RelativeDeltaActionSpace(
|
|
torch.tensor(expert_config["action_mean"], device=device),
|
|
torch.tensor(expert_config["action_scale"], device=device),
|
|
).to(device)
|
|
history, future = sample.load_trajectory()
|
|
history_xy = history[None, :, :2].to(device)
|
|
target_xy = future[None, :, :2].to(device)
|
|
target_action = action_space.trajectory_to_action(history_xy, target_xy).float()
|
|
|
|
rollouts: list[dict[str, float | int]] = []
|
|
first_prediction: tuple[torch.Tensor, torch.Tensor] | None = None
|
|
for index in range(num_rollouts):
|
|
seed = base_seed + index
|
|
seed_everything(seed)
|
|
predicted_action = flow.sample(
|
|
batch_size=1,
|
|
step_fn=partial(expert, prefix=prefix),
|
|
device=device,
|
|
dtype=dtype,
|
|
).float()
|
|
predicted_xy = action_space.action_to_trajectory(predicted_action, history_xy)
|
|
record: dict[str, float | int] = {
|
|
"seed": seed,
|
|
"action_mse": float(
|
|
torch.nn.functional.mse_loss(predicted_action, target_action).cpu()
|
|
),
|
|
**trajectory_metrics(predicted_xy, target_xy),
|
|
}
|
|
if not torch.isfinite(predicted_action).all() or not torch.isfinite(predicted_xy).all():
|
|
raise RuntimeError(f"Non-finite reload prediction for seed {seed}")
|
|
rollouts.append(record)
|
|
if first_prediction is None:
|
|
first_prediction = (predicted_action, predicted_xy)
|
|
|
|
assert first_prediction is not None
|
|
first_action, first_xy = first_prediction
|
|
save_prediction_arrays(
|
|
stage2_dir / "reload_prediction.npz",
|
|
history[:, :2],
|
|
target_action[0],
|
|
first_action[0],
|
|
target_xy[0],
|
|
first_xy[0],
|
|
)
|
|
save_trajectory_plot(
|
|
stage2_dir / "reload_trajectory.png",
|
|
history[:, :2],
|
|
target_xy[0],
|
|
first_xy[0],
|
|
)
|
|
|
|
metric_names = ("action_mse", "ade_m", "fde_m", "max_error_m")
|
|
result = {
|
|
"stage2_dir": str(stage2_dir),
|
|
"checkpoint_steps": int(
|
|
expert_config.get("completed_steps", saved_metrics.get("steps", 0))
|
|
),
|
|
"expert_parameters": parameter_counts(expert),
|
|
"architecture": expert_config["architecture"],
|
|
"action_mean": expert_config["action_mean"],
|
|
"action_scale": expert_config["action_scale"],
|
|
"action_shape": list(action_shape),
|
|
"inference_steps": flow.inference_steps,
|
|
"num_rollouts": num_rollouts,
|
|
"rollouts": rollouts,
|
|
"summary": {
|
|
name: _summary([float(record[name]) for record in rollouts])
|
|
for name in metric_names
|
|
},
|
|
}
|
|
del expert
|
|
empty_device_cache(device)
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", default="configs/local_overfit.yaml")
|
|
parser.add_argument(
|
|
"--stage2-dir",
|
|
action="append",
|
|
dest="stage2_dirs",
|
|
help="Stage-2 directory to reload; repeat to compare checkpoints.",
|
|
)
|
|
parser.add_argument("--num-rollouts", type=int, default=16)
|
|
parser.add_argument("--base-seed", type=int, default=None)
|
|
parser.add_argument("--output", default="outputs/reload_evaluation.json")
|
|
args = parser.parse_args()
|
|
if args.num_rollouts <= 0:
|
|
raise ValueError("--num-rollouts must be positive")
|
|
|
|
cfg = load_config(args.config)
|
|
seed_everything(cfg.seed)
|
|
device = select_device()
|
|
dtype = resolve_dtype(cfg.vlm.dtype, device)
|
|
stage2_values = args.stage2_dirs or [cfg.paths.stage2_dir]
|
|
stage2_dirs = [_resolve(cfg, value) for value in stage2_values]
|
|
sample = load_manifest(cfg.resolve(cfg.paths.manifest))[0]
|
|
history, _ = sample.load_trajectory()
|
|
|
|
stage1_dir = cfg.resolve(cfg.paths.stage1_dir)
|
|
processor, vlm = load_vlm(stage1_dir, cfg.vlm, device, dtype)
|
|
vlm.requires_grad_(False)
|
|
vlm.eval()
|
|
generated_reasoning = generate_reasoning(
|
|
vlm,
|
|
processor,
|
|
sample.load_image(),
|
|
build_vlm_prefix(history),
|
|
device,
|
|
dtype,
|
|
cfg.vlm.max_reasoning_tokens,
|
|
)
|
|
prefix = extract_prefix_kv(
|
|
vlm,
|
|
processor,
|
|
sample.load_image(),
|
|
build_expert_prefix(history, generated_reasoning),
|
|
device,
|
|
dtype,
|
|
)
|
|
base_seed = args.base_seed if args.base_seed is not None else cfg.seed + 1
|
|
models = [
|
|
_evaluate_expert(
|
|
stage2_dir,
|
|
vlm,
|
|
prefix,
|
|
sample,
|
|
device,
|
|
dtype,
|
|
args.num_rollouts,
|
|
base_seed,
|
|
)
|
|
for stage2_dir in stage2_dirs
|
|
]
|
|
result = {
|
|
"stage": "reloaded_pipeline_evaluation",
|
|
"device": str(device),
|
|
"dtype": str(dtype),
|
|
"stage1_dir": str(stage1_dir),
|
|
"vlm_parameters": parameter_counts(vlm),
|
|
"generated_reasoning": generated_reasoning,
|
|
"target_reasoning": sample.reasoning,
|
|
"reasoning_exact_match": generated_reasoning == sample.reasoning,
|
|
"prefix_layers": prefix.num_layers,
|
|
"prefix_sequence_length": prefix.sequence_length,
|
|
"prefix_key_shape": list(prefix.keys[0].shape),
|
|
"models": models,
|
|
}
|
|
output = _resolve(cfg, args.output)
|
|
write_json(output, result)
|
|
print(json.dumps(result, indent=2, ensure_ascii=False))
|
|
del vlm
|
|
empty_device_cache(device)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|