workspace/scripts/eval_ctt_generated_rollout.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import importlib
import importlib.machinery
import json
import math
import os
import pickle
import shutil
import subprocess
import sys
import types
from dataclasses import dataclass
from pathlib import Path
from typing import Any

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import numpy as np  # noqa: E402
import torch  # noqa: E402

from cil.chart_features import build_chart_feature  # noqa: E402
from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer, UtilityEnergy  # noqa: E402
from dovla_cil.generation.maniskill_parallel import execute_grouped_action_lattice_batch  # noqa: E402
from dovla_cil.utils.io import read_json  # noqa: E402
from scripts.eval_metrics import main as eval_metrics_main  # noqa: E402


@dataclass(frozen=True)
class ChartItem:
    chart_id: str
    task_id: str
    seed: str
    state_hash: str
    instruction: str
    source_dataset: Path
    base_action: np.ndarray
    feature: np.ndarray
    positive_tangents: np.ndarray
    negative_tangents: np.ndarray
    hidden_utilities: list[float]
    hidden_candidate_types: list[str]
    stored_base_utility: float | None


@dataclass(frozen=True)
class Proposal:
    tangent: np.ndarray
    action: np.ndarray
    score: float
    source_chart_id: str
    source_task_id: str
    source_rank: int


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Generate CTT candidates, decode them to ManiSkill action chunks, "
            "and measure them with same-state simulator rollouts."
        )
    )
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--source-index", type=Path, default=Path("data/cil_charts/train/index.json"))
    parser.add_argument("--target-index", type=Path, default=Path("data/cil_charts/val/index.json"))
    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_residual_rollout_smoke"))
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--pool-size", type=int, default=0)
    parser.add_argument("--neighbors", type=int, default=8)
    parser.add_argument("--max-target-charts", type=int, default=8)
    parser.add_argument("--group-batch-size", type=int, default=1)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--sim-backend", default=None)
    parser.add_argument("--render-backend", default=None)
    parser.add_argument("--restore-tolerance", type=float, default=1.0e-5)
    parser.add_argument("--delta-scale", type=float, default=1.0)
    parser.add_argument(
        "--execution-action-scale",
        type=float,
        default=1.0,
        help=(
            "Multiply the full base/generated action chunk before action-bound "
            "diagnostics, optional clipping, and simulator execution. This is a "
            "diagnostic for action-representation scale mismatch; keep the default "
            "1.0 for the unscaled deployment convention."
        ),
    )
    parser.add_argument(
        "--execution-action-scale-vector",
        default="",
        help=(
            "Optional comma-separated per-action-dimension positive scale vector "
            "applied after --execution-action-scale and before the execution "
            "transform. This diagnoses dimension-wise action representation mismatch."
        ),
    )
    parser.add_argument(
        "--execution-action-transform",
        choices=("identity", "tanh", "env_clip"),
        default="identity",
        help=(
            "Action-convention transform applied before action-bound diagnostics, "
            "optional clipping, and simulator execution. 'identity' preserves "
            "decoded controls; 'tanh' smoothly maps controls into finite env bounds; "
            "'env_clip' clips decoded controls to the finite env bounds as an explicit "
            "logged decoder convention rather than silent simulator clipping."
        ),
    )
    parser.add_argument(
        "--disable-action-clipping",
        action="store_true",
        help=(
            "Execute decoded chunks without clipping to env.action_space. "
            "Action-bound diagnostics are still recorded."
        ),
    )
    parser.add_argument("--include-targets-without-positives", action="store_true")
    parser.add_argument(
        "--exclude-self-source",
        action="store_true",
        help=(
            "When source and target indexes overlap, exclude source charts with the "
            "same chart_id or state_hash as the target. Use this for train-split "
            "calibration rollouts so retrieval cannot copy the target chart's own positives."
        ),
    )
    parser.add_argument("--skip-metrics", action="store_true")
    parser.add_argument("--bootstrap-samples", type=int, default=200)
    parser.add_argument(
        "--no-markdown-report",
        action="store_true",
        help="Do not write report.md files; persistent prose lives in README.md.",
    )
    args = parser.parse_args(argv)

    if args.k <= 0:
        parser.error("--k must be positive")
    if args.neighbors <= 0:
        parser.error("--neighbors must be positive")
    if args.group_batch_size <= 0:
        parser.error("--group-batch-size must be positive")
    if args.max_target_charts <= 0:
        parser.error("--max-target-charts must be positive")
    if args.restore_tolerance <= 0.0:
        parser.error("--restore-tolerance must be positive")
    if args.execution_action_scale <= 0.0:
        parser.error("--execution-action-scale must be positive")
    try:
        execution_action_scale_vector = _parse_execution_action_scale_vector(
            args.execution_action_scale_vector
        )
    except ValueError as exc:
        parser.error(str(exc))

    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    _write_run_provenance(out_dir, args)
    log_path = out_dir / "run.log"
    _append_log(log_path, "start")
    _append_log(log_path, "importing gymnasium/mani_skill")
    try:
        import gymnasium as gym
        _register_required_maniskill_envs(log_path)
    except ImportError as exc:  # pragma: no cover - exercised in the Apptainer env
        raise ImportError(
            "CTT measured rollout requires gymnasium, mani_skill, numpy, and torch. "
            "Run this script through the ManiSkill Apptainer environment on HPC."
        ) from exc
    _append_log(log_path, "imported gymnasium/mani_skill")

    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    config = CTTConfig(**checkpoint["config"])
    chart_feature_mode = str(checkpoint.get("chart_feature_mode", "base"))
    encoder = ChartEncoder(config.chart_feature_dim, output_dim=config.chart_dim)
    ctt = CausalTangentTransport(config)
    utility_energy = UtilityEnergy(chart_dim=config.chart_dim, tangent_dim=config.tangent_dim)
    encoder.load_state_dict(checkpoint["chart_encoder"])
    ctt.load_state_dict(checkpoint["ctt"])
    if "utility_energy" not in checkpoint:
        raise SystemExit(f"{args.checkpoint} does not contain a utility_energy state")
    utility_energy.load_state_dict(checkpoint["utility_energy"])
    normalizer = TangentNormalizer.from_dict(checkpoint["normalizer"])
    encoder.eval()
    ctt.eval()
    utility_energy.eval()
    for module in (encoder, ctt, utility_energy):
        for parameter in module.parameters():
            parameter.requires_grad_(False)
    _append_log(log_path, f"loaded checkpoint={args.checkpoint}")

    source_charts, source_index = load_chart_items(
        args.source_index,
        max_charts=None,
        require_positive=True,
        include_hidden=False,
        include_metadata=True,
        chart_feature_mode=chart_feature_mode,
    )
    target_charts, target_index = load_chart_items(
        args.target_index,
        max_charts=args.max_target_charts,
        require_positive=not args.include_targets_without_positives,
        include_hidden=True,
        include_metadata=True,
        chart_feature_mode=chart_feature_mode,
    )
    _validate_indexes(args.source_index, source_index, args.target_index, target_index)
    if not target_charts:
        raise SystemExit("No target charts available after filtering")
    _append_log(
        log_path,
        f"loaded charts source={len(source_charts)} target={len(target_charts)}",
    )

    resolved_device = _resolve_device(args.device)
    encoder.to(resolved_device)
    ctt.to(resolved_device)
    utility_energy.to(resolved_device)
    source_by_task: dict[str, list[ChartItem]] = {}
    for chart in source_charts:
        source_by_task.setdefault(chart.task_id, []).append(chart)

    pool_size = int(args.pool_size) if args.pool_size > 0 else int(args.k)
    generated_cases = [
        (
            target,
            generate_proposals(
                target,
                source_charts=source_charts,
                source_by_task=source_by_task,
                encoder=encoder,
                ctt=ctt,
                utility_energy=utility_energy,
                normalizer=normalizer,
                device=resolved_device,
                neighbors=args.neighbors,
                pool_size=max(pool_size, args.k),
                k=args.k,
                delta_scale=args.delta_scale,
                exclude_self_source=args.exclude_self_source,
            ),
        )
        for target in target_charts
    ]
    _append_log(
        log_path,
        "generated proposals "
        f"rows={len(generated_cases)} total={sum(len(item[1]) for item in generated_cases)}",
    )
    rows = rollout_generated_cases(
        generated_cases,
        gym=gym,
        torch=torch,
        device=resolved_device,
        group_batch_size=args.group_batch_size,
        sim_backend=args.sim_backend,
        render_backend=args.render_backend,
        restore_tolerance=args.restore_tolerance,
        clip_actions=not args.disable_action_clipping,
        execution_action_scale=args.execution_action_scale,
        execution_action_scale_vector=execution_action_scale_vector,
        execution_action_transform=args.execution_action_transform,
        log_path=log_path,
    )
    _append_log(log_path, f"rollout complete rows={len(rows)}")

    payload = {
        "report_type": "ctt_generated_measured_rollout",
        "candidates_evaluated": True,
        "schema_version": 1,
        "checkpoint": str(args.checkpoint),
        "source_index": str(args.source_index),
        "target_index": str(args.target_index),
        "source_content_hash": source_index.get("content_hash"),
        "source_split_hash": source_index.get("split_hash"),
        "target_content_hash": target_index.get("content_hash"),
        "target_split_hash": target_index.get("split_hash"),
        "k": args.k,
        "neighbors": args.neighbors,
        "pool_size": max(pool_size, args.k),
        "exclude_self_source": bool(args.exclude_self_source),
        "decoder": {
            "name": "linear_keyframe_decode",
            "source_code": "spline_tangent_code stores start/mid/end residual keyframes",
            "lossless": False,
            "delta_scale": args.delta_scale,
            "execution_action_scale": args.execution_action_scale,
            "execution_action_scale_vector": (
                None
                if execution_action_scale_vector is None
                else execution_action_scale_vector.astype(float).tolist()
            ),
            "execution_action_transform": args.execution_action_transform,
            "action_clipping_enabled": not args.disable_action_clipping,
        },
        "rows": rows,
    }
    measured_path = out_dir / "measured_candidates.json"
    measured_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    report_path = out_dir / "report.md"
    if args.no_markdown_report:
        report_path.unlink(missing_ok=True)
    else:
        report_path.write_text(_report(payload) + "\n")

    metrics_dir = out_dir / "measured_metrics"
    if not args.skip_metrics:
        metric_args = [
            "--input",
            str(measured_path),
            "--out-dir",
            str(metrics_dir),
            "--mode",
            "measured",
            "--k",
            str(args.k),
            "--bootstrap-samples",
            str(args.bootstrap_samples),
        ]
        if args.no_markdown_report:
            metric_args.append("--no-markdown-report")
        eval_metrics_main(metric_args)
    _write_required_artifacts(
        out_dir,
        payload,
        source_index=source_index,
        target_index=target_index,
        metrics_dir=metrics_dir if metrics_dir.exists() else None,
    )

    print(
        json.dumps(
            {
                "out_dir": str(out_dir),
                "num_rows": len(rows),
                "measured_candidates": str(measured_path),
            },
            indent=2,
        )
    )
    return 0


