File size: 2,361 Bytes
87d0bc2 | 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 | """Build a compact training cache from the frozen ogbench-cube-success release.
Reads ONLY the small state fields (observation/action/success) per episode in
frozen subset order, skips the terminal NaN-action row, and uses each episode's
final observation as the (hindsight) goal for goal-conditioned BC.
"""
import argparse, json, os
from pathlib import Path
import numpy as np
BUDGETS = (10, 50, 100, 1000)
def build(root, budget, out):
root = Path(root)
rows = [json.loads(l) for l in (root / "subsets" / f"subset-{budget:04d}.jsonl").read_text().splitlines()]
assert len(rows) == budget, f"expected {budget} rows, got {len(rows)}"
obs_l, act_l, goal_l, succ_l, tid_l = [], [], [], [], []
for ti, row in enumerate(rows):
with np.load(root / "data" / row["archive_file"], allow_pickle=False) as z:
o = z["observation"].astype(np.float32) # (T+1, 28)
a = z["action"].astype(np.float32) # (T+1, 5), last row NaN
s = z["success"].astype(np.bool_) # (T+1,)
T = len(a) - 1 # executed transitions
assert np.isfinite(a[:T]).all(), "unexpected NaN in executed actions"
goal = o[-1] # final/achieved state = goal
obs_l.append(o[:T]); act_l.append(a[:T])
goal_l.append(np.broadcast_to(goal, (T, o.shape[1])))
succ_l.append(s[:T]); tid_l.append(np.full(T, ti, np.int32))
OBS = np.concatenate(obs_l); ACT = np.concatenate(act_l)
GOAL = np.concatenate(goal_l); SUC = np.concatenate(succ_l); TID = np.concatenate(tid_l)
os.makedirs(Path(out).parent, exist_ok=True)
np.savez(out, observation=OBS, action=ACT, goal=GOAL, success=SUC, traj_id=TID)
print(json.dumps({"budget": budget, "episodes": budget, "transitions": int(len(OBS)),
"obs_dim": int(OBS.shape[1]), "act_dim": int(ACT.shape[1]),
"success_transitions": int(SUC.sum()), "cache": out}))
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--root", default="data/bucket/ogbench-cube-success-1000-v1")
ap.add_argument("--budget", type=int, choices=BUDGETS, default=1000)
ap.add_argument("--out", default=None)
a = ap.parse_args()
build(a.root, a.budget, a.out or f"cache/cube_state_b{a.budget:04d}.npz")
|