workspace/scripts/export_cil_charts.py

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

import argparse
import hashlib
import json
import math
import sys
from collections import Counter
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))

from dovla_cil.data.datasets import CILDataset  # noqa: E402
from dovla_cil.data.schema import CILRecord, OutcomeVector  # noqa: E402
from dovla_cil.generation.tangent_targets import (  # noqa: E402
    DEFAULT_BASE_CANDIDATE_TYPES,
    action_matrix,
    action_shape,
    choose_base_record,
    label_from_delta_utility,
    record_utility,
    spline_tangent_summary,
    subtract_actions,
)


LABEL_TO_ID = {"negative": -1, "neutral": 0, "positive": 1, "hidden": 9}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Export same-state CIL charts to NPZ shards for train-only CIL-Atlas "
            "retrieval/generator training, with an index that can be leakage-audited."
        )
    )
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("data/cil_charts/train"),
        help="Output split directory, e.g. data/cil_charts/train.",
    )
    parser.add_argument("--split", default="train")
    parser.add_argument(
        "--split-policy",
        choices=("explicit", "stable-hash"),
        default="explicit",
        help="Use explicit split for all groups or deterministic group-id hash splits.",
    )
    parser.add_argument(
        "--split-fractions",
        default="0.70,0.15,0.15",
        help="Train,val,test fractions for --split-policy stable-hash.",
    )
    parser.add_argument("--split-seed", type=int, default=0)
    parser.add_argument("--epsilon", type=float, default=0.0)
    parser.add_argument("--max-groups", type=int, default=None)
    parser.add_argument("--shard-size", type=int, default=50000)
    parser.add_argument(
        "--base-candidate-types",
        default=",".join(DEFAULT_BASE_CANDIDATE_TYPES),
    )
    parser.add_argument(
        "--include-outcomes",
        action="store_true",
        help=(
            "Include measured utilities/outcome vectors. Defaults to true only for "
            "the train split; non-train outcome exports are evaluator-only."
        ),
    )
    parser.add_argument(
        "--no-outcomes",
        action="store_true",
        help="Force-hide utilities/outcomes even for train export.",
    )
    args = parser.parse_args(argv)

    if args.epsilon < 0.0:
        parser.error("--epsilon must be non-negative")
    if args.max_groups is not None and args.max_groups <= 0:
        parser.error("--max-groups must be positive when provided")
    if args.shard_size <= 0:
        parser.error("--shard-size must be positive")
    split_fractions = _parse_split_fractions(args.split_fractions)

    try:
        import numpy as np
    except ImportError as exc:  # pragma: no cover - environment guard
        raise SystemExit("export_cil_charts.py requires numpy to write .npz shards") from exc

    include_outcomes = (args.split == "train" or args.include_outcomes) and not args.no_outcomes
    dataset = CILDataset(args.dataset)
    out_dir = args.out_dir
    out_dir.mkdir(parents=True, exist_ok=True)
    for stale in list(out_dir.glob("charts_*.npz")) + [
        out_dir / "index.json",
        out_dir / "candidate_type_counts.json",
    ]:
        if stale.exists():
            stale.unlink()
    rows: list[dict[str, Any]] = []
    shard_paths: list[dict[str, Any]] = []
    counters: Counter[str] = Counter()
    task_counts: Counter[str] = Counter()
    label_counts: Counter[str] = Counter()
    group_hashes: set[str] = set()
    state_hashes: set[str] = set()
    max_flat_dim = 0
    group_count = 0
    skipped_by_split = 0
    base_candidate_types = _parse_csv(args.base_candidate_types)

    for group in dataset.iter_groups():
        if args.max_groups is not None and group_count >= args.max_groups:
            break
        group_count += 1
        if not group:
            counters["empty_group"] += 1
            continue
        assigned_split = _assign_split(
            group[0].group_id,
            policy=args.split_policy,
            split=args.split,
            fractions=split_fractions,
            seed=args.split_seed,
        )
        if assigned_split != args.split:
            skipped_by_split += 1
            continue
        base = choose_base_record(group, base_candidate_types=base_candidate_types)
        if base is None:
            counters["missing_base"] += 1
            continue
        base_action = action_matrix(base)
        if not base_action:
            counters["empty_base_action"] += 1
            continue
        base_shape = action_shape(base_action)
        base_utility = record_utility(base) if include_outcomes else math.nan
        emitted_for_group = 0
        for record in group:
            action = action_matrix(record)
            if not action:
                counters["empty_action"] += 1
                continue
            if action_shape(action) != base_shape:
                counters["action_shape_mismatch"] += 1
                continue
            utility = record_utility(record) if include_outcomes else math.nan
            delta_utility = utility - base_utility if include_outcomes else math.nan
            label = (
                label_from_delta_utility(delta_utility, epsilon=args.epsilon)
                if include_outcomes and record.record_id != base.record_id
                else ("neutral" if include_outcomes else "hidden")
            )
            delta_action = subtract_actions(action, base_action)
            row = _row_from_record(
                record,
                base=base,
                split=args.split,
                action=action,
                base_action=base_action,
                delta_action=delta_action,
                utility=utility,
                base_utility=base_utility,
                delta_utility=delta_utility,
                label=label,
                include_outcomes=include_outcomes,
            )
            max_flat_dim = max(max_flat_dim, len(row["action_flat"]))
            rows.append(row)
            emitted_for_group += 1
            task_counts[str(record.task_id)] += 1
            label_counts[label] += 1
        if emitted_for_group:
            group_hashes.add(_hash_id(group[0].group_id))
            state_hashes.add(_hash_id(group[0].state_hash))
        if len(rows) >= args.shard_size:
            shard_paths.append(_write_shard(np, out_dir, rows, len(shard_paths)))
            rows = []
    if rows:
        shard_paths.append(_write_shard(np, out_dir, rows, len(shard_paths)))

    index = {
        "schema_version": 1,
        "format": "cil_charts_npz",
        "dataset": str(args.dataset),
        "split": args.split,
        "split_policy": args.split_policy,
        "split_fractions": split_fractions,
        "split_seed": args.split_seed,
        "epsilon": args.epsilon,
        "include_outcomes": include_outcomes,
        "audience": "train_retrieval" if args.split == "train" else "evaluator_only",
        "retrieval_index_allowed": args.split == "train",
        "deployment_clean": args.split == "train",
        "base_candidate_types": list(base_candidate_types),
        "num_groups_scanned": group_count,
        "num_groups_skipped_by_split": skipped_by_split,
        "num_groups_exported": len(group_hashes),
        "num_rows": sum(int(item["num_rows"]) for item in shard_paths),
        "max_flat_action_dim": max_flat_dim,
        "label_counts": dict(sorted(label_counts.items())),
        "task_counts": dict(sorted(task_counts.items())),
        "candidate_type_counts": _sum_shard_counter(out_dir, "candidate_type_counts.json"),
        "skip_counts": dict(sorted(counters.items())),
        "group_hashes": sorted(group_hashes),
        "state_hashes": sorted(state_hashes),
        "shards": shard_paths,
        "deployment_candidate_excludes_expert": True,
        "leakage_contract": {
            "train_split_only_for_retrieval": True,
            "nontrain_outcomes_are_evaluator_only": True,
            "deployment_must_not_read_outcomes": args.split != "train",
        },
    }
    index["split_hash"] = _split_hash(index)
    index["content_hash"] = _content_hash(index)
    (out_dir / "index.json").write_text(json.dumps(index, indent=2, sort_keys=True) + "\n")
    print(json.dumps({k: index[k] for k in ("split", "num_groups_exported", "num_rows", "content_hash")}, indent=2))
    return 0