def load_chart_items(
    index_path: Path,
    *,
    max_charts: int | None,
    require_positive: bool,
    include_hidden: bool,
    include_metadata: bool,
    chart_feature_mode: str = "base",
) -> tuple[list[ChartItem], dict[str, Any]]:
    index = json.loads(index_path.read_text())
    grouped: dict[str, dict[str, Any]] = {}
    for shard in index.get("shards", []):
        shard_path = index_path.parent / shard["path"]
        with np.load(shard_path, allow_pickle=False) as data:
            chart_ids = data["chart_id"]
            task_ids = data["task_id"]
            seeds = data["seed"]
            state_hashes = data["state_hash"]
            action_shapes = data["action_shape"]
            base_actions = data["base_action"]
            labels = data["label"]
            spline_tangents = data["spline_tangent_code"]
            is_base_branch = data["is_base_branch"]
            utilities = data["utility"] if include_hidden else None
            candidate_types = data["candidate_type"] if include_hidden else None
            metadata_values = data["metadata_json"] if include_metadata else None
            for row in range(chart_ids.shape[0]):
                chart_id = str(chart_ids[row])
                task_id = str(task_ids[row])
                metadata = (
                    _json_loads(str(metadata_values[row]))
                    if metadata_values is not None
                    else {}
                ) | {"_chart_root": str(index_path.parent)}
                shape = tuple(int(value) for value in action_shapes[row])
                flat_count = int(math.prod(shape))
                base_action = np.asarray(
                    base_actions[row][:flat_count], dtype=np.float32
                ).reshape(shape)
                item = grouped.setdefault(
                    chart_id,
                    {
                        "chart_id": chart_id,
                        "task_id": task_id,
                        "seed": str(seeds[row]),
                        "state_hash": str(state_hashes[row]),
                        "instruction": str(metadata.get("instruction", "")),
                        "metadata": metadata | {"task_id": task_id},
                        "source_dataset": _source_dataset_from_metadata(
                            metadata,
                            index=index,
                            task_id=task_id,
                        ),
                        "base_action": base_action,
                        "positive_tangents": [],
                        "negative_tangents": [],
                        "hidden_utilities": [],
                        "hidden_candidate_types": [],
                        "stored_base_utility": None,
                    },
                )
                label = str(labels[row])
                tangent = np.asarray(spline_tangents[row], dtype=np.float32)
                if label == "positive":
                    item["positive_tangents"].append(tangent)
                elif label == "negative":
                    item["negative_tangents"].append(tangent)
                utility = float(utilities[row]) if utilities is not None else math.nan
                if include_hidden and math.isfinite(utility):
                    item["hidden_utilities"].append(utility)
                    item["hidden_candidate_types"].append(str(candidate_types[row]))
                if bool(is_base_branch[row]):
                    item["stored_base_utility"] = utility if math.isfinite(utility) else None
                    item["base_action"] = base_action
                    item["metadata"] = metadata | {"task_id": task_id}

    charts: list[ChartItem] = []
    for chart_id, item in sorted(grouped.items()):
        positives = _matrix_or_empty(item["positive_tangents"], width=21)
        negatives = _matrix_or_empty(item["negative_tangents"], width=21)
        if require_positive and not len(positives):
            continue
        base_action = np.asarray(item["base_action"], dtype=np.float32)
        charts.append(
            ChartItem(
                chart_id=chart_id,
                task_id=str(item["task_id"]),
                seed=str(item["seed"]),
                state_hash=str(item["state_hash"]),
                instruction=str(item["instruction"]),
                source_dataset=Path(item["source_dataset"]).resolve(),
                base_action=base_action,
                feature=build_chart_feature(
                    base_action,
                    item.get("metadata", {}),
                    mode=chart_feature_mode,
                ),
                positive_tangents=positives,
                negative_tangents=negatives,
                hidden_utilities=[float(value) for value in item["hidden_utilities"]],
                hidden_candidate_types=[str(value) for value in item["hidden_candidate_types"]],
                stored_base_utility=item["stored_base_utility"],
            )
        )
        if max_charts is not None and len(charts) >= int(max_charts):
            break
    return charts, index


