auto-sync 2026-07-02T23:07:00Z workspace (part 3)

workspace/scripts/export_cil_charts.py

@@ -23,6 +23,7 @@ from dovla_cil.generation.tangent_targets import (  # noqa: E402
     choose_base_record,
     label_from_delta_utility,
     record_utility,
+    spline_tangent_summary,
     subtract_actions,
 )
 
@@ -45,6 +46,18 @@ def main(argv: list[str] | None = None) -> int:
         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)
@@ -73,6 +86,7 @@ def main(argv: list[str] | None = None) -> int:
         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
@@ -83,6 +97,12 @@ def main(argv: list[str] | None = None) -> int:
     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()
@@ -92,6 +112,7 @@ def main(argv: list[str] | None = None) -> int:
     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():
@@ -101,6 +122,16 @@ def main(argv: list[str] | None = None) -> int:
         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
@@ -160,6 +191,9 @@ def main(argv: list[str] | None = None) -> int:
         "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",
@@ -167,21 +201,25 @@ def main(argv: list[str] | None = None) -> int:
         "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))
@@ -203,36 +241,59 @@ def _row_from_record(
     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,
@@ -253,12 +314,20 @@ def _row_from_record(
 
 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),
@@ -268,8 +337,13 @@ def _write_shard(np: Any, out_dir: Path, rows: list[dict[str, Any]], shard_index
         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]),
     )
@@ -291,10 +365,73 @@ 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()
 
@@ -313,5 +450,32 @@ def _content_hash(index: dict[str, Any]) -> str:
     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())