| |
| """Build the dual-rate numeric cache directly from packed val tarballs. |
| |
| Streams scene JSONL files out of data_scenarios_packed/val/<type>.tar.gz |
| (no raw extraction — disk is tight) and writes one compressed .npz per scene: |
| |
| states (T, N, 6) float32 [px, py, vx, vy, angle, angvel] |
| static (N, 8) float32 [is_circle, is_box, log_mass, friction, |
| elasticity, size0/100, size1/100, 0] |
| ids (N,) int32 object ids (for obj_{id} text rendering) |
| gravity (2,) float32 |
| dt () float32 |
| header_text str pre-rendered scene header (baseline format) |
| |
| Scenes are split train/val by seed hash (seed % 10 == 0 -> val). |
| OOD scenario types are never processed here. |
| """ |
|
|
| import argparse |
| import json |
| import math |
| import sys |
| import tarfile |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| sys.path.insert(0, str(Path(__file__).parent.parent)) |
|
|
| from src.training.data_loader import _format_header_text |
|
|
| OOD_TYPES = {"pong", "bowling", "ramp_roll", "angry_birds", "hourglass", "newtons_cradle"} |
|
|
|
|
| def scene_to_arrays(lines): |
| header = json.loads(lines[0]) |
| if header.get("type") != "scene_header": |
| raise ValueError("first line is not scene_header") |
| frames = [json.loads(ln) for ln in lines[1:]] |
|
|
| ids = [obj["id"] for obj in header["objects"]] |
| order = {oid: i for i, oid in enumerate(ids)} |
| n = len(ids) |
| t = len(frames) |
|
|
| states = np.zeros((t, n, 6), dtype=np.float32) |
| for fi, frame in enumerate(frames): |
| for obj in frame["objects"]: |
| i = order[obj["id"]] |
| states[fi, i, 0] = obj["position"]["x"] |
| states[fi, i, 1] = obj["position"]["y"] |
| states[fi, i, 2] = obj["velocity"]["x"] |
| states[fi, i, 3] = obj["velocity"]["y"] |
| states[fi, i, 4] = obj.get("angle", 0.0) |
| states[fi, i, 5] = obj.get("angular_velocity", 0.0) |
|
|
| static = np.zeros((n, 8), dtype=np.float32) |
| for obj in header["objects"]: |
| i = order[obj["id"]] |
| is_circle = 1.0 if obj.get("type") == "circle" else 0.0 |
| mat = obj.get("material", {}) |
| if "radius" in obj: |
| s0 = s1 = float(obj["radius"]) |
| else: |
| s0 = float(obj.get("width", obj.get("size", 0.0)) or 0.0) |
| s1 = float(obj.get("height", obj.get("size", 0.0)) or 0.0) |
| static[i] = [ |
| is_circle, |
| 1.0 - is_circle, |
| math.log(max(mat.get("mass", 1.0), 1e-3)), |
| mat.get("friction", 0.0), |
| mat.get("elasticity", 0.0), |
| s0 / 100.0, |
| s1 / 100.0, |
| 0.0, |
| ] |
|
|
| gravity = np.array([header["gravity"]["x"], header["gravity"]["y"]], dtype=np.float32) |
| dt = np.float32(header["timestep"]) |
| header_text = _format_header_text(header) |
| return header, states, static, np.array(ids, dtype=np.int32), gravity, dt, header_text |
|
|
|
|
| def process_tarball(args): |
| tar_path, out_dir, scenes_per_type = args |
| scenario = Path(tar_path).name.replace(".tar.gz", "") |
| done = skipped = 0 |
| with tarfile.open(tar_path, "r:gz") as tf: |
| for member in tf: |
| if done >= scenes_per_type: |
| break |
| if not member.name.endswith(".jsonl"): |
| continue |
| try: |
| lines = tf.extractfile(member).read().decode().splitlines() |
| header, states, static, ids, gravity, dt, header_text = scene_to_arrays(lines) |
| if states.shape[0] < 50: |
| skipped += 1 |
| continue |
| seed = header["seed"] |
| split = "val" if seed % 10 == 0 else "train" |
| out_path = Path(out_dir) / split / f"{scenario}_{seed}.npz" |
| np.savez_compressed( |
| out_path, |
| states=states, |
| static=static, |
| ids=ids, |
| gravity=gravity, |
| dt=dt, |
| header_text=header_text, |
| ) |
| done += 1 |
| except Exception as e: |
| skipped += 1 |
| if skipped <= 3: |
| print(f"[{scenario}] skip {member.name}: {e}", flush=True) |
| print(f"[{scenario}] cached {done}, skipped {skipped}", flush=True) |
| return scenario, done, skipped |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--packed-dir", default="data_scenarios_packed/val") |
| ap.add_argument("--out", default="data_dualrate/cache") |
| ap.add_argument("--scenes-per-type", type=int, default=2500) |
| ap.add_argument("--workers", type=int, default=12) |
| args = ap.parse_args() |
|
|
| out_dir = Path(args.out) |
| (out_dir / "train").mkdir(parents=True, exist_ok=True) |
| (out_dir / "val").mkdir(parents=True, exist_ok=True) |
|
|
| tarballs = sorted( |
| p for p in Path(args.packed_dir).glob("*.tar.gz") |
| if p.name.replace(".tar.gz", "") not in OOD_TYPES |
| ) |
| print(f"{len(tarballs)} seen-type tarballs, {args.scenes_per_type}/type -> {out_dir}") |
|
|
| jobs = [(str(p), str(out_dir), args.scenes_per_type) for p in tarballs] |
| with Pool(args.workers) as pool: |
| results = pool.map(process_tarball, jobs) |
|
|
| total = sum(d for _, d, _ in results) |
| print(f"TOTAL cached {total} scenes " |
| f"(train {len(list((out_dir / 'train').glob('*.npz')))}, " |
| f"val {len(list((out_dir / 'val').glob('*.npz')))})") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