def generate_proposals(
    target: ChartItem,
    *,
    source_charts: list[ChartItem],
    source_by_task: dict[str, list[ChartItem]],
    encoder: ChartEncoder,
    ctt: CausalTangentTransport,
    utility_energy: UtilityEnergy,
    normalizer: TangentNormalizer,
    device: str,
    neighbors: int,
    pool_size: int,
    k: int,
    delta_scale: float,
    exclude_self_source: bool = False,
) -> list[Proposal]:
    task_pool = source_by_task.get(target.task_id) or source_charts
    pool = _source_pool_for_target(
        target,
        task_pool=task_pool,
        source_charts=source_charts,
        exclude_self_source=exclude_self_source,
    )
    target_feature = torch.as_tensor(target.feature, dtype=torch.float32, device=device)
    ranked_sources = sorted(
        pool,
        key=lambda source: float(
            torch.linalg.vector_norm(
                torch.as_tensor(source.feature, dtype=torch.float32, device=device)
                - target_feature
            )
            .detach()
            .cpu()
        ),
    )[:neighbors]
    target_z = encoder(target_feature.unsqueeze(0))
    proposals: list[Proposal] = []
    with torch.no_grad():
        for source_rank, source in enumerate(ranked_sources):
            if len(proposals) >= pool_size:
                break
            source_feature = torch.as_tensor(source.feature, dtype=torch.float32, device=device)
            source_z = encoder(source_feature.unsqueeze(0))
            for xi_source_raw in source.positive_tangents:
                if len(proposals) >= pool_size:
                    break
                xi_source = torch.as_tensor(
                    xi_source_raw, dtype=torch.float32, device=device
                ).unsqueeze(0)
                xi_source_norm = normalizer.transform(xi_source)
                xi_hat_norm = ctt(source_z, target_z, xi_source_norm)
                score = float(utility_energy(target_z, xi_hat_norm).squeeze(0).detach().cpu())
                xi_hat = normalizer.inverse_transform(xi_hat_norm).squeeze(0).detach().cpu().numpy()
                action_delta = decode_linear_keyframe_tangent(
                    xi_hat,
                    horizon=target.base_action.shape[0],
                    action_dim=target.base_action.shape[1],
                )
                action = target.base_action + float(delta_scale) * action_delta
                proposals.append(
                    Proposal(
                        tangent=xi_hat.astype(np.float32, copy=False),
                        action=action.astype(np.float32, copy=False),
                        score=score,
                        source_chart_id=source.chart_id,
                        source_task_id=source.task_id,
                        source_rank=source_rank,
                    )
                )
    proposals.sort(key=lambda proposal: proposal.score, reverse=True)
    return proposals[:k]


def _source_pool_for_target(
    target: ChartItem,
    *,
    task_pool: list[ChartItem],
    source_charts: list[ChartItem],
    exclude_self_source: bool,
) -> list[ChartItem]:
    if not exclude_self_source:
        return task_pool

    def is_not_self(source: ChartItem) -> bool:
        return source.chart_id != target.chart_id and source.state_hash != target.state_hash

    filtered = [source for source in task_pool if is_not_self(source)]
    if filtered:
        return filtered
    fallback = [source for source in source_charts if is_not_self(source)]
    return fallback or task_pool


