File size: 1,929 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 | """Cache builder for OGBench play datasets (scene, puzzle-3x3).
Play trajectories do not terminate at a goal, so the goal for each transition is
a hindsight future observation sampled within the same trajectory (fixed seed).
Nested budgets use the same fixed-permutation-prefix scheme as the cube release.
"""
import argparse, json, os
import numpy as np, ogbench
from common import get_traj_bounds, nested_traj_order
ap = argparse.ArgumentParser()
ap.add_argument("--dataset", required=True) # e.g. scene-play-v0
ap.add_argument("--budget", type=int, required=True) # trajectories
ap.add_argument("--split_seed", type=int, default=0)
ap.add_argument("--goal_seed", type=int, default=0)
ap.add_argument("--outdir", default="cache")
a = ap.parse_args()
_, train_ds, _ = ogbench.make_env_and_datasets(a.dataset, compact_dataset=False)
obs = train_ds["observations"].astype(np.float32)
act = train_ds["actions"].astype(np.float32)
term = train_ds["terminals"].astype(np.float32)
bounds = get_traj_bounds(term)
order = nested_traj_order(len(bounds), seed=a.split_seed)
sel = np.sort(order[:a.budget])
rng = np.random.default_rng(a.goal_seed)
O, A, G = [], [], []
for ti in sel:
s, e = bounds[int(ti)]
idx = np.arange(s, e)
# hindsight future goal: uniform index in [t, e-1] for each t
u = rng.random(len(idx))
gj = (idx + (u * (e - idx)).astype(np.int64)).clip(max=e - 1)
O.append(obs[idx]); A.append(act[idx]); G.append(obs[gj])
O = np.concatenate(O); A = np.concatenate(A); G = np.concatenate(G)
os.makedirs(a.outdir, exist_ok=True)
tag = a.dataset.replace("-play-v0", "")
out = f"{a.outdir}/{tag}_state_b{a.budget:04d}.npz"
np.savez(out, observation=O, action=A, goal=G)
print(json.dumps({"dataset": a.dataset, "budget": a.budget, "trajectories": len(sel),
"transitions": int(len(O)), "obs_dim": int(O.shape[1]),
"act_dim": int(A.shape[1]), "cache": out}))
|