File size: 4,992 Bytes
dae7cbe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
Generate the proof-gated planned comparison dataset.
The dataset contains one calibration push per episode and a held-out query
displacement target under a fixed query force.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
from planned_comparison_common import (
OOD_CAMERAS,
TRAIN_CAMERAS,
generate_episode,
)
def write_episode(path: Path, episode: dict) -> None:
np.savez_compressed(
path,
rgb_pre=episode["rgb_pre"],
rgb_post=episode["rgb_post"],
camera_intrinsics=episode["camera_intrinsics"],
camera_extrinsics=episode["camera_extrinsics"],
calibration_force=episode["calibration_force"],
query_force=episode["query_force"],
structured_pre=episode["structured_pre"],
structured_post=episode["structured_post"],
structured_delta_x=episode["structured_delta_x"],
structured_mobility=episode["structured_mobility"],
target_query_dx=episode["target_query_dx"],
detected_pre_uv=episode["detected_pre_uv"],
detected_post_uv=episode["detected_post_uv"],
mass=episode["mass"],
damping=episode["damping"],
gt_pre_world=episode["gt_pre_world"],
gt_post_world=episode["gt_post_world"],
gt_delta_x=episode["gt_delta_x"],
)
def generate_split(
rng: np.random.Generator,
output_dir: Path,
split: str,
camera_name: str,
count: int,
starting_index: int,
) -> tuple[list[dict], int]:
manifest_entries: list[dict] = []
index = starting_index
extraction_errors = []
for _ in range(count):
episode = generate_episode(rng, camera_name)
file_name = f"{split}_{camera_name}_{index:05d}.npz"
write_episode(output_dir / file_name, episode)
extraction_errors.append(abs(float(episode["structured_delta_x"]) - float(episode["gt_delta_x"])))
manifest_entries.append(
{
"file": file_name,
"split": split,
"camera_name": camera_name,
"material_name": episode["material_name"],
"calibration_force": float(episode["calibration_force"]),
"query_force": float(episode["query_force"]),
"target_query_dx": float(episode["target_query_dx"]),
}
)
index += 1
if extraction_errors:
avg_error = float(np.mean(extraction_errors))
max_error = float(np.max(extraction_errors))
print(
f"[generate_planned_comparison_data] split={split} camera={camera_name} "
f"count={count} avg_state_dx_error={avg_error:.6f} max_state_dx_error={max_error:.6f}"
)
return manifest_entries, index
def main() -> None:
parser = argparse.ArgumentParser(description="Generate planned comparison episodes.")
parser.add_argument("--output_dir", default="training/data/planned_comparison")
parser.add_argument("--train_count", type=int, default=300)
parser.add_argument("--val_count", type=int, default=80)
parser.add_argument("--test_id_count", type=int, default=80)
parser.add_argument("--test_ood_per_camera", type=int, default=120)
parser.add_argument("--seed", type=int, default=0)
args = parser.parse_args()
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(args.seed)
manifest: list[dict] = []
index = 0
for camera_name in TRAIN_CAMERAS:
entries, index = generate_split(rng, output_dir, "train", camera_name, args.train_count, index)
manifest.extend(entries)
entries, index = generate_split(rng, output_dir, "val", camera_name, args.val_count, index)
manifest.extend(entries)
entries, index = generate_split(rng, output_dir, "test_id", camera_name, args.test_id_count, index)
manifest.extend(entries)
for camera_name in OOD_CAMERAS:
entries, index = generate_split(
rng,
output_dir,
"test_ood",
camera_name,
args.test_ood_per_camera,
index,
)
manifest.extend(entries)
with open(output_dir / "manifest.json", "w", encoding="utf-8") as handle:
json.dump(manifest, handle, indent=2)
metadata = {
"seed": args.seed,
"train_cameras": list(TRAIN_CAMERAS),
"ood_cameras": list(OOD_CAMERAS),
"counts": {
"train_per_camera": args.train_count,
"val_per_camera": args.val_count,
"test_id_per_camera": args.test_id_count,
"test_ood_per_camera": args.test_ood_per_camera,
},
}
with open(output_dir / "metadata.json", "w", encoding="utf-8") as handle:
json.dump(metadata, handle, indent=2)
print(f"[generate_planned_comparison_data] wrote {len(manifest)} episodes to {output_dir}")
if __name__ == "__main__":
main()
|