def decode_linear_keyframe_tangent(
    tangent_code: np.ndarray,
    *,
    horizon: int,
    action_dim: int,
) -> np.ndarray:
    """Decode the public 21D CIL keyframe code into a full residual action chunk.

    The chart exporter stores start/mid/end residual rows for the common H=16, D=7
    chunk. This decoder linearly interpolates those rows; it is intentionally marked
    non-lossless in rollout metadata.
    """

    if horizon <= 0 or action_dim <= 0:
        raise ValueError("horizon and action_dim must be positive")
    code = np.asarray(tangent_code, dtype=np.float32).reshape(-1)
    key_dim = min(action_dim, max(1, min(7, code.shape[0] // 3)))
    keyframes = np.zeros((3, action_dim), dtype=np.float32)
    usable = min(3 * key_dim, code.shape[0])
    keyframes[:, :key_dim] = code[:usable].reshape(3, key_dim)
    if horizon == 1:
        return keyframes[:1]
    mid = horizon // 2
    positions = np.asarray([0, mid, horizon - 1], dtype=np.float32)
    timeline = np.arange(horizon, dtype=np.float32)
    decoded = np.zeros((horizon, action_dim), dtype=np.float32)
    for dim in range(action_dim):
        decoded[:, dim] = np.interp(timeline, positions, keyframes[:, dim])
    return decoded


def rollout_generated_cases(
    generated_cases: list[tuple[ChartItem, list[Proposal]]],
    *,
    gym: Any,
    torch: Any,
    device: str,
    group_batch_size: int,
    sim_backend: str | None,
    render_backend: str | None,
    restore_tolerance: float,
    clip_actions: bool,
    execution_action_scale: float = 1.0,
    execution_action_scale_vector: np.ndarray | None = None,
    execution_action_transform: str = "identity",
    log_path: Path | None = None,
) -> list[dict[str, Any]]:
    archives: dict[Path, dict[str, Any]] = {}
    rows: list[dict[str, Any]] = []
    by_task: dict[str, list[tuple[ChartItem, list[Proposal]]]] = {}
    for item in generated_cases:
        by_task.setdefault(item[0].task_id, []).append(item)
    for task_id, cases in sorted(by_task.items()):
        for start in range(0, len(cases), group_batch_size):
            batch = cases[start : start + group_batch_size]
            source_summary = _source_summary(batch[0][0].source_dataset)
            resolved_render_backend = (
                render_backend
                if render_backend is not None
                else source_summary.get("render_backend") or "none"
            )
            max_candidate_count = max(1 + len(proposals) for _target, proposals in batch)
            env_kwargs = {
                "num_envs": len(batch) * max_candidate_count,
                "obs_mode": "state",
                "control_mode": source_summary.get("control_mode", "pd_ee_delta_pose"),
                "render_mode": None,
                "sim_backend": sim_backend or source_summary.get("sim_backend", "physx_cuda"),
                "render_backend": resolved_render_backend,
                "reward_mode": "normalized_dense",
            }
            if _uses_single_env_cpu_backend(env_kwargs["sim_backend"]) and (
                len(batch) * max_candidate_count > 1
            ):
                rows.extend(
                    _rollout_cpu_sequential_batch(
                        task_id,
                        batch,
                        gym=gym,
                        torch=torch,
                        device=device,
                        env_kwargs=dict(env_kwargs) | {"num_envs": 1},
                        archives=archives,
                        restore_tolerance=restore_tolerance,
                        clip_actions=clip_actions,
                        execution_action_scale=execution_action_scale,
                        execution_action_scale_vector=execution_action_scale_vector,
                        execution_action_transform=execution_action_transform,
                        log_path=log_path,
                    )
                )
                continue
            _append_log(
                log_path,
                f"env init task={task_id} start={start} batch={len(batch)} "
                f"candidates={max_candidate_count} sim={env_kwargs['sim_backend']} "
                f"render={env_kwargs['render_backend']}",
            )
            env = gym.make(task_id, **env_kwargs)
            base_env = env.unwrapped
            try:
                env_device = getattr(base_env, "device", torch.device(device))
                env_dim = _env_action_dim(env)
                states: list[dict[str, Any]] = []
                action_groups: list[np.ndarray] = []
                valid_counts: list[int] = []
                for target, proposals in batch:
                    archive = archives.setdefault(
                        target.source_dataset, _load_state_archive(target.source_dataset)
                    )
                    states.append(archive["initial"][target.chart_id])
                    group_actions = [target.base_action] + [proposal.action for proposal in proposals]
                    valid_counts.append(len(group_actions))
                    while len(group_actions) < max_candidate_count:
                        group_actions.append(target.base_action)
                    action_groups.append(np.stack(group_actions, axis=0))
                candidate_values = np.stack(action_groups, axis=0).astype(np.float32)
                candidate_values = _adapt_action_dim_4d(candidate_values, env_dim)
                candidate_values = _prepare_actions_for_execution_4d(
                    candidate_values,
                    env,
                    execution_action_scale=execution_action_scale,
                    execution_action_scale_vector=execution_action_scale_vector,
                    execution_action_transform=execution_action_transform,
                )
                safety = _action_bound_diagnostics_4d(candidate_values, env)
                if clip_actions:
                    candidate_values = _clip_to_action_space_4d(candidate_values, env)
                _append_log(
                    log_path,
                    f"execute task={task_id} start={start} shape={candidate_values.shape} "
                    f"clip_actions={clip_actions} execution_action_scale={execution_action_scale} "
                    f"execution_action_transform={execution_action_transform}",
                )
                _after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
                    base_env,
                    states,
                    candidate_values,
                    torch=torch,
                    device=env_device,
                    restore_tolerance=restore_tolerance,
                )
                for index, (target, proposals) in enumerate(batch):
                    valid = valid_counts[index]
                    progress = [
                        float(max(0.0, min(1.0, rewards[index, candidate_index])))
                        for candidate_index in range(valid)
                    ]
                    success = [
                        bool(successes[index, candidate_index])
                        for candidate_index in range(valid)
                    ]
                    utilities = [
                        progress_value + (1.0 if success_value else 0.0)
                        for progress_value, success_value in zip(progress, success, strict=True)
                    ]
                    rows.append(
                        _measured_row_from_rollout(
                            target,
                            proposals,
                            progress=progress,
                            success=success,
                            utilities=utilities,
                            safety_violations=_slice_safety_labels(
                                safety,
                                group_index=index,
                                count=valid,
                            ),
                            action_clip_max_abs=_slice_clip_distances(
                                safety,
                                group_index=index,
                                count=valid,
                            ),
                            restore_error=float(restore_error),
                            execution_action_scale=execution_action_scale,
                            execution_action_scale_vector=execution_action_scale_vector,
                            execution_action_transform=execution_action_transform,
                        )
                    )
            finally:
                env.close()
            _append_log(log_path, f"batch done task={task_id} start={start}")
    return rows


def _rollout_cpu_sequential_batch(
    task_id: str,
    batch: list[tuple[ChartItem, list[Proposal]]],
    *,
    gym: Any,
    torch: Any,
    device: str,
    env_kwargs: dict[str, Any],
    archives: dict[Path, dict[str, Any]],
    restore_tolerance: float,
    clip_actions: bool,
    execution_action_scale: float,
    execution_action_scale_vector: np.ndarray | None,
    execution_action_transform: str,
    log_path: Path | None,
) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    _append_log(
        log_path,
        f"env init sequential task={task_id} batch={len(batch)} sim={env_kwargs['sim_backend']} "
        f"render={env_kwargs['render_backend']}",
    )
    env = gym.make(task_id, **env_kwargs)
    base_env = env.unwrapped
    try:
        env_device = getattr(base_env, "device", torch.device(device))
        env_dim = _env_action_dim(env)
        for target, proposals in batch:
            archive = archives.setdefault(
                target.source_dataset, _load_state_archive(target.source_dataset)
            )
            state = archive["initial"][target.chart_id]
            candidate_actions = [target.base_action] + [proposal.action for proposal in proposals]
            progress: list[float] = []
            success: list[bool] = []
            safety_violations: list[bool | None] = []
            action_clip_max_abs: list[float | None] = []
            restore_errors: list[float] = []
            for candidate_index, action in enumerate(candidate_actions):
                candidate_values = np.asarray(action, dtype=np.float32).reshape(
                    1, 1, *action.shape
                )
                candidate_values = _adapt_action_dim_4d(candidate_values, env_dim)
                candidate_values = _prepare_actions_for_execution_4d(
                    candidate_values,
                    env,
                    execution_action_scale=execution_action_scale,
                    execution_action_scale_vector=execution_action_scale_vector,
                    execution_action_transform=execution_action_transform,
                )
                safety = _action_bound_diagnostics_4d(candidate_values, env)
                if clip_actions:
                    candidate_values = _clip_to_action_space_4d(candidate_values, env)
                _append_log(
                    log_path,
                    f"execute sequential task={task_id} chart={target.chart_id} "
                    f"candidate={candidate_index} shape={candidate_values.shape} "
                    f"clip_actions={clip_actions} execution_action_scale={execution_action_scale} "
                    f"execution_action_transform={execution_action_transform}",
                )
                _after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
                    base_env,
                    [state],
                    candidate_values,
                    torch=torch,
                    device=env_device,
                    restore_tolerance=restore_tolerance,
                )
                progress.append(float(max(0.0, min(1.0, rewards[0, 0]))))
                success.append(bool(successes[0, 0]))
                safety_violations.append(_single_safety_label(safety))
                action_clip_max_abs.append(_single_clip_distance(safety))
                restore_errors.append(float(restore_error))
            utilities = [
                progress_value + (1.0 if success_value else 0.0)
                for progress_value, success_value in zip(progress, success, strict=True)
            ]
            rows.append(
                _measured_row_from_rollout(
                    target,
                    proposals,
                    progress=progress,
                    success=success,
                    utilities=utilities,
                    safety_violations=safety_violations,
                    action_clip_max_abs=action_clip_max_abs,
                    restore_error=max(restore_errors, default=0.0),
                    execution_action_scale=execution_action_scale,
                    execution_action_scale_vector=execution_action_scale_vector,
                    execution_action_transform=execution_action_transform,
                )
            )
    finally:
        env.close()
    _append_log(log_path, f"batch done sequential task={task_id}")
    return rows


def _measured_row_from_rollout(
    target: ChartItem,
    proposals: list[Proposal],
    *,
    progress: list[float],
    success: list[bool],
    utilities: list[float],
    safety_violations: list[bool | None] | None = None,
    action_clip_max_abs: list[float | None] | None = None,
    restore_error: float,
    execution_action_scale: float = 1.0,
    execution_action_scale_vector: np.ndarray | None = None,
    execution_action_transform: str = "identity",
) -> dict[str, Any]:
    if safety_violations is None:
        safety_violations = [None] * len(utilities)
    if action_clip_max_abs is None:
        action_clip_max_abs = [None] * len(utilities)
    if len(safety_violations) != len(utilities):
        raise ValueError("safety_violations must match utilities length")
    if len(action_clip_max_abs) != len(utilities):
        raise ValueError("action_clip_max_abs must match utilities length")
    base_utility = float(utilities[0])
    generated_utilities = [float(value) for value in utilities[1:]]
    predicted_scores = [float(proposal.score) for proposal in proposals]
    return {
        "chart_id": target.chart_id,
        "group_id": target.chart_id,
        "task_id": target.task_id,
        "seed": target.seed,
        "state_hash": target.state_hash,
        "instruction": target.instruction,
        "candidates_evaluated": True,
        "execution_action_scale": float(execution_action_scale),
        "execution_action_scale_vector": (
            None
            if execution_action_scale_vector is None
            else execution_action_scale_vector.astype(float).tolist()
        ),
        "execution_action_transform": str(execution_action_transform),
        "selected_index": 0,
        "base_utility": base_utility,
        "stored_base_utility": target.stored_base_utility,
        "generated_utilities": generated_utilities,
        "hidden_chart_utilities": target.hidden_utilities,
        "hidden_candidate_types": target.hidden_candidate_types,
        "outcome_vector_schema": [
            "success",
            "progress",
            "contact_quality",
            "safety_violation",
            "stage_progress",
            "smoothness",
            "recovery",
            "utility_scalar",
        ],
        "base_outcome": _outcome_payload(
            success=success[0],
            progress=progress[0],
            utility=utilities[0],
            safety_violation=safety_violations[0],
            action_clip_max_abs=action_clip_max_abs[0],
        ),
        "candidate_outcomes": [
            _outcome_payload(
                success=success_value,
                progress=progress_value,
                utility=utility,
                safety_violation=safety_violation,
                action_clip_max_abs=clip_distance,
            )
            for success_value, progress_value, utility, safety_violation, clip_distance in zip(
                success[1:],
                progress[1:],
                utilities[1:],
                safety_violations[1:],
                action_clip_max_abs[1:],
                strict=True,
            )
        ],
        "predicted_scores": predicted_scores,
        "generated_tangents": [
            proposal.tangent.astype(float).tolist() for proposal in proposals
        ],
        "positive_tangents": target.positive_tangents.astype(float).tolist(),
        "negative_tangents": target.negative_tangents.astype(float).tolist(),
        "candidate_types": [
            f"ctt_transport_rank{proposal.source_rank}" for proposal in proposals
        ],
        "candidate_source_chart_ids": [proposal.source_chart_id for proposal in proposals],
        "candidate_source_task_ids": [proposal.source_task_id for proposal in proposals],
        "candidate_progress": progress[1:],
        "candidate_success": success[1:],
        "candidate_safety_violation": safety_violations[1:],
        "candidate_action_clip_max_abs": action_clip_max_abs[1:],
        "base_progress": progress[0],
        "base_success": success[0],
        "base_safety_violation": safety_violations[0],
        "base_action_clip_max_abs": action_clip_max_abs[0],
        "restore_error": float(restore_error),
        "num_generated": len(proposals),
        "num_executed_including_base": len(utilities),
    }


def _outcome_payload(
    *,
    success: bool,
    progress: float,
    utility: float,
    safety_violation: bool | None,
    action_clip_max_abs: float | None,
) -> dict[str, Any]:
    action_bound_violation = None if safety_violation is None else bool(safety_violation)
    return {
        "success": bool(success),
        "progress": float(progress),
        "contact_quality": None,
        "safety_violation": action_bound_violation,
        "safety_violation_source": (
            "action_bounds" if action_bound_violation is not None else None
        ),
        "action_bound_violation": action_bound_violation,
        "action_clip_max_abs": (
            None if action_clip_max_abs is None else float(action_clip_max_abs)
        ),
        "stage_progress": None,
        "smoothness": None,
        "recovery": None,
        "utility_scalar": float(utility),
    }


def _uses_single_env_cpu_backend(sim_backend: Any) -> bool:
    value = str(sim_backend or "").lower()
    return value in {"cpu", "physx_cpu"} or value.endswith("_cpu")


def _register_required_maniskill_envs(log_path: Path | None = None) -> None:
    """Register only the tabletop ManiSkill tasks used by the CIL diagnostic.

    ManiSkill 3.0.1 imports all task families from its top-level package,
    including digital-twin Bridge tasks that import cv2. On the current HPC
    container that full import can hang before our five tabletop tasks are
    registered. A lightweight package stub lets Python resolve ManiSkill
    submodules without executing the top-level eager task import.
    """

    package = sys.modules.get("mani_skill")
    if package is None:
        spec = importlib.machinery.PathFinder.find_spec("mani_skill")
        if spec is None or not spec.submodule_search_locations:
            raise ImportError("Could not locate installed mani_skill package")
        package_dir = Path(list(spec.submodule_search_locations)[0]).resolve()
        package = types.ModuleType("mani_skill")
        package.__file__ = str(package_dir / "__init__.py")
        package.__path__ = [str(package_dir)]
        package.__package__ = "mani_skill"
        package.__version__ = "3.0.1"
        package.PACKAGE_DIR = package_dir
        package.PACKAGE_ASSET_DIR = package_dir / "assets"
        asset_root = Path(
            os.getenv("MS_ASSET_DIR", os.path.join(os.path.expanduser("~"), ".maniskill"))
        )
        package.ASSET_DIR = asset_root / "data"
        package.DEMO_DIR = asset_root / "demos"
        package.format_path = lambda p: p.format(
            PACKAGE_DIR=package.PACKAGE_DIR,
            PACKAGE_ASSET_DIR=package.PACKAGE_ASSET_DIR,
            ASSET_DIR=package.ASSET_DIR,
        )
        sys.modules["mani_skill"] = package
        _install_package_stub("mani_skill.envs", package_dir / "envs")
        _install_package_stub("mani_skill.envs.tasks", package_dir / "envs" / "tasks")
        _install_package_stub(
            "mani_skill.envs.tasks.tabletop",
            package_dir / "envs" / "tasks" / "tabletop",
        )
        package.logger = importlib.import_module("mani_skill.utils.logging_utils").logger
        _append_log(log_path, "installed lightweight mani_skill tabletop import stub")

    modules = [
        "mani_skill.envs.tasks.tabletop.pick_cube",
        "mani_skill.envs.tasks.tabletop.pull_cube",
        "mani_skill.envs.tasks.tabletop.push_cube",
        "mani_skill.envs.tasks.tabletop.stack_cube",
        "mani_skill.envs.tasks.tabletop.lift_peg_upright",
    ]
    for module_name in modules:
        importlib.import_module(module_name)
    _append_log(log_path, "registered tabletop ManiSkill tasks selectively")


def _install_package_stub(name: str, path: Path) -> None:
    if name in sys.modules:
        return
    module = types.ModuleType(name)
    module.__path__ = [str(path)]
    module.__package__ = name
    sys.modules[name] = module


def _validate_indexes(
    source_path: Path,
    source_index: dict[str, Any],
    target_path: Path,
    target_index: dict[str, Any],
) -> None:
    if source_index.get("split") != "train" or not source_index.get("retrieval_index_allowed"):
        raise SystemExit(
            f"{source_path} is not a train-only retrieval index; CTT rollout sources "
            "must come from train split only"
        )
    if not source_index.get("include_outcomes"):
        raise SystemExit(f"{source_path} must include train outcomes for source positives")
    if not target_index.get("include_outcomes"):
        raise SystemExit(f"{target_path} must include evaluator-only outcomes")
    if target_index.get("split") != "train" and target_index.get("retrieval_index_allowed"):
        raise SystemExit(f"{target_path} is non-train but marked retrieval_index_allowed")


def _adapt_action_dim_4d(actions: np.ndarray, action_dim: int) -> np.ndarray:
    if actions.ndim != 4:
        raise ValueError("actions must have shape [B,K,H,D]")
    if actions.shape[-1] == action_dim:
        return actions.astype(np.float32, copy=False)
    if actions.shape[-1] > action_dim:
        return actions[..., :action_dim].astype(np.float32, copy=False)
    pad = np.zeros((*actions.shape[:-1], action_dim - actions.shape[-1]), dtype=np.float32)
    return np.concatenate([actions.astype(np.float32, copy=False), pad], axis=-1)


def _parse_execution_action_scale_vector(raw: str) -> np.ndarray | None:
    text = str(raw or "").strip()
    if not text:
        return None
    values = [float(item.strip()) for item in text.split(",") if item.strip()]
    if not values:
        return None
    vector = np.asarray(values, dtype=np.float32)
    if np.any(~np.isfinite(vector)) or np.any(vector <= 0.0):
        raise ValueError("--execution-action-scale-vector values must be positive and finite")
    return vector


def _prepare_actions_for_execution_4d(
    actions: np.ndarray,
    env: Any,
    *,
    execution_action_scale: float,
    execution_action_scale_vector: np.ndarray | None,
    execution_action_transform: str,
) -> np.ndarray:
    scaled = _scale_actions_4d(
        actions,
        execution_action_scale,
        scale_vector=execution_action_scale_vector,
    )
    return _transform_actions_4d(scaled, env, execution_action_transform)


def _scale_actions_4d(
    actions: np.ndarray,
    scale: float,
    *,
    scale_vector: np.ndarray | None = None,
) -> np.ndarray:
    if not math.isfinite(float(scale)) or float(scale) <= 0.0:
        raise ValueError("execution action scale must be positive and finite")
    output = actions.astype(np.float32, copy=False)
    if float(scale) != 1.0:
        output = output * float(scale)
    if scale_vector is not None:
        vector = np.asarray(scale_vector, dtype=np.float32).reshape(-1)
        if vector.size < output.shape[-1]:
            pad = np.ones(output.shape[-1] - vector.size, dtype=np.float32)
            vector = np.concatenate([vector, pad], axis=0)
        elif vector.size > output.shape[-1]:
            vector = vector[: output.shape[-1]]
        if np.any(~np.isfinite(vector)) or np.any(vector <= 0.0):
            raise ValueError("execution action scale vector must be positive and finite")
        output = output * vector.reshape(1, 1, 1, -1)
    return output.astype(np.float32, copy=False)


def _transform_actions_4d(actions: np.ndarray, env: Any, transform: str) -> np.ndarray:
    value = str(transform or "identity").strip().lower()
    if value == "identity":
        return actions.astype(np.float32, copy=False)
    if value == "env_clip":
        return _clip_to_action_space_4d(actions.astype(np.float32, copy=False), env)
    if value != "tanh":
        raise ValueError(f"Unknown execution action transform: {transform}")
    squashed = np.tanh(actions.astype(np.float32, copy=False))
    bounds = _action_space_bounds(actions, env)
    if bounds is None:
        return squashed.astype(np.float32, copy=False)
    low_arr, high_arr = bounds
    center = ((low_arr + high_arr) * 0.5).reshape(1, 1, 1, -1)
    half_range = ((high_arr - low_arr) * 0.5).reshape(1, 1, 1, -1)
    return (center + half_range * squashed).astype(np.float32, copy=False)


def _clip_to_action_space_4d(actions: np.ndarray, env: Any) -> np.ndarray:
    bounds = _action_space_bounds(actions, env)
    if bounds is None:
        return actions
    low_arr, high_arr = bounds
    return np.clip(actions, low_arr.reshape(1, 1, 1, -1), high_arr.reshape(1, 1, 1, -1))


def _action_bound_diagnostics_4d(actions: np.ndarray, env: Any) -> dict[str, np.ndarray] | None:
    """Measure action-space bound violations before clipping.

    This is a validity/safety label source, not a collision/contact label.
    """

    bounds = _action_space_bounds(actions, env)
    if bounds is None:
        return None
    low_arr, high_arr = bounds
    low_view = low_arr.reshape(1, 1, 1, -1)
    high_view = high_arr.reshape(1, 1, 1, -1)
    below = np.maximum(low_view - actions, 0.0)
    above = np.maximum(actions - high_view, 0.0)
    excess = np.maximum(below, above)
    max_abs = np.max(excess, axis=(2, 3))
    return {
        "violation": max_abs > 1.0e-6,
        "max_abs": max_abs.astype(np.float32, copy=False),
    }


def _action_space_bounds(actions: np.ndarray, env: Any) -> tuple[np.ndarray, np.ndarray] | None:
    space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None)
    low = getattr(space, "low", None)
    high = getattr(space, "high", None)
    if low is None or high is None:
        return None
    action_dim = int(actions.shape[-1])
    low_arr = _bound_vector(low, action_dim)
    high_arr = _bound_vector(high, action_dim)
    if low_arr is None or high_arr is None:
        return None
    if np.any(~np.isfinite(low_arr)) or np.any(~np.isfinite(high_arr)):
        return None
    return low_arr, high_arr


