| """CPU check of the GR00T data path for one task view: loader shapes, v2 augmentation, processor patch import.""" |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| H = Path(__file__).resolve().parent |
| sys.path.insert(0, str(H)) |
| os.environ["HB_PATCH_ONLY"] = "1" |
| import hb_launch |
| import hb_aug_v2 as A |
| from gr00t.data.dataset.lerobot_episode_loader import LeRobotEpisodeLoader |
| from gr00t.model.gr00t_n1d7.processing_gr00t_n1d7 import Gr00tN1d7Processor |
|
|
| view = Path(sys.argv[1]) |
| if view.name == "a18": |
| import h2_body31_hand44_h16 as C |
| cfg = C.h2_body31_hand44_h16 |
| else: |
| import g1_feedback76_h16 as C |
| cfg = C.feedback76_h16 |
| loader = LeRobotEpisodeLoader(view, cfg) |
| df = loader[0] |
| cols = list(df.columns) |
| row = df.iloc[20] |
| out = {"episodes": len(loader), "columns": cols[:40], "rows_ep0": len(df)} |
| state_dims = {c: int(np.asarray(row[c]).size) for c in cols if c.startswith("state.")} |
| action_dims = {c: int(np.asarray(row[c]).size) for c in cols if c.startswith("action.")} |
| out.update(state_dims=state_dims, action_dims=action_dims, state_total=sum(state_dims.values()), |
| action_total=sum(action_dims.values())) |
| vids = [c for c in cols if c.startswith("video.")] |
| imgs = [np.asarray(row[c]) for c in vids] |
| out["video"] = {c: list(i.shape) for c, i in zip(vids, imgs)} |
| x = torch.stack([torch.from_numpy(i).permute(2, 0, 1).float().div(255) for i in imgs])[None] |
| rng = np.random.default_rng(0) |
| boosts = {} |
| for _ in range(60): |
| y, p = A.augment(x, rng) |
| boosts[p["boost"]] = boosts.get(p["boost"], 0) + 1 |
| assert y.shape == x.shape and torch.isfinite(y).all() and float(y.min()) >= 0 and float(y.max()) <= 1 |
| out["aug_boost_counts_60"] = boosts |
| out["mean_abs_change_last"] = round(float((y - x).abs().mean()), 4) |
| q = np.zeros((1, 29)) |
| dq = A.perturb_joints(q, rng) |
| out["proprio_max_abs_deg"] = round(float(np.degrees(np.abs(dq).max())), 3) |
| out["patched"] = Gr00tN1d7Processor._get_vlm_inputs.__name__ == "_get_vlm_inputs" and \ |
| Gr00tN1d7Processor.__call__.__module__ == "hb_launch" |
| print("CPU_VERIFY", json.dumps(out), flush=True) |
|
|
