File size: 5,609 Bytes
7e86b86 | 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 146 147 148 149 150 151 152 | #!/usr/bin/env python3
"""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 # noqa: E402
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: # noqa: BLE001 — skip malformed scenes, count them
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()
|