#!/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())