scripts/eval_ctt_generated_rollout.py

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

import argparse
import json
import math
import pickle
import shutil
import subprocess
import sys
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("--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)
    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")

    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
        import mani_skill  # noqa: F401 - importing registers environments
    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,
        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,
        },
        "rows": rows,
    }
    measured_path = out_dir / "measured_candidates.json"
    measured_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    (out_dir / "report.md").write_text(_report(payload) + "\n")

    metrics_dir = out_dir / "measured_metrics"
    if not args.skip_metrics:
        eval_metrics_main(
            [
                "--input",
                str(measured_path),
                "--out-dir",
                str(metrics_dir),
                "--mode",
                "measured",
                "--k",
                str(args.k),
                "--bootstrap-samples",
                str(args.bootstrap_samples),
            ]
        )
    _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,
    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,
                        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 = _clip_to_action_space_4d(candidate_values, env)
                _append_log(
                    log_path,
                    f"execute task={task_id} start={start} shape={candidate_values.shape}",
                )
                _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,
                            restore_error=float(restore_error),
                        )
                    )
            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,
    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] = []
            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 = _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}",
                )
                _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]))
                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,
                    restore_error=max(restore_errors, default=0.0),
                )
            )
    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],
    restore_error: float,
) -> dict[str, Any]:
    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,
        "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],
        ),
        "candidate_outcomes": [
            _outcome_payload(
                success=success_value,
                progress=progress_value,
                utility=utility,
            )
            for success_value, progress_value, utility in zip(
                success[1:],
                progress[1:],
                utilities[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:],
        "base_progress": progress[0],
        "base_success": success[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) -> dict[str, Any]:
    return {
        "success": bool(success),
        "progress": float(progress),
        "contact_quality": None,
        "safety_violation": None,
        "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 _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 _clip_to_action_space_4d(actions: np.ndarray, env: Any) -> np.ndarray:
    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 actions
    low_arr = np.asarray(low, dtype=np.float32).reshape(-1)[-actions.shape[-1] :]
    high_arr = np.asarray(high, dtype=np.float32).reshape(-1)[-actions.shape[-1] :]
    if low_arr.shape[0] != actions.shape[-1] or high_arr.shape[0] != actions.shape[-1]:
        return actions
    return np.clip(actions, low_arr.reshape(1, 1, 1, -1), high_arr.reshape(1, 1, 1, -1))


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())