| """CPU checks of the HB ACT pipeline for one task (no GPU): args -> tyro -> DynamicLaunchConfig (same parse as |
| scripts/train.py), real Psi0 view samples (2 cameras 224x224, raw state 44/78 without velocity, chunk 16), HB v2 |
| augmentation on / native ColorJitter off, and one ACTPolicy forward (ResNet18 ImageNet weights from TORCH_HOME). |
| Usage: <psi0 python> verify_cpu.py --task trolley [--samples 4]""" |
| import argparse |
| import os |
| import sys |
| import tempfile |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| H = Path(__file__).resolve().parent |
| sys.path.insert(0, str(H)) |
| sys.path.insert(0, str(H.parent / "hb_psi0")) |
| FAILS = [] |
|
|
|
|
| def check(c, m): |
| print((" ok " if c else " FAIL ") + m, flush=True) |
| if not c: |
| FAILS.append(m) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--task", default="trolley") |
| ap.add_argument("--samples", type=int, default=4) |
| a = ap.parse_args() |
| import tyro |
| import hb_act_args as AA |
| import hb_act_config as C |
| sys.modules["psi.config.train.hb_act_config"] = C |
| r = AA.HA.receipt(a.task) |
| SD, AD = int(r["state_dim"]), int(r["action_dim"]) |
| args = AA.build_args(a.task, a.samples, 100, 50, 3, tempfile.mkdtemp(prefix="hb_act_verify_"), "verify", 2, resume=False) |
| cfg = tyro.cli(C.DynamicLaunchConfig, config=(tyro.conf.ConsolidateSubcommandArgs,), args=args) |
| m, mt = cfg.model, cfg.data.transform.model |
| print(f"[1] task {a.task}: robot {r.get('robot')} state {SD} action {AD} cams {r['image_keys']}") |
| check((m.action_dim, m.state_dim, m.chunk_size, m.n_action_steps) == (AD, SD, 16, 16), f"model A/S/chunk {m.action_dim}/{m.state_dim}/{m.chunk_size}") |
| check(SD in (44, 78), f"raw state {SD} = 44 (G1) / 78 (H2), no joint velocity") |
| check(not mt.img_aug and mt.hb_aug, "native ColorJitter off, HB v2 aug on") |
| check(tuple(mt.resize.size) == (256, 480) and tuple(mt.center_crop.size) == (224, 224), f"native geometry resize {mt.resize.size} crop {mt.center_crop.size}") |
| check(m.use_vae and m.kl_weight == 10.0 and cfg.train.learning_rate == 1e-4, "VAE kl 10, lr 1e-4 (native)") |
| ds = cfg.data(split="train") |
| check(len(ds) == r["frames"], f"len {len(ds)} == frames {r['frames']}") |
| items = [ds[int(i)] for i in np.linspace(0, len(ds) - 1, a.samples).astype(int)] |
| s = items[0] |
| print(" sample keys", sorted(s)) |
| im = torch.as_tensor(s["observation.images"]) |
| check(tuple(im.shape) == (len(r["image_keys"]), 3, 224, 224), f"images {tuple(im.shape)} = (n_cameras, 3, 224, 224)") |
| check(-2.2 < float(im.min()) and float(im.max()) < 2.7, f"ImageNet-normalised range [{float(im.min()):.2f}, {float(im.max()):.2f}]") |
| st, ac = np.asarray(s["observation.state"]), np.asarray(s["action"]) |
| check(st.shape == (1, SD) and ac.shape == (16, AD), f"state {st.shape} action {ac.shape}") |
| check(np.abs(ac).max() <= 1.0 + 1e-4, f"actions bounds-normalised (max |a| {np.abs(ac).max():.3f})") |
| check(not np.asarray(s["action_is_pad"]).any(), "no action padding") |
| batch = torch.utils.data.default_collate(items) |
| from act.models.act import ACTConfig, ACTPolicy |
| ac_cfg = ACTConfig(chunk_size=m.chunk_size, n_action_steps=m.n_action_steps, action_dim=m.action_dim, state_dim=m.state_dim, |
| dim_model=m.dim_model, n_heads=m.n_heads, dim_feedforward=m.dim_feedforward, |
| feedforward_activation=m.feedforward_activation, n_encoder_layers=m.n_encoder_layers, |
| n_decoder_layers=m.n_decoder_layers, pre_norm=m.pre_norm, dropout=m.dropout, use_vae=m.use_vae, |
| latent_dim=m.latent_dim, n_vae_encoder_layers=m.n_vae_encoder_layers, kl_weight=m.kl_weight, |
| temporal_ensemble_coeff=m.temporal_ensemble_coeff) |
| pol = ACTPolicy(config=ac_cfg, dataset_stats=None) |
| if batch["observation.state"].dim() == 3: |
| batch["observation.state"] = batch["observation.state"].squeeze(1) |
| batch = {k: (v.float() if torch.is_tensor(v) and v.is_floating_point() else v) for k, v in batch.items()} |
| out = pol(batch) |
| check(all(torch.isfinite(torch.as_tensor(v)).all() for v in out.values() if torch.is_tensor(v)), f"ACT forward {{k: float(v) for k, v in out.items() if torch.is_tensor(v)}}".replace("{k", "").replace("}", "") + " " + str({k: round(float(v), 4) for k, v in out.items() if torch.is_tensor(v) or isinstance(v, float)})) |
| print(f" ACT params {sum(p.numel() for p in pol.parameters()) / 1e6:.1f} M") |
| print("VERIFY_" + ("OK" if not FAILS else f"FAILED ({len(FAILS)})")) |
| sys.exit(1 if FAILS else 0) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