def _row_from_record(
    record: CILRecord,
    *,
    base: CILRecord,
    split: str,
    action: list[list[float]],
    base_action: list[list[float]],
    delta_action: list[list[float]],
    utility: float,
    base_utility: float,
    delta_utility: float,
    label: str,
    include_outcomes: bool,
) -> dict[str, Any]:
    outcome = OutcomeVector.from_reward(record.reward, failure=record.failure)
    seed = record.metadata.get("episode_seed", record.metadata.get("random_seed"))
    branch_family = str(record.candidate_type)
    is_base_branch = record.record_id == base.record_id
    is_expert_branch = branch_family == "expert" or branch_family.endswith("_expert")
    tangent_summary = spline_tangent_summary(delta_action)
    source_policy_name = str(record.metadata.get("source_policy_name", "unknown"))
    metadata = {
        "chart_id": record.group_id,
        "group_id": record.group_id,
        "split": split,
        "state_hash": record.state_hash,
        "task_id": record.task_id,
        "seed": seed,
        "split_id": split,
        "instruction": record.instruction,
        "observation_embedding_path": record.metadata.get("observation_embedding_path"),
        "observation_ref": record.observation_ref,
        "record_id": record.record_id,
        "candidate_type": record.candidate_type,
        "branch_family": branch_family,
        "base_record_id": base.record_id,
        "base_candidate_type": base.candidate_type,
        "source_policy_name": source_policy_name,
        "action_id": record.action_chunk.action_id,
        "base_action_id": base.action_chunk.action_id,
        "scene_id": record.scene_id,
        "rank_within_group": record.rank_within_group,
        "failure_type": record.failure.type if record.failure else None,
        "is_expert_branch": is_expert_branch,
        "is_base_branch": is_base_branch,
        "xi_obj": None,
        "source_dataset": record.metadata.get("source_dataset"),
    }
    return {
        "chart_id": record.group_id,
        "group_id": record.group_id,
        "split": split,
        "state_hash": record.state_hash,
        "task_id": record.task_id,
        "seed": "" if seed is None else str(seed),
        "record_id": record.record_id,
        "candidate_type": record.candidate_type,
        "branch_family": branch_family,
        "source_policy_name": source_policy_name,
        "is_expert_branch": is_expert_branch,
        "is_base_branch": is_base_branch,
        "label": label,
        "label_id": LABEL_TO_ID[label],
        "shape": list(action_shape(action)),
        "action_flat": _flatten(action),
        "base_action_flat": _flatten(base_action),
        "delta_action_flat": _flatten(delta_action),
        "spline_tangent_code": _spline_tangent_code(delta_action),
        "utility": utility,
        "base_utility": base_utility,
        "delta_utility": delta_utility,
        "outcome_vector": [
            outcome.success,
            outcome.progress,
            outcome.contact_quality,
            outcome.safety_violation,
            outcome.task_stage_quality,
            outcome.smoothness,
            outcome.recovery,
        ]
        if include_outcomes
        else [math.nan] * 7,
        "metadata_json": json.dumps(metadata, sort_keys=True),
    }