def _bound_vector(value: Any, action_dim: int) -> np.ndarray | None:
    array = np.asarray(value, dtype=np.float32).reshape(-1)
    if array.size == 1:
        return np.repeat(array, action_dim).astype(np.float32, copy=False)
    if array.size < action_dim:
        return None
    return array[-action_dim:].astype(np.float32, copy=False)


def _slice_safety_labels(
    safety: dict[str, np.ndarray] | None,
    *,
    group_index: int,
    count: int,
) -> list[bool | None]:
    if safety is None:
        return [None] * count
    values = safety["violation"][group_index, :count]
    return [bool(value) for value in values.tolist()]


def _slice_clip_distances(
    safety: dict[str, np.ndarray] | None,
    *,
    group_index: int,
    count: int,
) -> list[float | None]:
    if safety is None:
        return [None] * count
    values = safety["max_abs"][group_index, :count]
    return [float(value) for value in values.tolist()]


def _single_safety_label(safety: dict[str, np.ndarray] | None) -> bool | None:
    if safety is None:
        return None
    return bool(safety["violation"][0, 0])


def _single_clip_distance(safety: dict[str, np.ndarray] | None) -> float | None:
    if safety is None:
        return None
    return float(safety["max_abs"][0, 0])


def _env_action_dim(env: Any) -> int:
    for space in (
        getattr(env, "single_action_space", None),
        getattr(env.unwrapped, "single_action_space", None),
        getattr(env, "action_space", None),
    ):
        shape = getattr(space, "shape", None)
        if shape:
            return int(shape[-1])
    raise ValueError("Could not infer ManiSkill action dimension from environment")


