anhtld commited on
Commit
c6553a4
·
verified ·
1 Parent(s): 8465bae

ctt train-calibration hygiene 2026-07-03T16:31:19Z: scripts/eval_ctt_generated_rollout.py

Browse files
Files changed (1) hide show
  1. scripts/eval_ctt_generated_rollout.py +1112 -0
scripts/eval_ctt_generated_rollout.py ADDED
@@ -0,0 +1,1112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ import math
7
+ import pickle
8
+ import shutil
9
+ import subprocess
10
+ import sys
11
+ from dataclasses import dataclass
12
+ from pathlib import Path
13
+ from typing import Any
14
+
15
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
16
+ if str(PROJECT_ROOT) not in sys.path:
17
+ sys.path.insert(0, str(PROJECT_ROOT))
18
+
19
+ import numpy as np # noqa: E402
20
+ import torch # noqa: E402
21
+
22
+ from cil.chart_features import build_chart_feature # noqa: E402
23
+ from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer, UtilityEnergy # noqa: E402
24
+ from dovla_cil.generation.maniskill_parallel import execute_grouped_action_lattice_batch # noqa: E402
25
+ from dovla_cil.utils.io import read_json # noqa: E402
26
+ from scripts.eval_metrics import main as eval_metrics_main # noqa: E402
27
+
28
+
29
+ @dataclass(frozen=True)
30
+ class ChartItem:
31
+ chart_id: str
32
+ task_id: str
33
+ seed: str
34
+ state_hash: str
35
+ instruction: str
36
+ source_dataset: Path
37
+ base_action: np.ndarray
38
+ feature: np.ndarray
39
+ positive_tangents: np.ndarray
40
+ negative_tangents: np.ndarray
41
+ hidden_utilities: list[float]
42
+ hidden_candidate_types: list[str]
43
+ stored_base_utility: float | None
44
+
45
+
46
+ @dataclass(frozen=True)
47
+ class Proposal:
48
+ tangent: np.ndarray
49
+ action: np.ndarray
50
+ score: float
51
+ source_chart_id: str
52
+ source_task_id: str
53
+ source_rank: int
54
+
55
+
56
+ def main(argv: list[str] | None = None) -> int:
57
+ parser = argparse.ArgumentParser(
58
+ description=(
59
+ "Generate CTT candidates, decode them to ManiSkill action chunks, "
60
+ "and measure them with same-state simulator rollouts."
61
+ )
62
+ )
63
+ parser.add_argument("--checkpoint", type=Path, required=True)
64
+ parser.add_argument("--source-index", type=Path, default=Path("data/cil_charts/train/index.json"))
65
+ parser.add_argument("--target-index", type=Path, default=Path("data/cil_charts/val/index.json"))
66
+ parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_residual_rollout_smoke"))
67
+ parser.add_argument("--k", type=int, default=16)
68
+ parser.add_argument("--pool-size", type=int, default=0)
69
+ parser.add_argument("--neighbors", type=int, default=8)
70
+ parser.add_argument("--max-target-charts", type=int, default=8)
71
+ parser.add_argument("--group-batch-size", type=int, default=1)
72
+ parser.add_argument("--device", default="auto")
73
+ parser.add_argument("--sim-backend", default=None)
74
+ parser.add_argument("--render-backend", default=None)
75
+ parser.add_argument("--restore-tolerance", type=float, default=1.0e-5)
76
+ parser.add_argument("--delta-scale", type=float, default=1.0)
77
+ parser.add_argument("--include-targets-without-positives", action="store_true")
78
+ parser.add_argument(
79
+ "--exclude-self-source",
80
+ action="store_true",
81
+ help=(
82
+ "When source and target indexes overlap, exclude source charts with the "
83
+ "same chart_id or state_hash as the target. Use this for train-split "
84
+ "calibration rollouts so retrieval cannot copy the target chart's own positives."
85
+ ),
86
+ )
87
+ parser.add_argument("--skip-metrics", action="store_true")
88
+ parser.add_argument("--bootstrap-samples", type=int, default=200)
89
+ args = parser.parse_args(argv)
90
+
91
+ if args.k <= 0:
92
+ parser.error("--k must be positive")
93
+ if args.neighbors <= 0:
94
+ parser.error("--neighbors must be positive")
95
+ if args.group_batch_size <= 0:
96
+ parser.error("--group-batch-size must be positive")
97
+ if args.max_target_charts <= 0:
98
+ parser.error("--max-target-charts must be positive")
99
+ if args.restore_tolerance <= 0.0:
100
+ parser.error("--restore-tolerance must be positive")
101
+
102
+ out_dir = args.out_dir
103
+ out_dir.mkdir(parents=True, exist_ok=True)
104
+ _write_run_provenance(out_dir, args)
105
+ log_path = out_dir / "run.log"
106
+ _append_log(log_path, "start")
107
+ _append_log(log_path, "importing gymnasium/mani_skill")
108
+ try:
109
+ import gymnasium as gym
110
+ import mani_skill # noqa: F401 - importing registers environments
111
+ except ImportError as exc: # pragma: no cover - exercised in the Apptainer env
112
+ raise ImportError(
113
+ "CTT measured rollout requires gymnasium, mani_skill, numpy, and torch. "
114
+ "Run this script through the ManiSkill Apptainer environment on HPC."
115
+ ) from exc
116
+ _append_log(log_path, "imported gymnasium/mani_skill")
117
+
118
+ checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
119
+ config = CTTConfig(**checkpoint["config"])
120
+ chart_feature_mode = str(checkpoint.get("chart_feature_mode", "base"))
121
+ encoder = ChartEncoder(config.chart_feature_dim, output_dim=config.chart_dim)
122
+ ctt = CausalTangentTransport(config)
123
+ utility_energy = UtilityEnergy(chart_dim=config.chart_dim, tangent_dim=config.tangent_dim)
124
+ encoder.load_state_dict(checkpoint["chart_encoder"])
125
+ ctt.load_state_dict(checkpoint["ctt"])
126
+ if "utility_energy" not in checkpoint:
127
+ raise SystemExit(f"{args.checkpoint} does not contain a utility_energy state")
128
+ utility_energy.load_state_dict(checkpoint["utility_energy"])
129
+ normalizer = TangentNormalizer.from_dict(checkpoint["normalizer"])
130
+ encoder.eval()
131
+ ctt.eval()
132
+ utility_energy.eval()
133
+ for module in (encoder, ctt, utility_energy):
134
+ for parameter in module.parameters():
135
+ parameter.requires_grad_(False)
136
+ _append_log(log_path, f"loaded checkpoint={args.checkpoint}")
137
+
138
+ source_charts, source_index = load_chart_items(
139
+ args.source_index,
140
+ max_charts=None,
141
+ require_positive=True,
142
+ include_hidden=False,
143
+ include_metadata=True,
144
+ chart_feature_mode=chart_feature_mode,
145
+ )
146
+ target_charts, target_index = load_chart_items(
147
+ args.target_index,
148
+ max_charts=args.max_target_charts,
149
+ require_positive=not args.include_targets_without_positives,
150
+ include_hidden=True,
151
+ include_metadata=True,
152
+ chart_feature_mode=chart_feature_mode,
153
+ )
154
+ _validate_indexes(args.source_index, source_index, args.target_index, target_index)
155
+ if not target_charts:
156
+ raise SystemExit("No target charts available after filtering")
157
+ _append_log(
158
+ log_path,
159
+ f"loaded charts source={len(source_charts)} target={len(target_charts)}",
160
+ )
161
+
162
+ resolved_device = _resolve_device(args.device)
163
+ encoder.to(resolved_device)
164
+ ctt.to(resolved_device)
165
+ utility_energy.to(resolved_device)
166
+ source_by_task: dict[str, list[ChartItem]] = {}
167
+ for chart in source_charts:
168
+ source_by_task.setdefault(chart.task_id, []).append(chart)
169
+
170
+ pool_size = int(args.pool_size) if args.pool_size > 0 else int(args.k)
171
+ generated_cases = [
172
+ (
173
+ target,
174
+ generate_proposals(
175
+ target,
176
+ source_charts=source_charts,
177
+ source_by_task=source_by_task,
178
+ encoder=encoder,
179
+ ctt=ctt,
180
+ utility_energy=utility_energy,
181
+ normalizer=normalizer,
182
+ device=resolved_device,
183
+ neighbors=args.neighbors,
184
+ pool_size=max(pool_size, args.k),
185
+ k=args.k,
186
+ delta_scale=args.delta_scale,
187
+ exclude_self_source=args.exclude_self_source,
188
+ ),
189
+ )
190
+ for target in target_charts
191
+ ]
192
+ _append_log(
193
+ log_path,
194
+ "generated proposals "
195
+ f"rows={len(generated_cases)} total={sum(len(item[1]) for item in generated_cases)}",
196
+ )
197
+ rows = rollout_generated_cases(
198
+ generated_cases,
199
+ gym=gym,
200
+ torch=torch,
201
+ device=resolved_device,
202
+ group_batch_size=args.group_batch_size,
203
+ sim_backend=args.sim_backend,
204
+ render_backend=args.render_backend,
205
+ restore_tolerance=args.restore_tolerance,
206
+ log_path=log_path,
207
+ )
208
+ _append_log(log_path, f"rollout complete rows={len(rows)}")
209
+
210
+ payload = {
211
+ "report_type": "ctt_generated_measured_rollout",
212
+ "candidates_evaluated": True,
213
+ "schema_version": 1,
214
+ "checkpoint": str(args.checkpoint),
215
+ "source_index": str(args.source_index),
216
+ "target_index": str(args.target_index),
217
+ "source_content_hash": source_index.get("content_hash"),
218
+ "source_split_hash": source_index.get("split_hash"),
219
+ "target_content_hash": target_index.get("content_hash"),
220
+ "target_split_hash": target_index.get("split_hash"),
221
+ "k": args.k,
222
+ "neighbors": args.neighbors,
223
+ "pool_size": max(pool_size, args.k),
224
+ "exclude_self_source": bool(args.exclude_self_source),
225
+ "decoder": {
226
+ "name": "linear_keyframe_decode",
227
+ "source_code": "spline_tangent_code stores start/mid/end residual keyframes",
228
+ "lossless": False,
229
+ "delta_scale": args.delta_scale,
230
+ },
231
+ "rows": rows,
232
+ }
233
+ measured_path = out_dir / "measured_candidates.json"
234
+ measured_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
235
+ (out_dir / "report.md").write_text(_report(payload) + "\n")
236
+
237
+ metrics_dir = out_dir / "measured_metrics"
238
+ if not args.skip_metrics:
239
+ eval_metrics_main(
240
+ [
241
+ "--input",
242
+ str(measured_path),
243
+ "--out-dir",
244
+ str(metrics_dir),
245
+ "--mode",
246
+ "measured",
247
+ "--k",
248
+ str(args.k),
249
+ "--bootstrap-samples",
250
+ str(args.bootstrap_samples),
251
+ ]
252
+ )
253
+ _write_required_artifacts(
254
+ out_dir,
255
+ payload,
256
+ source_index=source_index,
257
+ target_index=target_index,
258
+ metrics_dir=metrics_dir if metrics_dir.exists() else None,
259
+ )
260
+
261
+ print(
262
+ json.dumps(
263
+ {
264
+ "out_dir": str(out_dir),
265
+ "num_rows": len(rows),
266
+ "measured_candidates": str(measured_path),
267
+ },
268
+ indent=2,
269
+ )
270
+ )
271
+ return 0
272
+
273
+
274
+ def load_chart_items(
275
+ index_path: Path,
276
+ *,
277
+ max_charts: int | None,
278
+ require_positive: bool,
279
+ include_hidden: bool,
280
+ include_metadata: bool,
281
+ chart_feature_mode: str = "base",
282
+ ) -> tuple[list[ChartItem], dict[str, Any]]:
283
+ index = json.loads(index_path.read_text())
284
+ grouped: dict[str, dict[str, Any]] = {}
285
+ for shard in index.get("shards", []):
286
+ shard_path = index_path.parent / shard["path"]
287
+ with np.load(shard_path, allow_pickle=False) as data:
288
+ chart_ids = data["chart_id"]
289
+ task_ids = data["task_id"]
290
+ seeds = data["seed"]
291
+ state_hashes = data["state_hash"]
292
+ action_shapes = data["action_shape"]
293
+ base_actions = data["base_action"]
294
+ labels = data["label"]
295
+ spline_tangents = data["spline_tangent_code"]
296
+ is_base_branch = data["is_base_branch"]
297
+ utilities = data["utility"] if include_hidden else None
298
+ candidate_types = data["candidate_type"] if include_hidden else None
299
+ metadata_values = data["metadata_json"] if include_metadata else None
300
+ for row in range(chart_ids.shape[0]):
301
+ chart_id = str(chart_ids[row])
302
+ task_id = str(task_ids[row])
303
+ metadata = (
304
+ _json_loads(str(metadata_values[row]))
305
+ if metadata_values is not None
306
+ else {}
307
+ ) | {"_chart_root": str(index_path.parent)}
308
+ shape = tuple(int(value) for value in action_shapes[row])
309
+ flat_count = int(math.prod(shape))
310
+ base_action = np.asarray(
311
+ base_actions[row][:flat_count], dtype=np.float32
312
+ ).reshape(shape)
313
+ item = grouped.setdefault(
314
+ chart_id,
315
+ {
316
+ "chart_id": chart_id,
317
+ "task_id": task_id,
318
+ "seed": str(seeds[row]),
319
+ "state_hash": str(state_hashes[row]),
320
+ "instruction": str(metadata.get("instruction", "")),
321
+ "metadata": metadata | {"task_id": task_id},
322
+ "source_dataset": _source_dataset_from_metadata(
323
+ metadata,
324
+ index=index,
325
+ task_id=task_id,
326
+ ),
327
+ "base_action": base_action,
328
+ "positive_tangents": [],
329
+ "negative_tangents": [],
330
+ "hidden_utilities": [],
331
+ "hidden_candidate_types": [],
332
+ "stored_base_utility": None,
333
+ },
334
+ )
335
+ label = str(labels[row])
336
+ tangent = np.asarray(spline_tangents[row], dtype=np.float32)
337
+ if label == "positive":
338
+ item["positive_tangents"].append(tangent)
339
+ elif label == "negative":
340
+ item["negative_tangents"].append(tangent)
341
+ utility = float(utilities[row]) if utilities is not None else math.nan
342
+ if include_hidden and math.isfinite(utility):
343
+ item["hidden_utilities"].append(utility)
344
+ item["hidden_candidate_types"].append(str(candidate_types[row]))
345
+ if bool(is_base_branch[row]):
346
+ item["stored_base_utility"] = utility if math.isfinite(utility) else None
347
+ item["base_action"] = base_action
348
+ item["metadata"] = metadata | {"task_id": task_id}
349
+
350
+ charts: list[ChartItem] = []
351
+ for chart_id, item in sorted(grouped.items()):
352
+ positives = _matrix_or_empty(item["positive_tangents"], width=21)
353
+ negatives = _matrix_or_empty(item["negative_tangents"], width=21)
354
+ if require_positive and not len(positives):
355
+ continue
356
+ base_action = np.asarray(item["base_action"], dtype=np.float32)
357
+ charts.append(
358
+ ChartItem(
359
+ chart_id=chart_id,
360
+ task_id=str(item["task_id"]),
361
+ seed=str(item["seed"]),
362
+ state_hash=str(item["state_hash"]),
363
+ instruction=str(item["instruction"]),
364
+ source_dataset=Path(item["source_dataset"]).resolve(),
365
+ base_action=base_action,
366
+ feature=build_chart_feature(
367
+ base_action,
368
+ item.get("metadata", {}),
369
+ mode=chart_feature_mode,
370
+ ),
371
+ positive_tangents=positives,
372
+ negative_tangents=negatives,
373
+ hidden_utilities=[float(value) for value in item["hidden_utilities"]],
374
+ hidden_candidate_types=[str(value) for value in item["hidden_candidate_types"]],
375
+ stored_base_utility=item["stored_base_utility"],
376
+ )
377
+ )
378
+ if max_charts is not None and len(charts) >= int(max_charts):
379
+ break
380
+ return charts, index
381
+
382
+
383
+ def generate_proposals(
384
+ target: ChartItem,
385
+ *,
386
+ source_charts: list[ChartItem],
387
+ source_by_task: dict[str, list[ChartItem]],
388
+ encoder: ChartEncoder,
389
+ ctt: CausalTangentTransport,
390
+ utility_energy: UtilityEnergy,
391
+ normalizer: TangentNormalizer,
392
+ device: str,
393
+ neighbors: int,
394
+ pool_size: int,
395
+ k: int,
396
+ delta_scale: float,
397
+ exclude_self_source: bool = False,
398
+ ) -> list[Proposal]:
399
+ task_pool = source_by_task.get(target.task_id) or source_charts
400
+ pool = _source_pool_for_target(
401
+ target,
402
+ task_pool=task_pool,
403
+ source_charts=source_charts,
404
+ exclude_self_source=exclude_self_source,
405
+ )
406
+ target_feature = torch.as_tensor(target.feature, dtype=torch.float32, device=device)
407
+ ranked_sources = sorted(
408
+ pool,
409
+ key=lambda source: float(
410
+ torch.linalg.vector_norm(
411
+ torch.as_tensor(source.feature, dtype=torch.float32, device=device)
412
+ - target_feature
413
+ )
414
+ .detach()
415
+ .cpu()
416
+ ),
417
+ )[:neighbors]
418
+ target_z = encoder(target_feature.unsqueeze(0))
419
+ proposals: list[Proposal] = []
420
+ with torch.no_grad():
421
+ for source_rank, source in enumerate(ranked_sources):
422
+ if len(proposals) >= pool_size:
423
+ break
424
+ source_feature = torch.as_tensor(source.feature, dtype=torch.float32, device=device)
425
+ source_z = encoder(source_feature.unsqueeze(0))
426
+ for xi_source_raw in source.positive_tangents:
427
+ if len(proposals) >= pool_size:
428
+ break
429
+ xi_source = torch.as_tensor(
430
+ xi_source_raw, dtype=torch.float32, device=device
431
+ ).unsqueeze(0)
432
+ xi_source_norm = normalizer.transform(xi_source)
433
+ xi_hat_norm = ctt(source_z, target_z, xi_source_norm)
434
+ score = float(utility_energy(target_z, xi_hat_norm).squeeze(0).detach().cpu())
435
+ xi_hat = normalizer.inverse_transform(xi_hat_norm).squeeze(0).detach().cpu().numpy()
436
+ action_delta = decode_linear_keyframe_tangent(
437
+ xi_hat,
438
+ horizon=target.base_action.shape[0],
439
+ action_dim=target.base_action.shape[1],
440
+ )
441
+ action = target.base_action + float(delta_scale) * action_delta
442
+ proposals.append(
443
+ Proposal(
444
+ tangent=xi_hat.astype(np.float32, copy=False),
445
+ action=action.astype(np.float32, copy=False),
446
+ score=score,
447
+ source_chart_id=source.chart_id,
448
+ source_task_id=source.task_id,
449
+ source_rank=source_rank,
450
+ )
451
+ )
452
+ proposals.sort(key=lambda proposal: proposal.score, reverse=True)
453
+ return proposals[:k]
454
+
455
+
456
+ def _source_pool_for_target(
457
+ target: ChartItem,
458
+ *,
459
+ task_pool: list[ChartItem],
460
+ source_charts: list[ChartItem],
461
+ exclude_self_source: bool,
462
+ ) -> list[ChartItem]:
463
+ if not exclude_self_source:
464
+ return task_pool
465
+
466
+ def is_not_self(source: ChartItem) -> bool:
467
+ return source.chart_id != target.chart_id and source.state_hash != target.state_hash
468
+
469
+ filtered = [source for source in task_pool if is_not_self(source)]
470
+ if filtered:
471
+ return filtered
472
+ fallback = [source for source in source_charts if is_not_self(source)]
473
+ return fallback or task_pool
474
+
475
+
476
+ def decode_linear_keyframe_tangent(
477
+ tangent_code: np.ndarray,
478
+ *,
479
+ horizon: int,
480
+ action_dim: int,
481
+ ) -> np.ndarray:
482
+ """Decode the public 21D CIL keyframe code into a full residual action chunk.
483
+
484
+ The chart exporter stores start/mid/end residual rows for the common H=16, D=7
485
+ chunk. This decoder linearly interpolates those rows; it is intentionally marked
486
+ non-lossless in rollout metadata.
487
+ """
488
+
489
+ if horizon <= 0 or action_dim <= 0:
490
+ raise ValueError("horizon and action_dim must be positive")
491
+ code = np.asarray(tangent_code, dtype=np.float32).reshape(-1)
492
+ key_dim = min(action_dim, max(1, min(7, code.shape[0] // 3)))
493
+ keyframes = np.zeros((3, action_dim), dtype=np.float32)
494
+ usable = min(3 * key_dim, code.shape[0])
495
+ keyframes[:, :key_dim] = code[:usable].reshape(3, key_dim)
496
+ if horizon == 1:
497
+ return keyframes[:1]
498
+ mid = horizon // 2
499
+ positions = np.asarray([0, mid, horizon - 1], dtype=np.float32)
500
+ timeline = np.arange(horizon, dtype=np.float32)
501
+ decoded = np.zeros((horizon, action_dim), dtype=np.float32)
502
+ for dim in range(action_dim):
503
+ decoded[:, dim] = np.interp(timeline, positions, keyframes[:, dim])
504
+ return decoded
505
+
506
+
507
+ def rollout_generated_cases(
508
+ generated_cases: list[tuple[ChartItem, list[Proposal]]],
509
+ *,
510
+ gym: Any,
511
+ torch: Any,
512
+ device: str,
513
+ group_batch_size: int,
514
+ sim_backend: str | None,
515
+ render_backend: str | None,
516
+ restore_tolerance: float,
517
+ log_path: Path | None = None,
518
+ ) -> list[dict[str, Any]]:
519
+ archives: dict[Path, dict[str, Any]] = {}
520
+ rows: list[dict[str, Any]] = []
521
+ by_task: dict[str, list[tuple[ChartItem, list[Proposal]]]] = {}
522
+ for item in generated_cases:
523
+ by_task.setdefault(item[0].task_id, []).append(item)
524
+ for task_id, cases in sorted(by_task.items()):
525
+ for start in range(0, len(cases), group_batch_size):
526
+ batch = cases[start : start + group_batch_size]
527
+ source_summary = _source_summary(batch[0][0].source_dataset)
528
+ resolved_render_backend = (
529
+ render_backend
530
+ if render_backend is not None
531
+ else source_summary.get("render_backend") or "none"
532
+ )
533
+ max_candidate_count = max(1 + len(proposals) for _target, proposals in batch)
534
+ env_kwargs = {
535
+ "num_envs": len(batch) * max_candidate_count,
536
+ "obs_mode": "state",
537
+ "control_mode": source_summary.get("control_mode", "pd_ee_delta_pose"),
538
+ "render_mode": None,
539
+ "sim_backend": sim_backend or source_summary.get("sim_backend", "physx_cuda"),
540
+ "render_backend": resolved_render_backend,
541
+ "reward_mode": "normalized_dense",
542
+ }
543
+ if _uses_single_env_cpu_backend(env_kwargs["sim_backend"]) and (
544
+ len(batch) * max_candidate_count > 1
545
+ ):
546
+ rows.extend(
547
+ _rollout_cpu_sequential_batch(
548
+ task_id,
549
+ batch,
550
+ gym=gym,
551
+ torch=torch,
552
+ device=device,
553
+ env_kwargs=dict(env_kwargs) | {"num_envs": 1},
554
+ archives=archives,
555
+ restore_tolerance=restore_tolerance,
556
+ log_path=log_path,
557
+ )
558
+ )
559
+ continue
560
+ _append_log(
561
+ log_path,
562
+ f"env init task={task_id} start={start} batch={len(batch)} "
563
+ f"candidates={max_candidate_count} sim={env_kwargs['sim_backend']} "
564
+ f"render={env_kwargs['render_backend']}",
565
+ )
566
+ env = gym.make(task_id, **env_kwargs)
567
+ base_env = env.unwrapped
568
+ try:
569
+ env_device = getattr(base_env, "device", torch.device(device))
570
+ env_dim = _env_action_dim(env)
571
+ states: list[dict[str, Any]] = []
572
+ action_groups: list[np.ndarray] = []
573
+ valid_counts: list[int] = []
574
+ for target, proposals in batch:
575
+ archive = archives.setdefault(
576
+ target.source_dataset, _load_state_archive(target.source_dataset)
577
+ )
578
+ states.append(archive["initial"][target.chart_id])
579
+ group_actions = [target.base_action] + [proposal.action for proposal in proposals]
580
+ valid_counts.append(len(group_actions))
581
+ while len(group_actions) < max_candidate_count:
582
+ group_actions.append(target.base_action)
583
+ action_groups.append(np.stack(group_actions, axis=0))
584
+ candidate_values = np.stack(action_groups, axis=0).astype(np.float32)
585
+ candidate_values = _adapt_action_dim_4d(candidate_values, env_dim)
586
+ candidate_values = _clip_to_action_space_4d(candidate_values, env)
587
+ _append_log(
588
+ log_path,
589
+ f"execute task={task_id} start={start} shape={candidate_values.shape}",
590
+ )
591
+ _after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
592
+ base_env,
593
+ states,
594
+ candidate_values,
595
+ torch=torch,
596
+ device=env_device,
597
+ restore_tolerance=restore_tolerance,
598
+ )
599
+ for index, (target, proposals) in enumerate(batch):
600
+ valid = valid_counts[index]
601
+ progress = [
602
+ float(max(0.0, min(1.0, rewards[index, candidate_index])))
603
+ for candidate_index in range(valid)
604
+ ]
605
+ success = [
606
+ bool(successes[index, candidate_index])
607
+ for candidate_index in range(valid)
608
+ ]
609
+ utilities = [
610
+ progress_value + (1.0 if success_value else 0.0)
611
+ for progress_value, success_value in zip(progress, success, strict=True)
612
+ ]
613
+ rows.append(
614
+ _measured_row_from_rollout(
615
+ target,
616
+ proposals,
617
+ progress=progress,
618
+ success=success,
619
+ utilities=utilities,
620
+ restore_error=float(restore_error),
621
+ )
622
+ )
623
+ finally:
624
+ env.close()
625
+ _append_log(log_path, f"batch done task={task_id} start={start}")
626
+ return rows
627
+
628
+
629
+ def _rollout_cpu_sequential_batch(
630
+ task_id: str,
631
+ batch: list[tuple[ChartItem, list[Proposal]]],
632
+ *,
633
+ gym: Any,
634
+ torch: Any,
635
+ device: str,
636
+ env_kwargs: dict[str, Any],
637
+ archives: dict[Path, dict[str, Any]],
638
+ restore_tolerance: float,
639
+ log_path: Path | None,
640
+ ) -> list[dict[str, Any]]:
641
+ rows: list[dict[str, Any]] = []
642
+ _append_log(
643
+ log_path,
644
+ f"env init sequential task={task_id} batch={len(batch)} sim={env_kwargs['sim_backend']} "
645
+ f"render={env_kwargs['render_backend']}",
646
+ )
647
+ env = gym.make(task_id, **env_kwargs)
648
+ base_env = env.unwrapped
649
+ try:
650
+ env_device = getattr(base_env, "device", torch.device(device))
651
+ env_dim = _env_action_dim(env)
652
+ for target, proposals in batch:
653
+ archive = archives.setdefault(
654
+ target.source_dataset, _load_state_archive(target.source_dataset)
655
+ )
656
+ state = archive["initial"][target.chart_id]
657
+ candidate_actions = [target.base_action] + [proposal.action for proposal in proposals]
658
+ progress: list[float] = []
659
+ success: list[bool] = []
660
+ restore_errors: list[float] = []
661
+ for candidate_index, action in enumerate(candidate_actions):
662
+ candidate_values = np.asarray(action, dtype=np.float32).reshape(
663
+ 1, 1, *action.shape
664
+ )
665
+ candidate_values = _adapt_action_dim_4d(candidate_values, env_dim)
666
+ candidate_values = _clip_to_action_space_4d(candidate_values, env)
667
+ _append_log(
668
+ log_path,
669
+ f"execute sequential task={task_id} chart={target.chart_id} "
670
+ f"candidate={candidate_index} shape={candidate_values.shape}",
671
+ )
672
+ _after_state, rewards, successes, restore_error = execute_grouped_action_lattice_batch(
673
+ base_env,
674
+ [state],
675
+ candidate_values,
676
+ torch=torch,
677
+ device=env_device,
678
+ restore_tolerance=restore_tolerance,
679
+ )
680
+ progress.append(float(max(0.0, min(1.0, rewards[0, 0]))))
681
+ success.append(bool(successes[0, 0]))
682
+ restore_errors.append(float(restore_error))
683
+ utilities = [
684
+ progress_value + (1.0 if success_value else 0.0)
685
+ for progress_value, success_value in zip(progress, success, strict=True)
686
+ ]
687
+ rows.append(
688
+ _measured_row_from_rollout(
689
+ target,
690
+ proposals,
691
+ progress=progress,
692
+ success=success,
693
+ utilities=utilities,
694
+ restore_error=max(restore_errors, default=0.0),
695
+ )
696
+ )
697
+ finally:
698
+ env.close()
699
+ _append_log(log_path, f"batch done sequential task={task_id}")
700
+ return rows
701
+
702
+
703
+ def _measured_row_from_rollout(
704
+ target: ChartItem,
705
+ proposals: list[Proposal],
706
+ *,
707
+ progress: list[float],
708
+ success: list[bool],
709
+ utilities: list[float],
710
+ restore_error: float,
711
+ ) -> dict[str, Any]:
712
+ base_utility = float(utilities[0])
713
+ generated_utilities = [float(value) for value in utilities[1:]]
714
+ predicted_scores = [float(proposal.score) for proposal in proposals]
715
+ return {
716
+ "chart_id": target.chart_id,
717
+ "group_id": target.chart_id,
718
+ "task_id": target.task_id,
719
+ "seed": target.seed,
720
+ "state_hash": target.state_hash,
721
+ "instruction": target.instruction,
722
+ "candidates_evaluated": True,
723
+ "selected_index": 0,
724
+ "base_utility": base_utility,
725
+ "stored_base_utility": target.stored_base_utility,
726
+ "generated_utilities": generated_utilities,
727
+ "hidden_chart_utilities": target.hidden_utilities,
728
+ "hidden_candidate_types": target.hidden_candidate_types,
729
+ "outcome_vector_schema": [
730
+ "success",
731
+ "progress",
732
+ "contact_quality",
733
+ "safety_violation",
734
+ "stage_progress",
735
+ "smoothness",
736
+ "recovery",
737
+ "utility_scalar",
738
+ ],
739
+ "base_outcome": _outcome_payload(
740
+ success=success[0],
741
+ progress=progress[0],
742
+ utility=utilities[0],
743
+ ),
744
+ "candidate_outcomes": [
745
+ _outcome_payload(
746
+ success=success_value,
747
+ progress=progress_value,
748
+ utility=utility,
749
+ )
750
+ for success_value, progress_value, utility in zip(
751
+ success[1:],
752
+ progress[1:],
753
+ utilities[1:],
754
+ strict=True,
755
+ )
756
+ ],
757
+ "predicted_scores": predicted_scores,
758
+ "generated_tangents": [
759
+ proposal.tangent.astype(float).tolist() for proposal in proposals
760
+ ],
761
+ "positive_tangents": target.positive_tangents.astype(float).tolist(),
762
+ "negative_tangents": target.negative_tangents.astype(float).tolist(),
763
+ "candidate_types": [
764
+ f"ctt_transport_rank{proposal.source_rank}" for proposal in proposals
765
+ ],
766
+ "candidate_source_chart_ids": [proposal.source_chart_id for proposal in proposals],
767
+ "candidate_source_task_ids": [proposal.source_task_id for proposal in proposals],
768
+ "candidate_progress": progress[1:],
769
+ "candidate_success": success[1:],
770
+ "base_progress": progress[0],
771
+ "base_success": success[0],
772
+ "restore_error": float(restore_error),
773
+ "num_generated": len(proposals),
774
+ "num_executed_including_base": len(utilities),
775
+ }
776
+
777
+
778
+ def _outcome_payload(*, success: bool, progress: float, utility: float) -> dict[str, Any]:
779
+ return {
780
+ "success": bool(success),
781
+ "progress": float(progress),
782
+ "contact_quality": None,
783
+ "safety_violation": None,
784
+ "stage_progress": None,
785
+ "smoothness": None,
786
+ "recovery": None,
787
+ "utility_scalar": float(utility),
788
+ }
789
+
790
+
791
+ def _uses_single_env_cpu_backend(sim_backend: Any) -> bool:
792
+ value = str(sim_backend or "").lower()
793
+ return value in {"cpu", "physx_cpu"} or value.endswith("_cpu")
794
+
795
+
796
+ def _validate_indexes(
797
+ source_path: Path,
798
+ source_index: dict[str, Any],
799
+ target_path: Path,
800
+ target_index: dict[str, Any],
801
+ ) -> None:
802
+ if source_index.get("split") != "train" or not source_index.get("retrieval_index_allowed"):
803
+ raise SystemExit(
804
+ f"{source_path} is not a train-only retrieval index; CTT rollout sources "
805
+ "must come from train split only"
806
+ )
807
+ if not source_index.get("include_outcomes"):
808
+ raise SystemExit(f"{source_path} must include train outcomes for source positives")
809
+ if not target_index.get("include_outcomes"):
810
+ raise SystemExit(f"{target_path} must include evaluator-only outcomes")
811
+ if target_index.get("split") != "train" and target_index.get("retrieval_index_allowed"):
812
+ raise SystemExit(f"{target_path} is non-train but marked retrieval_index_allowed")
813
+
814
+
815
+ def _adapt_action_dim_4d(actions: np.ndarray, action_dim: int) -> np.ndarray:
816
+ if actions.ndim != 4:
817
+ raise ValueError("actions must have shape [B,K,H,D]")
818
+ if actions.shape[-1] == action_dim:
819
+ return actions.astype(np.float32, copy=False)
820
+ if actions.shape[-1] > action_dim:
821
+ return actions[..., :action_dim].astype(np.float32, copy=False)
822
+ pad = np.zeros((*actions.shape[:-1], action_dim - actions.shape[-1]), dtype=np.float32)
823
+ return np.concatenate([actions.astype(np.float32, copy=False), pad], axis=-1)
824
+
825
+
826
+ def _clip_to_action_space_4d(actions: np.ndarray, env: Any) -> np.ndarray:
827
+ space = getattr(env, "single_action_space", None) or getattr(env, "action_space", None)
828
+ low = getattr(space, "low", None)
829
+ high = getattr(space, "high", None)
830
+ if low is None or high is None:
831
+ return actions
832
+ low_arr = np.asarray(low, dtype=np.float32).reshape(-1)[-actions.shape[-1] :]
833
+ high_arr = np.asarray(high, dtype=np.float32).reshape(-1)[-actions.shape[-1] :]
834
+ if low_arr.shape[0] != actions.shape[-1] or high_arr.shape[0] != actions.shape[-1]:
835
+ return actions
836
+ return np.clip(actions, low_arr.reshape(1, 1, 1, -1), high_arr.reshape(1, 1, 1, -1))
837
+
838
+
839
+ def _env_action_dim(env: Any) -> int:
840
+ for space in (
841
+ getattr(env, "single_action_space", None),
842
+ getattr(env.unwrapped, "single_action_space", None),
843
+ getattr(env, "action_space", None),
844
+ ):
845
+ shape = getattr(space, "shape", None)
846
+ if shape:
847
+ return int(shape[-1])
848
+ raise ValueError("Could not infer ManiSkill action dimension from environment")
849
+
850
+
851
+ def _load_state_archive(source_dataset: Path) -> dict[str, Any]:
852
+ archive_path = source_dataset / "state_archive.pkl"
853
+ if not archive_path.exists():
854
+ summary_path = source_dataset / "generation_summary.json"
855
+ if summary_path.exists():
856
+ raw_path = read_json(summary_path).get("state_archive")
857
+ if raw_path:
858
+ archive_path = Path(str(raw_path))
859
+ if not archive_path.exists():
860
+ raise FileNotFoundError(f"ManiSkill state archive not found for {source_dataset}")
861
+ with archive_path.open("rb") as handle:
862
+ archive = pickle.load(handle)
863
+ if not isinstance(archive, dict) or "initial" not in archive:
864
+ raise ValueError(f"Invalid ManiSkill state archive: {archive_path}")
865
+ return archive
866
+
867
+
868
+ def _source_summary(source_dataset: Path) -> dict[str, Any]:
869
+ path = source_dataset / "generation_summary.json"
870
+ return read_json(path) if path.exists() else {}
871
+
872
+
873
+ def _source_dataset_from_metadata(
874
+ metadata: dict[str, Any],
875
+ *,
876
+ index: dict[str, Any],
877
+ task_id: str,
878
+ ) -> Path:
879
+ raw = metadata.get("source_dataset")
880
+ if raw:
881
+ return Path(str(raw))
882
+ dataset = index.get("dataset")
883
+ if dataset:
884
+ candidate = Path(str(dataset)) / task_id
885
+ if candidate.exists():
886
+ return candidate
887
+ raise ValueError(f"Could not resolve source_dataset for chart task {task_id}")
888
+
889
+
890
+ def _matrix_or_empty(items: list[np.ndarray], *, width: int) -> np.ndarray:
891
+ if not items:
892
+ return np.zeros((0, width), dtype=np.float32)
893
+ return np.asarray(items, dtype=np.float32)
894
+
895
+
896
+ def _json_loads(raw: str) -> dict[str, Any]:
897
+ try:
898
+ payload = json.loads(raw)
899
+ except json.JSONDecodeError:
900
+ return {}
901
+ return payload if isinstance(payload, dict) else {}
902
+
903
+
904
+ def _resolve_device(device: str) -> str:
905
+ if device != "auto":
906
+ return device
907
+ return "cuda" if torch.cuda.is_available() else "cpu"
908
+
909
+
910
+ def _write_run_provenance(out_dir: Path, args: argparse.Namespace) -> None:
911
+ config = {key: str(value) for key, value in sorted(vars(args).items())}
912
+ (out_dir / "config.yaml").write_text(
913
+ "\n".join(f"{key}: {value}" for key, value in config.items()) + "\n"
914
+ )
915
+ (out_dir / "command.txt").write_text(
916
+ "python scripts/eval_ctt_generated_rollout.py " + " ".join(sys.argv[1:]) + "\n"
917
+ )
918
+ (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
919
+ for name, path in {
920
+ "source_index.json": args.source_index,
921
+ "target_index.json": args.target_index,
922
+ }.items():
923
+ try:
924
+ (out_dir / name).write_text(Path(path).read_text())
925
+ except OSError:
926
+ pass
927
+
928
+
929
+ def _write_required_artifacts(
930
+ out_dir: Path,
931
+ payload: dict[str, Any],
932
+ *,
933
+ source_index: dict[str, Any],
934
+ target_index: dict[str, Any],
935
+ metrics_dir: Path | None,
936
+ ) -> None:
937
+ (out_dir / "data_hash.txt").write_text(str(target_index.get("content_hash", "")) + "\n")
938
+ (out_dir / "split_hash.txt").write_text(str(target_index.get("split_hash", "")) + "\n")
939
+ (out_dir / "source_data_hash.txt").write_text(str(source_index.get("content_hash", "")) + "\n")
940
+ (out_dir / "source_split_hash.txt").write_text(str(source_index.get("split_hash", "")) + "\n")
941
+ (out_dir / "train.log").write_text(
942
+ "CTT checkpoint trained separately; rollout used checkpoint="
943
+ f"{payload.get('checkpoint', '')}\n"
944
+ )
945
+ run_log = out_dir / "run.log"
946
+ (out_dir / "eval.log").write_text(run_log.read_text() if run_log.exists() else "eval log unavailable\n")
947
+ summary = {
948
+ "report_type": payload.get("report_type"),
949
+ "schema_version": payload.get("schema_version"),
950
+ "k": payload.get("k"),
951
+ "checkpoint": payload.get("checkpoint"),
952
+ "source_content_hash": payload.get("source_content_hash"),
953
+ "source_split_hash": payload.get("source_split_hash"),
954
+ "target_content_hash": payload.get("target_content_hash"),
955
+ "target_split_hash": payload.get("target_split_hash"),
956
+ "num_rows": len(payload.get("rows", [])),
957
+ "success_summary": _success_summary(payload.get("rows", []), k=int(payload.get("k", 0))),
958
+ }
959
+ if metrics_dir is not None:
960
+ metric_path = metrics_dir / "metrics.json"
961
+ if metric_path.exists():
962
+ metrics = json.loads(metric_path.read_text())
963
+ summary["measured_metric_summary"] = metrics.get("summary", {})
964
+ for filename in ("metrics_by_task.json", "metrics_by_seed.json", "table.tex"):
965
+ src = metrics_dir / filename
966
+ if src.exists():
967
+ shutil.copyfile(src, out_dir / filename)
968
+ (out_dir / "metrics.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n")
969
+ for filename in ("metrics_by_task.json", "metrics_by_seed.json", "table.tex"):
970
+ path = out_dir / filename
971
+ if not path.exists():
972
+ path.write_text("{}\n" if filename.endswith(".json") else "% metrics not computed\n")
973
+
974
+
975
+ def _success_summary(rows: list[dict[str, Any]], *, k: int) -> dict[str, Any]:
976
+ if k <= 0:
977
+ k = 10**9
978
+ base_success = []
979
+ selected_success = []
980
+ oracle_success = []
981
+ base_utility = []
982
+ selected_utility = []
983
+ oracle_utility = []
984
+ hidden_oracle_utility = []
985
+ hidden_oracle_success = []
986
+ success_support_gap = []
987
+ success_selector_gap = []
988
+ selected_success_gain = []
989
+ proposal_oracle_success_gain = []
990
+ restore_errors = []
991
+ for row in rows:
992
+ generated_utilities = [float(value) for value in row.get("generated_utilities", [])[:k]]
993
+ generated_success = [bool(value) for value in row.get("candidate_success", [])[:k]]
994
+ selected_index = int(row.get("selected_index", 0))
995
+ selected_success_value: float | None = None
996
+ proposal_oracle_success_value: float | None = None
997
+ base_success_value: float | None = None
998
+ if "base_success" in row:
999
+ base_success_value = float(bool(row["base_success"]))
1000
+ base_success.append(base_success_value)
1001
+ if selected_index < len(generated_success):
1002
+ selected_success_value = float(generated_success[selected_index])
1003
+ selected_success.append(selected_success_value)
1004
+ if generated_success:
1005
+ proposal_oracle_success_value = float(any(generated_success))
1006
+ oracle_success.append(proposal_oracle_success_value)
1007
+ if "base_utility" in row:
1008
+ base_utility.append(float(row["base_utility"]))
1009
+ if selected_index < len(generated_utilities):
1010
+ selected_utility.append(float(generated_utilities[selected_index]))
1011
+ if generated_utilities:
1012
+ oracle_utility.append(max(generated_utilities))
1013
+ hidden = [float(value) for value in row.get("hidden_chart_utilities", [])]
1014
+ if hidden:
1015
+ hidden_oracle_utility.append(max(hidden))
1016
+ hidden_success_value = float(any(value >= 1.0 for value in hidden))
1017
+ hidden_oracle_success.append(hidden_success_value)
1018
+ if proposal_oracle_success_value is not None:
1019
+ success_support_gap.append(
1020
+ max(0.0, hidden_success_value - proposal_oracle_success_value)
1021
+ )
1022
+ if (
1023
+ proposal_oracle_success_value is not None
1024
+ and selected_success_value is not None
1025
+ ):
1026
+ success_selector_gap.append(
1027
+ max(0.0, proposal_oracle_success_value - selected_success_value)
1028
+ )
1029
+ if base_success_value is not None and selected_success_value is not None:
1030
+ selected_success_gain.append(selected_success_value - base_success_value)
1031
+ if base_success_value is not None and proposal_oracle_success_value is not None:
1032
+ proposal_oracle_success_gain.append(
1033
+ proposal_oracle_success_value - base_success_value
1034
+ )
1035
+ if "restore_error" in row:
1036
+ restore_errors.append(float(row["restore_error"]))
1037
+ return {
1038
+ "base_success_rate": _mean(base_success),
1039
+ "selected_success_rate": _mean(selected_success),
1040
+ "proposal_oracle_success_rate": _mean(oracle_success),
1041
+ "hidden_chart_oracle_success_rate": _mean(hidden_oracle_success),
1042
+ "selected_success_gain_over_base": _mean(selected_success_gain),
1043
+ "proposal_oracle_success_gain_over_base": _mean(proposal_oracle_success_gain),
1044
+ "success_support_gap": _mean(success_support_gap),
1045
+ "success_selector_gap": _mean(success_selector_gap),
1046
+ "base_utility_mean": _mean(base_utility),
1047
+ "selected_utility_mean": _mean(selected_utility),
1048
+ "proposal_oracle_utility_mean": _mean(oracle_utility),
1049
+ "hidden_chart_oracle_utility_mean": _mean(hidden_oracle_utility),
1050
+ "max_restore_error": max(restore_errors) if restore_errors else None,
1051
+ }
1052
+
1053
+
1054
+ def _append_log(path: Path | None, message: str) -> None:
1055
+ if path is None:
1056
+ return
1057
+ with path.open("a") as handle:
1058
+ handle.write(message + "\n")
1059
+
1060
+
1061
+ def _run(command: list[str]) -> str:
1062
+ try:
1063
+ return subprocess.check_output(command, cwd=PROJECT_ROOT, text=True).strip()
1064
+ except (subprocess.CalledProcessError, FileNotFoundError):
1065
+ return ""
1066
+
1067
+
1068
+ def _report(payload: dict[str, Any]) -> str:
1069
+ rows = payload["rows"]
1070
+ outcome = [
1071
+ 1.0
1072
+ if any(
1073
+ float(value) > float(row["base_utility"])
1074
+ for value in row.get("generated_utilities", [])[: int(payload["k"])]
1075
+ )
1076
+ else 0.0
1077
+ for row in rows
1078
+ ]
1079
+ selector_regrets = []
1080
+ support_gaps = []
1081
+ for row in rows:
1082
+ generated = row.get("generated_utilities", [])[: int(payload["k"])]
1083
+ if generated:
1084
+ selector_regrets.append(max(generated) - generated[int(row.get("selected_index", 0))])
1085
+ hidden = row.get("hidden_chart_utilities", [])
1086
+ if hidden:
1087
+ support_gaps.append(max(hidden) - max(generated))
1088
+ lines = [
1089
+ "# CTT Generated Measured Rollout",
1090
+ "",
1091
+ f"Rows: `{len(rows)}`",
1092
+ f"K: `{payload['k']}`",
1093
+ f"Checkpoint: `{payload['checkpoint']}`",
1094
+ "",
1095
+ "| Metric | Mean |",
1096
+ "| --- | ---: |",
1097
+ f"| OutcomePTR@K | {_mean(outcome):.4f} |",
1098
+ f"| SelectorRegret@K | {_mean(selector_regrets):.4f} |",
1099
+ f"| SupportGap@K | {_mean(support_gaps):.4f} |",
1100
+ "",
1101
+ "These are measured simulator rollouts of decoded CTT candidates, not PPTC proxies.",
1102
+ ]
1103
+ return "\n".join(lines)
1104
+
1105
+
1106
+ def _mean(values: list[float]) -> float:
1107
+ clean = [float(value) for value in values if math.isfinite(float(value))]
1108
+ return sum(clean) / len(clean) if clean else math.nan
1109
+
1110
+
1111
+ if __name__ == "__main__":
1112
+ raise SystemExit(main())