| """Smoke test for 3-PoinTr: load a tiny subset of our prepared Stack-Can data, |
| build the model, run 1 forward pass, verify loss is finite.""" |
| import os, sys, json, glob |
| sys.path.insert(0, "/home/hongyi/PWM/reward_model/3PoinTr_repo") |
|
|
| import torch |
| from util.training_utils.train_util import ( |
| build_common_arg_parser, prepare_datasets_and_loaders, normalize_and_augment, |
| extract_diffusion_kwargs, |
| ) |
| from util.training_utils.arch import TrajectoryModel |
|
|
|
|
| def main(): |
| DATA_ROOT = "/home/hongyi/PWM/external/3pointr_data" |
| SPLIT_FILE = os.path.join(DATA_ROOT, "smoke_test_split.json") |
|
|
| eps = sorted(glob.glob(os.path.join(DATA_ROOT, "no_actions/stack_can/[0-9][0-9][0-9][0-9]"))) |
| if len(eps) < 4: |
| raise RuntimeError(f"Need >=4 episodes built; only {len(eps)} available.") |
| train_eps = eps[:max(2, len(eps) - 1)] |
| val_eps = eps[-1:] |
| payload = { |
| "actions_train": [], "actions_val": [], "actions_test": [], |
| "flows_train": train_eps, "flows_val": val_eps, "flows_test": [], |
| "meta": {"smoke": True} |
| } |
| with open(SPLIT_FILE, "w") as f: |
| json.dump(payload, f, indent=2) |
| print(f"[smoke] split: train={len(train_eps)}, val={len(val_eps)}") |
|
|
| parser = build_common_arg_parser("smoke") |
| args = parser.parse_args([ |
| "--dataset", DATA_ROOT + "/", |
| "--split_file", SPLIT_FILE, |
| "--subsample", "512", |
| "--batch_size", "2", |
| "--epochs", "1", |
| "--num_queries", "10", |
| "--arch_dim", "256", "--arch_depth", "2", "--arch_heads", "4", |
| "--arch_query_depth", "1", "--arch_query_heads", "1", |
| "--flow_data_only", |
| "--tasks", "stack_can", |
| "--lr", "1e-4", |
| ]) |
|
|
| arts = prepare_datasets_and_loaders(args) |
| train_loader = arts.train_loader |
| val_loader = arts.val_loader |
| num_tasks = len(arts.tasks) |
| traj_len = arts.train_set.traj_len |
| print(f"[smoke] traj_len={traj_len}, num_tasks={num_tasks}, train_size={len(arts.train_set)}") |
| dataset_artifacts = type("X", (), {"traj_len": traj_len, "num_tasks": num_tasks})() |
|
|
| diff_kwargs = extract_diffusion_kwargs(args) |
| diff_kwargs["n_action_steps"] = dataset_artifacts.traj_len - 1 |
| model = TrajectoryModel( |
| dim=args.arch_dim, depth=args.arch_depth, num_heads=args.arch_heads, |
| query_depth=args.arch_query_depth, query_heads=args.arch_query_heads, |
| horizon=dataset_artifacts.traj_len - 1, |
| traj_len=dataset_artifacts.traj_len, |
| use_task_conditioning=False, num_tasks=dataset_artifacts.num_tasks, |
| decoder_only_modules=("decoder_only.point_flow_head",), |
| **diff_kwargs, |
| ).cuda().train() |
| print(f"[smoke] model params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M") |
|
|
| batch = next(iter(train_loader)) |
| (x0, cum, ee_pos, ee_ori6d, grip, step_mask, task_ids, action_mask, point_mask) = batch |
| print(f"[smoke] batch: x0={tuple(x0.shape)} cum={tuple(cum.shape)} step_mask={tuple(step_mask.shape)} pt_mask={tuple(point_mask.shape)}") |
|
|
| norm = normalize_and_augment(args, x0, cum, None, None, None, |
| step_mask, action_mask, point_mask, "cuda") |
| feed = { |
| "x0": norm["x0n"], |
| "cum_target": norm["cum_gtn"], |
| "step_mask": norm["step_mask"], |
| "point_mask": norm["point_mask"], |
| "num_steps": norm["step_mask"].shape[1], |
| "flow_loss_w": 10.0, |
| } |
| out = model(feed, mode="train") |
| print(f"[smoke] outputs:", {k: (v.item() if v.dim()==0 else tuple(v.shape)) for k, v in out.items() if torch.is_tensor(v)}) |
| print("[smoke] OK") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