def _load_state_archive(source_dataset: Path) -> dict[str, Any]:
    archive_path = source_dataset / "state_archive.pkl"
    if not archive_path.exists():
        summary_path = source_dataset / "generation_summary.json"
        if summary_path.exists():
            raw_path = read_json(summary_path).get("state_archive")
            if raw_path:
                archive_path = Path(str(raw_path))
    if not archive_path.exists():
        raise FileNotFoundError(f"ManiSkill state archive not found for {source_dataset}")
    with archive_path.open("rb") as handle:
        archive = pickle.load(handle)
    if not isinstance(archive, dict) or "initial" not in archive:
        raise ValueError(f"Invalid ManiSkill state archive: {archive_path}")
    return archive


def _source_summary(source_dataset: Path) -> dict[str, Any]:
    path = source_dataset / "generation_summary.json"
    return read_json(path) if path.exists() else {}


def _source_dataset_from_metadata(
    metadata: dict[str, Any],
    *,
    index: dict[str, Any],
    task_id: str,
) -> Path:
    raw = metadata.get("source_dataset")
    if raw:
        return Path(str(raw))
    dataset = index.get("dataset")
    if dataset:
        candidate = Path(str(dataset)) / task_id
        if candidate.exists():
            return candidate
    raise ValueError(f"Could not resolve source_dataset for chart task {task_id}")


