ctt artifacts 2026-07-02 workspace/scripts/train_ctt.py

workspace/scripts/train_ctt.py

@@ -0,0 +1,303 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+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.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer  # noqa: E402
+from cil.models.ctt import chamfer_to_target_set, negative_boundary_loss  # noqa: E402
+
+
+@dataclass
+class Chart:
+    chart_id: str
+    task_id: str
+    feature: torch.Tensor
+    positives: torch.Tensor
+    negatives: torch.Tensor
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Train Causal Tangent Transport on chart DB.")
+    parser.add_argument("--config", type=Path, default=None)
+    parser.add_argument("--dataset", type=Path, default=Path("data/cil_charts/train/index.json"))
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/ctt_residual_smoke"))
+    parser.add_argument("--variant", choices=("residual", "gated_residual"), default="residual")
+    parser.add_argument("--epochs", type=int, default=2)
+    parser.add_argument("--max-charts", type=int, default=64)
+    parser.add_argument("--neighbors", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=1.0e-3)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--pos-alignment", type=float, default=1.0)
+    parser.add_argument("--negative-boundary", type=float, default=0.25)
+    parser.add_argument("--cycle", type=float, default=0.1)
+    parser.add_argument("--diversity", type=float, default=0.05)
+    parser.add_argument("--negative-margin", type=float, default=0.2)
+    args = parser.parse_args(argv)
+    if args.config:
+        args = _merge_config(args, _load_simple_yaml(args.config))
+
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+
+    out_dir = args.out_dir
+    out_dir.mkdir(parents=True, exist_ok=True)
+    _write_run_provenance(out_dir, args)
+    charts, index = load_charts(args.dataset, max_charts=args.max_charts)
+    if len(charts) < 2:
+        raise SystemExit("CTT training requires at least two charts with positive tangents")
+
+    all_xi = torch.cat([chart.positives for chart in charts] + [chart.negatives for chart in charts if chart.negatives.numel()], dim=0)
+    normalizer = TangentNormalizer.fit(all_xi)
+    for chart in charts:
+        chart.positives[:] = normalizer.transform(chart.positives)
+        if chart.negatives.numel():
+            chart.negatives[:] = normalizer.transform(chart.negatives)
+
+    feature_dim = int(charts[0].feature.numel())
+    config = CTTConfig(
+        chart_feature_dim=feature_dim,
+        tangent_dim=int(charts[0].positives.shape[1]),
+        variant=args.variant,
+    )
+    encoder = ChartEncoder(feature_dim, output_dim=config.chart_dim)
+    ctt = CausalTangentTransport(config)
+    optimizer = torch.optim.Adam(list(encoder.parameters()) + list(ctt.parameters()), lr=args.lr)
+    pairs = _neighbor_pairs(charts, neighbors=args.neighbors)
+    log_lines = [
+        f"charts={len(charts)} pairs={len(pairs)} variant={args.variant}",
+        f"loss_weights pos={args.pos_alignment} neg={args.negative_boundary} cycle={args.cycle} diversity={args.diversity}",
+    ]
+
+    for epoch in range(args.epochs):
+        random.shuffle(pairs)
+        losses = []
+        for source_idx, target_idx in pairs:
+            source = charts[source_idx]
+            target = charts[target_idx]
+            xi_source = source.positives[random.randrange(source.positives.shape[0])].unsqueeze(0)
+            z_source = encoder(source.feature.unsqueeze(0))
+            z_target = encoder(target.feature.unsqueeze(0))
+            xi_hat = ctt(z_source, z_target, xi_source)
+            loss_pos = chamfer_to_target_set(xi_hat, target.positives)
+            loss_neg = negative_boundary_loss(
+                xi_hat,
+                target.negatives,
+                margin=args.negative_margin,
+            )
+            xi_cycle = ctt(z_target, z_source, xi_hat)
+            loss_cycle = (xi_cycle - xi_source).pow(2).mean()
+            loss = (
+                args.pos_alignment * loss_pos
+                + args.negative_boundary * loss_neg
+                + args.cycle * loss_cycle
+            )
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            losses.append(float(loss.detach()))
+        mean_loss = sum(losses) / len(losses) if losses else math.nan
+        log_lines.append(f"epoch={epoch + 1} mean_loss={mean_loss:.6f}")
+
+    checkpoint = {
+        "config": config.to_dict(),
+        "chart_encoder": encoder.state_dict(),
+        "ctt": ctt.state_dict(),
+        "normalizer": normalizer.to_dict(),
+        "source_index": str(args.dataset),
+        "data_hash": index.get("content_hash"),
+        "split_hash": index.get("split_hash"),
+    }
+    torch.save(checkpoint, out_dir / "model.pt")
+    metrics = {
+        "report_type": "ctt_train",
+        "variant": args.variant,
+        "epochs": args.epochs,
+        "num_charts": len(charts),
+        "num_pairs": len(pairs),
+        "final_mean_loss": mean_loss,
+        "data_hash": index.get("content_hash"),
+        "split_hash": index.get("split_hash"),
+        "loss_weights": {
+            "pos_alignment": args.pos_alignment,
+            "negative_boundary": args.negative_boundary,
+            "cycle": args.cycle,
+            "diversity": args.diversity,
+            "negative_margin": args.negative_margin,
+        },
+    }
+    (out_dir / "train.log").write_text("\n".join(log_lines) + "\n")
+    (out_dir / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")
+    (out_dir / "metrics_by_task.json").write_text("{}\n")
+    (out_dir / "metrics_by_seed.json").write_text("{}\n")
+    (out_dir / "eval.log").write_text("eval not run by train_ctt.py\n")
+    (out_dir / "table.tex").write_text(_table(metrics) + "\n")
+    (out_dir / "report.md").write_text(_report(metrics) + "\n")
+    print(json.dumps({"out_dir": str(out_dir), "final_mean_loss": mean_loss}, indent=2))
+    return 0
+
+
+def load_charts(index_path: Path, *, max_charts: int | None = None) -> tuple[list[Chart], dict[str, Any]]:
+    index = json.loads(index_path.read_text())
+    if not index.get("include_outcomes", False):
+        raise SystemExit(f"{index_path} does not include outcomes; CTT training requires train split")
+    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"]
+            base_actions = data["base_action"]
+            labels = data["label"]
+            is_base = data["is_base_branch"]
+            spline_codes = data["spline_tangent_code"]
+            for row in range(chart_ids.shape[0]):
+                chart_id = str(chart_ids[row])
+                item = grouped.setdefault(
+                    chart_id,
+                    {
+                        "task_id": str(task_ids[row]),
+                        "base": None,
+                        "positives": [],
+                        "negatives": [],
+                    },
+                )
+                if bool(is_base[row]):
+                    item["base"] = base_actions[row].astype("float32")
+                label = str(labels[row])
+                if label == "positive":
+                    item["positives"].append(spline_codes[row].astype("float32"))
+                elif label == "negative":
+                    item["negatives"].append(spline_codes[row].astype("float32"))
+    charts = []
+    for chart_id, item in sorted(grouped.items()):
+        if item["base"] is None or not item["positives"]:
+            continue
+        positives = torch.as_tensor(np.asarray(item["positives"]), dtype=torch.float32)
+        negatives = torch.as_tensor(np.asarray(item["negatives"]), dtype=torch.float32)
+        charts.append(
+            Chart(
+                chart_id=chart_id,
+                task_id=str(item["task_id"]),
+                feature=torch.as_tensor(item["base"], dtype=torch.float32),
+                positives=positives,
+                negatives=negatives,
+            )
+        )
+    if max_charts is not None:
+        charts = charts[: int(max_charts)]
+    return charts, index
+
+
+def _neighbor_pairs(charts: list[Chart], *, neighbors: int) -> list[tuple[int, int]]:
+    pairs: list[tuple[int, int]] = []
+    for target_idx, target in enumerate(charts):
+        candidates = []
+        for source_idx, source in enumerate(charts):
+            if source_idx == target_idx or source.task_id != target.task_id:
+                continue
+            distance = torch.linalg.vector_norm(source.feature - target.feature).item()
+            candidates.append((distance, source_idx))
+        for _, source_idx in sorted(candidates)[: int(neighbors)]:
+            pairs.append((source_idx, target_idx))
+    return pairs
+
+
+def _write_run_provenance(out_dir: Path, args: argparse.Namespace) -> None:
+    config = vars(args) | {"config": None if args.config is None else str(args.config)}
+    (out_dir / "config.yaml").write_text("\n".join(f"{k}: {v}" for k, v in sorted(config.items())) + "\n")
+    (out_dir / "command.txt").write_text("python scripts/train_ctt.py " + " ".join(sys.argv[1:]) + "\n")
+    (out_dir / "git_hash.txt").write_text(_run(["git", "rev-parse", "HEAD"]) + "\n")
+    index = json.loads(Path(args.dataset).read_text())
+    (out_dir / "data_hash.txt").write_text(str(index.get("content_hash", "")) + "\n")
+    (out_dir / "split_hash.txt").write_text(str(index.get("split_hash", "")) + "\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 _load_simple_yaml(path: Path) -> dict[str, Any]:
+    values: dict[str, Any] = {}
+    for line in path.read_text().splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith("#") or ":" not in stripped:
+            continue
+        key, raw = stripped.split(":", 1)
+        values[key.strip().replace("-", "_")] = _coerce(raw.strip())
+    return values
+
+
+def _merge_config(args: argparse.Namespace, config: dict[str, Any]) -> argparse.Namespace:
+    for key, value in config.items():
+        if hasattr(args, key):
+            if key in {"dataset", "out_dir"}:
+                value = Path(value)
+            setattr(args, key, value)
+    return args
+
+
+def _coerce(value: str) -> Any:
+    if value in {"true", "True"}:
+        return True
+    if value in {"false", "False"}:
+        return False
+    try:
+        if "." in value:
+            return float(value)
+        return int(value)
+    except ValueError:
+        return value
+
+
+def _table(metrics: dict[str, Any]) -> str:
+    return "\n".join(
+        [
+            "% Auto-generated by scripts/train_ctt.py",
+            "\\begin{tabular}{lrrr}",
+            "\\toprule",
+            "Variant & Charts & Pairs & Final loss \\\\",
+            "\\midrule",
+            f"{metrics['variant']} & {metrics['num_charts']} & {metrics['num_pairs']} & {metrics['final_mean_loss']:.4f} \\\\",
+            "\\bottomrule",
+            "\\end{tabular}",
+        ]
+    )
+
+
+def _report(metrics: dict[str, Any]) -> str:
+    return "\n".join(
+        [
+            "# CTT Train Report",
+            "",
+            f"Variant: `{metrics['variant']}`",
+            f"Charts: `{metrics['num_charts']}`",
+            f"Neighbor pairs: `{metrics['num_pairs']}`",
+            f"Final mean loss: `{metrics['final_mean_loss']:.6f}`",
+            "",
+            "This smoke run trains CTT from measured train positives; it is not a rollout success claim.",
+        ]
+    )
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())