def _write_shard(np: Any, out_dir: Path, rows: list[dict[str, Any]], shard_index: int) -> dict[str, Any]:
    max_dim = max((len(row["action_flat"]) for row in rows), default=0)
    candidate_type_counts = Counter(str(row["candidate_type"]) for row in rows)
    _write_counter(out_dir / "candidate_type_counts.json", candidate_type_counts)
    path = out_dir / f"charts_{shard_index:05d}.npz"
    np.savez_compressed(
        path,
        chart_id=np.asarray([row["chart_id"] for row in rows]),
        split=np.asarray([row["split"] for row in rows]),
        action=_pad(np, [row["action_flat"] for row in rows], max_dim),
        base_action=_pad(np, [row["base_action_flat"] for row in rows], max_dim),
        delta_action=_pad(np, [row["delta_action_flat"] for row in rows], max_dim),
        spline_tangent_code=np.asarray(
            [row["spline_tangent_code"] for row in rows],
            dtype=np.float32,
        ),
        action_shape=np.asarray([row["shape"] for row in rows], dtype=np.int32),
        utility=np.asarray([row["utility"] for row in rows], dtype=np.float32),
        base_utility=np.asarray([row["base_utility"] for row in rows], dtype=np.float32),
        delta_utility=np.asarray([row["delta_utility"] for row in rows], dtype=np.float32),
        label_id=np.asarray([row["label_id"] for row in rows], dtype=np.int8),
        outcome_vector=np.asarray([row["outcome_vector"] for row in rows], dtype=np.float32),
        group_id=np.asarray([row["group_id"] for row in rows]),
        state_hash=np.asarray([row["state_hash"] for row in rows]),
        task_id=np.asarray([row["task_id"] for row in rows]),
        seed=np.asarray([row["seed"] for row in rows]),
        record_id=np.asarray([row["record_id"] for row in rows]),
        candidate_type=np.asarray([row["candidate_type"] for row in rows]),
        branch_family=np.asarray([row["branch_family"] for row in rows]),
        source_policy_name=np.asarray([row["source_policy_name"] for row in rows]),
        is_expert_branch=np.asarray([row["is_expert_branch"] for row in rows], dtype=bool),
        is_base_branch=np.asarray([row["is_base_branch"] for row in rows], dtype=bool),
        label=np.asarray([row["label"] for row in rows]),
        metadata_json=np.asarray([row["metadata_json"] for row in rows]),
    )
    return {
        "path": path.name,
        "num_rows": len(rows),
        "sha256": _sha256(path),
    }