def _matrix_or_empty(items: list[np.ndarray], *, width: int) -> np.ndarray:
    if not items:
        return np.zeros((0, width), dtype=np.float32)
    return np.asarray(items, dtype=np.float32)


def _json_loads(raw: str) -> dict[str, Any]:
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError:
        return {}
    return payload if isinstance(payload, dict) else {}


def _resolve_device(device: str) -> str:
    if device != "auto":
        return device
    return "cuda" if torch.cuda.is_available() else "cpu"


def _write_run_provenance(out_dir: Path, args: argparse.Namespace) -> None:
    config = {key: str(value) for key, value in sorted(vars(args).items())}
    (out_dir / "config.yaml").write_text(
        "\n".join(f"{key}: {value}" for key, value in config.items()) + "\n"
    )
    (out_dir / "command.txt").write_text(
        "python scripts/eval_ctt_generated_rollout.py " + " ".join(sys.argv[1:]) + "\n"
    )
    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
    for name, path in {
        "source_index.json": args.source_index,
        "target_index.json": args.target_index,
    }.items():
        try:
            (out_dir / name).write_text(Path(path).read_text())
        except OSError:
            pass


def _write_required_artifacts(
    out_dir: Path,
    payload: dict[str, Any],
    *,
    source_index: dict[str, Any],
    target_index: dict[str, Any],
    metrics_dir: Path | None,
) -> None:
    (out_dir / "data_hash.txt").write_text(str(target_index.get("content_hash", "")) + "\n")
    (out_dir / "split_hash.txt").write_text(str(target_index.get("split_hash", "")) + "\n")
    (out_dir / "source_data_hash.txt").write_text(str(source_index.get("content_hash", "")) + "\n")
    (out_dir / "source_split_hash.txt").write_text(str(source_index.get("split_hash", "")) + "\n")
    (out_dir / "train.log").write_text(
        "CTT checkpoint trained separately; rollout used checkpoint="
        f"{payload.get('checkpoint', '')}\n"
    )
    run_log = out_dir / "run.log"
    (out_dir / "eval.log").write_text(run_log.read_text() if run_log.exists() else "eval log unavailable\n")
    summary = {
        "report_type": payload.get("report_type"),
        "schema_version": payload.get("schema_version"),
        "k": payload.get("k"),
        "checkpoint": payload.get("checkpoint"),
        "source_content_hash": payload.get("source_content_hash"),
        "source_split_hash": payload.get("source_split_hash"),
        "target_content_hash": payload.get("target_content_hash"),
        "target_split_hash": payload.get("target_split_hash"),
        "num_rows": len(payload.get("rows", [])),
        "success_summary": _success_summary(payload.get("rows", []), k=int(payload.get("k", 0))),
    }
    if metrics_dir is not None:
        metric_path = metrics_dir / "metrics.json"
        if metric_path.exists():
            metrics = json.loads(metric_path.read_text())
            summary["measured_metric_summary"] = metrics.get("summary", {})
        for filename in ("metrics_by_task.json", "metrics_by_seed.json", "table.tex"):
            src = metrics_dir / filename
            if src.exists():
                shutil.copyfile(src, out_dir / filename)
    (out_dir / "metrics.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
    for filename in ("metrics_by_task.json", "metrics_by_seed.json", "table.tex"):
        path = out_dir / filename
        if not path.exists():
            path.write_text("{}\n" if filename.endswith(".json") else "% metrics not computed\n")


def _success_summary(rows: list[dict[str, Any]], *, k: int) -> dict[str, Any]:
    if k <= 0:
        k = 10**9
    base_success = []
    selected_success = []
    oracle_success = []
    base_utility = []
    selected_utility = []
    oracle_utility = []
    hidden_oracle_utility = []
    hidden_oracle_success = []
    success_support_gap = []
    success_selector_gap = []
    selected_success_gain = []
    proposal_oracle_success_gain = []
    restore_errors = []
    for row in rows:
        generated_utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
        generated_success = [bool(value) for value in row.get("candidate_success", [])[:k]]
        selected_index = int(row.get("selected_index", 0))
        selected_success_value: float | None = None
        proposal_oracle_success_value: float | None = None
        base_success_value: float | None = None
        if "base_success" in row:
            base_success_value = float(bool(row["base_success"]))
            base_success.append(base_success_value)
        if selected_index < len(generated_success):
            selected_success_value = float(generated_success[selected_index])
            selected_success.append(selected_success_value)
        if generated_success:
            proposal_oracle_success_value = float(any(generated_success))
            oracle_success.append(proposal_oracle_success_value)
        if "base_utility" in row:
            base_utility.append(float(row["base_utility"]))
        if selected_index < len(generated_utilities):
            selected_utility.append(float(generated_utilities[selected_index]))
        if generated_utilities:
            oracle_utility.append(max(generated_utilities))
        hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
        if hidden:
            hidden_oracle_utility.append(max(hidden))
            hidden_success_value = float(any(value >= 1.0 for value in hidden))
            hidden_oracle_success.append(hidden_success_value)
            if proposal_oracle_success_value is not None:
                success_support_gap.append(
                    max(0.0, hidden_success_value - proposal_oracle_success_value)
                )
            if (
                proposal_oracle_success_value is not None
                and selected_success_value is not None
            ):
                success_selector_gap.append(
                    max(0.0, proposal_oracle_success_value - selected_success_value)
                )
        if base_success_value is not None and selected_success_value is not None:
            selected_success_gain.append(selected_success_value - base_success_value)
        if base_success_value is not None and proposal_oracle_success_value is not None:
            proposal_oracle_success_gain.append(
                proposal_oracle_success_value - base_success_value
            )
        if "restore_error" in row:
            restore_errors.append(float(row["restore_error"]))
    return {
        "base_success_rate": _mean(base_success),
        "selected_success_rate": _mean(selected_success),
        "proposal_oracle_success_rate": _mean(oracle_success),
        "hidden_chart_oracle_success_rate": _mean(hidden_oracle_success),
        "selected_success_gain_over_base": _mean(selected_success_gain),
        "proposal_oracle_success_gain_over_base": _mean(proposal_oracle_success_gain),
        "success_support_gap": _mean(success_support_gap),
        "success_selector_gap": _mean(success_selector_gap),
        "base_utility_mean": _mean(base_utility),
        "selected_utility_mean": _mean(selected_utility),
        "proposal_oracle_utility_mean": _mean(oracle_utility),
        "hidden_chart_oracle_utility_mean": _mean(hidden_oracle_utility),
        "max_restore_error": max(restore_errors) if restore_errors else None,
    }


def _append_log(path: Path | None, message: str) -> None:
    if path is None:
        return
    with path.open("a") as handle:
        handle.write(message + "\n")


def _run(command: list[str]) -> str:
    try:
        return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        return ""


def _report(payload: dict[str, Any]) -> str:
    rows = payload["rows"]
    outcome = [
        1.0
        if any(
            float(value) > float(row["base_utility"])
            for value in row.get("generated_utilities", [])[: int(payload["k"])]
        )
        else 0.0
        for row in rows
    ]
    selector_regrets = []
    support_gaps = []
    for row in rows:
        generated = row.get("generated_utilities", [])[: int(payload["k"])]
        if generated:
            selector_regrets.append(max(generated) - generated[int(row.get("selected_index", 0))])
            hidden = row.get("hidden_chart_utilities", [])
            if hidden:
                support_gaps.append(max(hidden) - max(generated))
    lines = [
        "# CTT Generated Measured Rollout",
        "",
        f"Rows: `{len(rows)}`",
        f"K: `{payload['k']}`",
        f"Checkpoint: `{payload['checkpoint']}`",
        "",
        "| Metric | Mean |",
        "| --- | ---: |",
        f"| OutcomePTR@K | {_mean(outcome):.4f} |",
        f"| SelectorRegret@K | {_mean(selector_regrets):.4f} |",
        f"| SupportGap@K | {_mean(support_gaps):.4f} |",
        "",
        "These are measured simulator rollouts of decoded CTT candidates, not PPTC proxies.",
    ]
    return "\n".join(lines)


def _mean(values: list[float]) -> float:
    clean = [float(value) for value in values if math.isfinite(float(value))]
    return sum(clean) / len(clean) if clean else math.nan


if __name__ == "__main__":
    raise SystemExit(main())