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