def _pad(np: Any, vectors: list[list[float]], width: int) -> Any:
    array = np.full((len(vectors), width), np.nan, dtype=np.float32)
    for index, vector in enumerate(vectors):
        array[index, : len(vector)] = np.asarray(vector, dtype=np.float32)
    return array


def _flatten(values: list[list[float]]) -> list[float]:
    return [float(value) for row in values for value in row]


def _spline_tangent_code(delta_action: list[list[float]]) -> list[float]:
    summary = spline_tangent_summary(delta_action)
    code = summary.get("spline_code", {}) if isinstance(summary, dict) else {}
    pieces = [
        code.get("endpoint_delta_position", []),
        code.get("midpoint_delta_position", []),
        code.get("endpoint_delta_rotation_approx", []),
        [code.get("gripper_gate_shift", 0.0)],
        [code.get("gripper_close_strength", 0.0)],
        [code.get("time_scale", 1.0)],
        [code.get("lift_bias", 0.0)],
        code.get("approach_axis_bias", []),
    ]
    flat = [float(value) for piece in pieces for value in piece]
    # Stable 21D keyframe code: start/mid/end full residual rows when available,
    # padded/truncated for common horizon-16, action-dim-7 chunks.
    keyframes: list[float] = []
    if delta_action:
        for row_index in (0, len(delta_action) // 2, len(delta_action) - 1):
            keyframes.extend(float(value) for value in delta_action[row_index])
    if keyframes:
        flat = keyframes
    if len(flat) < 21:
        flat.extend([0.0] * (21 - len(flat)))
    return flat[:21]


def _parse_csv(value: str) -> tuple[str, ...]:
    return tuple(item.strip() for item in value.split(",") if item.strip())


def _parse_split_fractions(value: str) -> dict[str, float]:
    parts = [float(item.strip()) for item in value.split(",") if item.strip()]
    if len(parts) != 3:
        raise ValueError("--split-fractions must contain train,val,test values")
    if any(part < 0.0 for part in parts) or sum(parts) <= 0.0:
        raise ValueError("--split-fractions must be non-negative with positive sum")
    total = sum(parts)
    train, val, test = [part / total for part in parts]
    return {"train": train, "val": val, "test": test}


def _assign_split(
    group_id: str,
    *,
    policy: str,
    split: str,
    fractions: dict[str, float],
    seed: int,
) -> str:
    if policy == "explicit":
        return split
    value = _stable_uniform(group_id, seed=seed)
    train_cut = fractions["train"]
    val_cut = train_cut + fractions["val"]
    if value < train_cut:
        return "train"
    if value < val_cut:
        return "val"
    return "test"


def _stable_uniform(value: str, *, seed: int) -> float:
    digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).digest()
    return int.from_bytes(digest[:8], "big") / float(2**64)


def _hash_id(value: str) -> str:
    return hashlib.sha256(str(value).encode()).hexdigest()


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _content_hash(index: dict[str, Any]) -> str:
    payload = dict(index)
    payload.pop("content_hash", None)
    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()


def _split_hash(index: dict[str, Any]) -> str:
    payload = {
        "split": index.get("split"),
        "split_policy": index.get("split_policy"),
        "split_fractions": index.get("split_fractions"),
        "split_seed": index.get("split_seed"),
        "group_hashes": index.get("group_hashes", []),
        "state_hashes": index.get("state_hashes", []),
    }
    return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()


def _write_counter(path: Path, counts: Counter[str]) -> None:
    existing: Counter[str] = Counter()
    if path.exists():
        existing.update(json.loads(path.read_text()))
    existing.update(counts)
    path.write_text(json.dumps(dict(sorted(existing.items())), indent=2, sort_keys=True) + "\n")


def _sum_shard_counter(out_dir: Path, filename: str) -> dict[str, int]:
    path = out_dir / filename
    if not path.exists():
        return {}
    return {str(key): int(value) for key, value in json.loads(path.read_text()).items()}


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