Upload UR5 full fine-tuned checkpoint: pi05_pour_full configs/check_policy_inference.py

configs/check_policy_inference.py

@@ -0,0 +1,568 @@
+import dataclasses
+import datetime as dt
+import contextlib
+import json
+import logging
+import pathlib
+import re
+import time
+from typing import Any
+
+import numpy as np
+import torch
+import tyro
+
+from openpi.policies import libero_policy
+from openpi.policies import policy_config as _policy_config
+from openpi.policies import ur5_policy
+from openpi.shared import normalize as _normalize
+from serve_policy_from_checkpoint import load_train_config_from_checkpoint
+
+
+@dataclasses.dataclass
+class Args:
+    checkpoint_dir: pathlib.Path
+    config_name: str | None = None
+    default_prompt: str | None = None
+    pytorch_device: str | None = None
+    num_denoise_steps: int = 2
+    lora_rank: int | None = None
+    lora_alpha: float | None = None
+    warmup_runs: int = 2
+    timed_runs: int = 10
+    profile_flops: bool = False
+    example: str = "auto"
+    train_log: pathlib.Path | None = None
+    metrics_out_path: pathlib.Path | None = None
+    table_out_path: pathlib.Path | None = None
+    norm_stats_path: pathlib.Path | None = None
+    model_label: str | None = None
+
+
+def main(args: Args) -> None:
+    checkpoint_metadata = _load_checkpoint_metadata(args.checkpoint_dir)
+    train_config = load_train_config_from_checkpoint(
+        args.checkpoint_dir,
+        config_name=args.config_name,
+        lora_rank=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+    )
+    norm_stats = _load_norm_stats_override(args.norm_stats_path)
+    policy = _policy_config.create_trained_policy(
+        train_config,
+        args.checkpoint_dir,
+        default_prompt=args.default_prompt,
+        sample_kwargs={"num_steps": args.num_denoise_steps},
+        norm_stats=norm_stats,
+        pytorch_device=args.pytorch_device,
+    )
+
+    example = _make_example(args.example, train_config)
+    last_result = None
+    device = _policy_device(policy)
+    if device is not None and device.type == "cuda":
+        torch.cuda.reset_peak_memory_stats(device)
+
+    with torch.inference_mode():
+        for _ in range(args.warmup_runs):
+            last_result = policy.infer(example)
+            _assert_finite_actions(last_result)
+    if torch.cuda.is_available() and (args.pytorch_device is None or "cuda" in args.pytorch_device):
+        torch.cuda.synchronize()
+
+    model_infer_ms = []
+    wall_ms = []
+    timing_values = {
+        "sample_total_ms": [],
+        "vlm_prefix_ms": [],
+        "denoising_ms": [],
+        "denoising_ms_per_step": [],
+        "denoising_steps": [],
+        "other_ms": [],
+    }
+    with torch.inference_mode():
+        for _ in range(args.timed_runs):
+            start_time = time.monotonic()
+            result = policy.infer(example)
+            if torch.cuda.is_available() and (args.pytorch_device is None or "cuda" in args.pytorch_device):
+                torch.cuda.synchronize()
+            wall_ms.append((time.monotonic() - start_time) * 1000)
+            _assert_finite_actions(result)
+            policy_timing = result.get("policy_timing", {})
+            model_infer_ms.append(float(policy_timing.get("infer_ms", float("nan"))))
+            _append_timing_values(timing_values, policy_timing)
+            last_result = result
+
+    if last_result is None:
+        raise RuntimeError("No inference runs were executed")
+
+    actions = np.asarray(last_result["actions"])
+    action_horizon = actions.shape[0]
+    model_infer_summary = _latency_summary(model_infer_ms)
+    wall_summary = _latency_summary(wall_ms)
+    model_metrics = _collect_model_metrics(policy, args.checkpoint_dir, checkpoint_metadata)
+    training_metrics = _collect_training_metrics(checkpoint_metadata, args.train_log)
+    approx_gflops = None
+    flop_profile = None
+    if args.profile_flops:
+        flop_profile = _profile_policy_flops(policy, example, device)
+        approx_gflops = flop_profile["total_gflops"]
+
+    metrics = {
+        "model_label": args.model_label or train_config.name,
+        "config": train_config.name,
+        "checkpoint": str(args.checkpoint_dir),
+        "norm_stats_path": str(args.norm_stats_path.resolve()) if args.norm_stats_path else None,
+        "global_step": checkpoint_metadata.get("global_step"),
+        "warmup_runs": args.warmup_runs,
+        "timed_runs": args.timed_runs,
+        "num_denoise_steps": args.num_denoise_steps,
+        "action_horizon": action_horizon,
+        "actions_shape": list(actions.shape),
+        "actions_mean": float(actions.mean()),
+        "actions_std": float(actions.std()),
+        "latency": {
+            "model_infer_ms": model_infer_summary,
+            "wall_ms": wall_summary,
+            "model_infer_per_action_ms": _latency_summary([value / action_horizon for value in model_infer_ms]),
+            "wall_per_action_ms": _latency_summary([value / action_horizon for value in wall_ms]),
+        },
+        "latency_decomposition": _latency_decomposition_summary(timing_values),
+        "frequency": _frequency_metrics(model_infer_summary, wall_summary, action_horizon),
+        "model": model_metrics,
+        "training": training_metrics,
+        "approx_gflops": approx_gflops,
+        "flop_profile": flop_profile,
+    }
+    metrics["table_row"] = _build_table_row(metrics)
+
+    print(f"config={train_config.name}")
+    print(f"checkpoint={args.checkpoint_dir}")
+    print(f"warmup_runs={args.warmup_runs}")
+    print(f"timed_runs={args.timed_runs}")
+    print(f"num_denoise_steps={args.num_denoise_steps}")
+    print(f"actions_shape={actions.shape}")
+    print(f"actions_mean={actions.mean():.6f}")
+    print(f"actions_std={actions.std():.6f}")
+    _print_latency("model_infer_ms", model_infer_ms)
+    _print_latency("wall_ms", wall_ms)
+    _print_latency("model_infer_per_action_ms", [value / action_horizon for value in model_infer_ms])
+    _print_latency("wall_per_action_ms", [value / action_horizon for value in wall_ms])
+    decomposition = metrics["latency_decomposition"]
+    if decomposition["available"]:
+        if timing_values["vlm_prefix_ms"]:
+            _print_latency("vlm_prefix_ms", timing_values["vlm_prefix_ms"])
+        if timing_values["denoising_ms"]:
+            _print_latency("denoising_ms", timing_values["denoising_ms"])
+        if timing_values["denoising_ms_per_step"]:
+            _print_latency("denoising_ms_per_step", timing_values["denoising_ms_per_step"])
+    print(f"model_query_hz={metrics['frequency']['model_query_hz']:.4f}")
+    print(f"model_action_generation_hz={metrics['frequency']['model_action_generation_hz']:.4f}")
+    print(f"wall_query_hz={metrics['frequency']['wall_query_hz']:.4f}")
+    print(f"wall_action_generation_hz={metrics['frequency']['wall_action_generation_hz']:.4f}")
+    print(f"model_total_params_b={model_metrics.get('model_total_params_b', float('nan')):.4f}")
+    print(f"model_trainable_params_m={model_metrics.get('model_trainable_params_m', float('nan')):.2f}")
+    print(f"checkpoint_size_gb={model_metrics.get('checkpoint_size_gb', float('nan')):.2f}")
+    if training_metrics.get("training_time_hours") is not None:
+        print(f"training_time_hours={training_metrics['training_time_hours']:.4f}")
+    if approx_gflops is not None:
+        print(f"approx_gflops={approx_gflops:.2f}")
+        if flop_profile is not None:
+            print(f"profiler_gflops={flop_profile['profiler_gflops']:.2f}")
+            print(f"manual_conv2d_gflops={flop_profile['manual_conv2d_gflops']:.2f}")
+            print(f"conv2d_modules_counted={flop_profile['manual_conv2d_modules_counted']}")
+    print("table_row_markdown=" + _format_table_row(metrics["table_row"]))
+
+    if args.metrics_out_path is not None:
+        args.metrics_out_path.parent.mkdir(parents=True, exist_ok=True)
+        args.metrics_out_path.write_text(json.dumps(metrics, indent=2) + "\n")
+        print(f"metrics_out_path={args.metrics_out_path}")
+    if args.table_out_path is not None:
+        args.table_out_path.parent.mkdir(parents=True, exist_ok=True)
+        args.table_out_path.write_text(_format_table(metrics["table_row"]))
+        print(f"table_out_path={args.table_out_path}")
+
+
+def _assert_finite_actions(result: dict) -> None:
+    actions = np.asarray(result["actions"])
+    if not np.isfinite(actions).all():
+        raise RuntimeError("Policy returned non-finite actions")
+
+
+def _make_example(example_name: str, train_config: Any) -> dict:
+    example_name = example_name.lower()
+    if example_name == "auto":
+        config_text = f"{train_config.name} {type(train_config.data).__name__}".lower()
+        example_name = "ur5" if "ur5" in config_text or "tube" in config_text or "pour" in config_text else "libero"
+    if example_name == "ur5":
+        return ur5_policy.make_ur5_example()
+    if example_name == "libero":
+        return libero_policy.make_libero_example()
+    raise ValueError(f"Unknown example={example_name!r}. Expected auto, ur5, or libero.")
+
+
+def _print_latency(name: str, values: list[float]) -> None:
+    summary = _latency_summary(values)
+    print(f"{name}_mean={summary['mean']:.2f}")
+    print(f"{name}_median={summary['median']:.2f}")
+    print(f"{name}_p95={summary['p95']:.2f}")
+    print(f"{name}_min={summary['min']:.2f}")
+    print(f"{name}_max={summary['max']:.2f}")
+
+
+def _append_timing_values(timing_values: dict[str, list[float]], policy_timing: dict[str, Any]) -> None:
+    for key in timing_values:
+        value = _finite_float(policy_timing.get(key))
+        if value is not None:
+            timing_values[key].append(value)
+
+
+def _latency_decomposition_summary(timing_values: dict[str, list[float]]) -> dict[str, Any]:
+    sample_total = _latency_summary_or_none(timing_values["sample_total_ms"])
+    vlm_prefix = _latency_summary_or_none(timing_values["vlm_prefix_ms"])
+    denoising = _latency_summary_or_none(timing_values["denoising_ms"])
+    denoising_per_step = _latency_summary_or_none(timing_values["denoising_ms_per_step"])
+    denoising_steps = _latency_summary_or_none(timing_values["denoising_steps"])
+    other = _latency_summary_or_none(timing_values["other_ms"])
+    return {
+        "available": bool(timing_values["sample_total_ms"]),
+        "measurement": (
+            "Cached-observation policy benchmark. VLM prefix is image/language prefix embedding plus "
+            "prefix transformer KV-cache construction. Denoising is the action denoising loop."
+        ),
+        "sample_total_ms": sample_total,
+        "vlm_prefix_ms": vlm_prefix,
+        "denoising_ms": denoising,
+        "denoising_ms_per_step": denoising_per_step,
+        "denoising_steps": denoising_steps,
+        "other_ms": other,
+        "fractions_of_mean_sample_total": {
+            "vlm_prefix": _fraction(vlm_prefix.get("mean"), sample_total.get("mean")),
+            "denoising": _fraction(denoising.get("mean"), sample_total.get("mean")),
+            "other": _fraction(other.get("mean"), sample_total.get("mean")),
+        },
+    }
+
+
+def _latency_summary_or_none(values: list[float]) -> dict[str, float | None]:
+    if not values:
+        return {
+            "mean": None,
+            "median": None,
+            "p95": None,
+            "min": None,
+            "max": None,
+        }
+    return _latency_summary(values)
+
+
+def _finite_float(value: Any) -> float | None:
+    try:
+        number = float(value)
+    except (TypeError, ValueError):
+        return None
+    return number if np.isfinite(number) else None
+
+
+def _fraction(numerator: Any, denominator: Any) -> float | None:
+    num = _finite_float(numerator)
+    den = _finite_float(denominator)
+    if num is None or den is None or den <= 0:
+        return None
+    return float(num / den)
+
+
+def _latency_summary(values: list[float]) -> dict[str, float]:
+    arr = np.asarray(values, dtype=np.float64)
+    return {
+        "mean": float(arr.mean()),
+        "median": float(np.percentile(arr, 50)),
+        "p95": float(np.percentile(arr, 95)),
+        "min": float(arr.min()),
+        "max": float(arr.max()),
+    }
+
+
+def _frequency_metrics(
+    model_infer_summary: dict[str, float],
+    wall_summary: dict[str, float],
+    action_horizon: int,
+) -> dict[str, float]:
+    model_query_hz = _hz_from_ms(model_infer_summary["median"])
+    wall_query_hz = _hz_from_ms(wall_summary["median"])
+    return {
+        "model_query_hz": model_query_hz,
+        "wall_query_hz": wall_query_hz,
+        "model_action_generation_hz": model_query_hz * action_horizon,
+        "wall_action_generation_hz": wall_query_hz * action_horizon,
+    }
+
+
+def _hz_from_ms(milliseconds: float) -> float:
+    return 1000.0 / milliseconds if milliseconds > 0 else float("nan")
+
+
+def _load_norm_stats_override(path: pathlib.Path | None):
+    if path is None:
+        return None
+    path = pathlib.Path(path).expanduser().resolve()
+    if path.is_file():
+        logging.info("Loaded explicit norm stats from %s", path)
+        return _normalize.deserialize_json(path.read_text(encoding="utf-8"))
+    logging.info("Loaded explicit norm stats from %s", path / "norm_stats.json")
+    return _normalize.load(path)
+
+
+def _load_checkpoint_metadata(checkpoint_dir: pathlib.Path) -> dict[str, Any]:
+    metadata_path = checkpoint_dir / "metadata.pt"
+    if not metadata_path.exists():
+        return {}
+    try:
+        return torch.load(metadata_path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(metadata_path, map_location="cpu")
+
+
+def _policy_device(policy) -> torch.device | None:
+    device_name = getattr(policy, "_pytorch_device", None)
+    if device_name is not None:
+        return torch.device(device_name)
+    model = getattr(policy, "_model", None)
+    if model is None:
+        return None
+    try:
+        return next(model.parameters()).device
+    except StopIteration:
+        return None
+
+
+def _collect_model_metrics(policy, checkpoint_dir: pathlib.Path, metadata: dict[str, Any]) -> dict[str, float | int | None]:
+    model = getattr(policy, "_model", None)
+    metrics = dict(metadata.get("efficiency_metrics") or {})
+    if model is not None:
+        if hasattr(model, "configure_trainable_parameters"):
+            model.configure_trainable_parameters()
+        params = list(model.parameters())
+        total_params = sum(param.numel() for param in params)
+        trainable_params = sum(param.numel() for param in params if param.requires_grad)
+        metrics.update(
+            {
+                "model_total_params": total_params,
+                "model_trainable_params": trainable_params,
+                "model_trainable_pct": trainable_params / total_params * 100 if total_params else 0.0,
+                "model_total_params_b": total_params / 1e9,
+                "model_trainable_params_m": trainable_params / 1e6,
+            }
+        )
+
+    weight_path = checkpoint_dir / "model.safetensors"
+    if weight_path.exists():
+        checkpoint_size_bytes = weight_path.stat().st_size
+        metrics["checkpoint_size_bytes"] = checkpoint_size_bytes
+        metrics["checkpoint_size_gb"] = checkpoint_size_bytes / 1024**3
+
+    device = _policy_device(policy)
+    if device is not None and device.type == "cuda" and torch.cuda.is_available():
+        metrics["cuda_peak_allocated_gb_inference"] = torch.cuda.max_memory_allocated(device) / 1024**3
+        metrics["cuda_peak_reserved_gb_inference"] = torch.cuda.max_memory_reserved(device) / 1024**3
+
+    return metrics
+
+
+def _collect_training_metrics(metadata: dict[str, Any], train_log: pathlib.Path | None) -> dict[str, Any]:
+    efficiency_metrics = metadata.get("efficiency_metrics") or {}
+    training_time_hours = (
+        efficiency_metrics.get("elapsed_training_hours")
+        or efficiency_metrics.get("elapsed_after_first_step_hours")
+        or efficiency_metrics.get("elapsed_run_hours")
+    )
+    metrics: dict[str, Any] = {
+        "training_time_hours": training_time_hours,
+        "source": "checkpoint_metadata" if training_time_hours is not None else None,
+    }
+    if train_log is not None:
+        parsed = _parse_training_time_from_log(train_log, metadata.get("global_step"))
+        metrics.update(parsed)
+        if parsed.get("training_time_hours") is not None:
+            metrics["training_time_hours"] = parsed["training_time_hours"]
+            metrics["source"] = "train_log"
+    return metrics
+
+
+def _parse_training_time_from_log(train_log: pathlib.Path, global_step: int | None) -> dict[str, Any]:
+    if not train_log.exists():
+        raise FileNotFoundError(f"Training log does not exist: {train_log}")
+
+    created = None
+    train_config = None
+    first_step = None
+    final_step = None
+    final_save = None
+    step_pattern = re.compile(rf"step={global_step}\s") if global_step is not None else None
+
+    for raw_line in train_log.read_text(errors="replace").replace("\r", "\n").splitlines():
+        timestamp = _parse_log_time(raw_line)
+        if timestamp is None:
+            continue
+        if "Created experiment checkpoint directory" in raw_line:
+            created = timestamp
+        if "Training config:" in raw_line:
+            train_config = timestamp
+        if re.search(r"step=1\s", raw_line):
+            first_step = timestamp
+        if step_pattern is not None and step_pattern.search(raw_line):
+            final_step = timestamp
+        if global_step is not None and f"Saved checkpoint at step {global_step}" in raw_line:
+            final_save = timestamp
+
+    end_time = final_save or final_step
+    start_time = train_config or first_step or created
+    training_time_hours = None if start_time is None or end_time is None else _elapsed_hours(start_time, end_time)
+    return {
+        "training_time_hours": training_time_hours,
+        "created_to_final_save_hours": None if created is None or final_save is None else _elapsed_hours(created, final_save),
+        "train_config_to_final_save_hours": (
+            None if train_config is None or final_save is None else _elapsed_hours(train_config, final_save)
+        ),
+        "first_step_to_final_save_hours": None if first_step is None or final_save is None else _elapsed_hours(first_step, final_save),
+        "first_step_to_final_step_hours": None if first_step is None or final_step is None else _elapsed_hours(first_step, final_step),
+        "log_path": str(train_log),
+    }
+
+
+def _parse_log_time(line: str) -> dt.datetime | None:
+    match = re.search(r"(\d{2}:\d{2}:\d{2}\.\d{3})", line)
+    if match is None:
+        return None
+    return dt.datetime.strptime(match.group(1), "%H:%M:%S.%f")
+
+
+def _elapsed_hours(start: dt.datetime, end: dt.datetime) -> float:
+    while end < start:
+        end += dt.timedelta(days=1)
+    return (end - start).total_seconds() / 3600
+
+
+def _profile_policy_flops(policy, example: dict, device: torch.device | None) -> dict[str, Any]:
+    activities = [torch.profiler.ProfilerActivity.CPU]
+    if device is not None and device.type == "cuda" and torch.cuda.is_available():
+        activities.append(torch.profiler.ProfilerActivity.CUDA)
+        torch.cuda.synchronize(device)
+
+    conv2d_counter = _Conv2dFlopCounter(policy)
+    with torch.inference_mode():
+        with conv2d_counter:
+            with torch.profiler.profile(activities=activities, with_flops=True) as profiler:
+                result = policy.infer(example)
+                _assert_finite_actions(result)
+                if device is not None and device.type == "cuda" and torch.cuda.is_available():
+                    torch.cuda.synchronize(device)
+
+    profiler_flops = sum((getattr(event, "flops", 0) or 0) for event in profiler.key_averages())
+    profiler_conv2d_flops = sum(
+        (getattr(event, "flops", 0) or 0)
+        for event in profiler.key_averages()
+        if "conv2d" in getattr(event, "key", "")
+    )
+    manual_conv2d_flops = conv2d_counter.total_flops
+
+    # PyTorch currently warns that some aten::conv2d FLOPs cannot be computed,
+    # and reports those as zero. Add our module-hook Conv2d count only when the
+    # profiler did not already account for Conv2d FLOPs to avoid double counting
+    # on future PyTorch versions.
+    added_manual_conv2d_flops = manual_conv2d_flops if profiler_conv2d_flops == 0 else 0
+    total_flops = profiler_flops + added_manual_conv2d_flops
+    return {
+        "total_flops": float(total_flops),
+        "total_gflops": float(total_flops) / 1e9,
+        "profiler_flops": float(profiler_flops),
+        "profiler_gflops": float(profiler_flops) / 1e9,
+        "profiler_conv2d_flops": float(profiler_conv2d_flops),
+        "profiler_conv2d_gflops": float(profiler_conv2d_flops) / 1e9,
+        "manual_conv2d_flops": float(manual_conv2d_flops),
+        "manual_conv2d_gflops": float(manual_conv2d_flops) / 1e9,
+        "manual_conv2d_flops_added": float(added_manual_conv2d_flops),
+        "manual_conv2d_gflops_added": float(added_manual_conv2d_flops) / 1e9,
+        "manual_conv2d_modules_counted": conv2d_counter.modules_counted,
+        "flop_convention": "2 FLOPs per multiply-add; bias additions are not counted.",
+        "note": (
+            "Total combines PyTorch profiler FLOPs with manual Conv2d FLOPs only "
+            "when profiler Conv2d FLOPs are zero."
+        ),
+    }
+
+
+class _Conv2dFlopCounter(contextlib.AbstractContextManager):
+    def __init__(self, policy) -> None:
+        self.model = getattr(policy, "_model", None)
+        self.handles: list[Any] = []
+        self.total_flops = 0.0
+        self.modules_counted = 0
+
+    def __enter__(self):
+        if self.model is None:
+            return self
+        for module in self.model.modules():
+            if isinstance(module, torch.nn.Conv2d):
+                self.handles.append(module.register_forward_hook(self._hook))
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> bool:
+        for handle in self.handles:
+            handle.remove()
+        self.handles.clear()
+        return False
+
+    def _hook(self, module: torch.nn.Conv2d, inputs: tuple[Any, ...], output: Any) -> None:
+        if not inputs:
+            return
+        x = inputs[0]
+        if not torch.is_tensor(x) or not torch.is_tensor(output) or output.ndim < 4:
+            return
+        batch = int(output.shape[0])
+        out_channels = int(output.shape[1])
+        out_h = int(output.shape[2])
+        out_w = int(output.shape[3])
+        kernel_h, kernel_w = module.kernel_size
+        in_channels = int(module.in_channels)
+        groups = int(module.groups)
+        macs_per_output = (in_channels // groups) * int(kernel_h) * int(kernel_w)
+        self.total_flops += float(batch * out_channels * out_h * out_w * macs_per_output * 2)
+        self.modules_counted += 1
+
+
+def _build_table_row(metrics: dict[str, Any]) -> dict[str, Any]:
+    model_metrics = metrics["model"]
+    training_metrics = metrics["training"]
+    return {
+        "Model": metrics["model_label"],
+        "Model Size": f"{model_metrics.get('model_total_params_b', float('nan')):.2f}B",
+        "Trainable Params": f"{model_metrics.get('model_trainable_params_m', float('nan')):.2f}M",
+        "Training Time (hours)": _format_optional(training_metrics.get("training_time_hours"), precision=3),
+        "GFLOPs": _format_optional(metrics.get("approx_gflops"), precision=2),
+        "Inference Speed (ms)": f"{metrics['latency']['model_infer_ms']['median']:.2f}",
+        "Inference Speed (Hz)": f"{metrics['frequency']['model_query_hz']:.3f}",
+        "Action Gen. (Hz)": f"{metrics['frequency']['model_action_generation_hz']:.3f}",
+        "Checkpoint Size": f"{model_metrics.get('checkpoint_size_gb', float('nan')):.2f}GB",
+    }
+
+
+def _format_optional(value: Any, *, precision: int) -> str:
+    if value is None:
+        return "n/a"
+    return f"{float(value):.{precision}f}"
+
+
+def _format_table_row(row: dict[str, Any]) -> str:
+    return "| " + " | ".join(str(value) for value in row.values()) + " |"
+
+
+def _format_table(row: dict[str, Any]) -> str:
+    header = "| " + " | ".join(row.keys()) + " |\n"
+    separator = "| " + " | ".join("---" for _ in row) + " |\n"
+    return header + separator + _format_table_row(row) + "\n"
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO, force=True)
+    main(tyro.cli(Args))