"""Minimal ACT training entrypoint. Usage: # Vision-only on finger 1 python train.py --data ./episodes \ --state-dim 22 --action-dim 22 \ --image-keys raw_1 --image-shape 3 240 320 --chunk-size 30 # Vision + 6-DoF wrench on finger 1 python train.py --data ./episodes \ --state-dim 22 --action-dim 22 \ --image-keys raw_1 \ --tactile-keys wrench_1 --tactile-dim 6 \ --image-shape 3 240 320 --chunk-size 30 # Two cams (raw + deform) on finger 1 python train.py --data ./episodes \ --state-dim 22 --action-dim 22 \ --image-keys raw_1 deform_1 \ --image-shape 3 240 320 --chunk-size 30 `--smoke` skips data loading and runs on synthetic batches. """ from __future__ import annotations import argparse import json import time from pathlib import Path import torch from torch.utils.data import DataLoader from torch.utils.tensorboard import SummaryWriter from act_policy import ACTConfig, ACTPolicy from data.dataset import DatasetSpec, EpisodicHDF5Dataset, RandomDemoDataset from data.stats import compute_stats def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser() p.add_argument("--data", type=str, default=None, help="Directory of episode_*.hdf5 files.") p.add_argument("--smoke", action="store_true", help="Use random data instead of disk.") p.add_argument("--episode-glob", type=str, default="episode_*.hdf5") p.add_argument("--state-dim", type=int, default=22, help="Per-frame qpos dim.") p.add_argument("--action-dim", type=int, default=22) p.add_argument("--image-keys", nargs="*", default=["raw_1"]) p.add_argument("--image-shape", nargs=3, type=int, default=[3, 240, 320], metavar=("C", "H", "W")) p.add_argument("--tactile-dim", type=int, default=0, help="Per-frame tactile dim (e.g. 30 for 5×wrench).") p.add_argument("--tactile-keys", nargs="*", default=[], help="HDF5 wrench keys to concatenate, e.g. wrench_1 or wrench_0 wrench_1 wrench_2.") p.add_argument("--tactile-image-keys", nargs="*", default=[], help="Per-finger image keys for the in-graph tactile-fusion branch " "(e.g. raw_0 ... raw_4). When set, --tactile-dim is auto-set to " "len(keys) * --tactile-encoder-feature-dim and the dataset emits " "raw image pairs + per-episode bg instead of a precomputed tactile vector.") p.add_argument("--tactile-image-pair", type=int, default=2, help="Number of consecutive frames stacked per sample for the tactile-fusion encoder.") p.add_argument("--tactile-image-group", default="images", help="HDF5 group for tactile image frames. Use 'tactile_imgs' after running " "tools/precompute_tactile_resized.py for much faster I/O.") p.add_argument("--tactile-bg-group", default="tf_bg", help="HDF5 group for per-episode bg images. Use 'tactile_bg_224' alongside " "the resized tactile_imgs.") p.add_argument("--tactile-encoder-modality", type=str, default="9dtact", choices=["9dtact", "gsmini", "sitr", "t3"], help="In-graph tactile encoder: tactile-fusion modality, 'sitr' " "(SITR_base, 768-d cls), or 't3' (T3 mini encoder + trunk; " "--tactile-encoder-ckpt = T3 release dir, feature dim selects size).") p.add_argument("--tactile-encoder-ckpt", type=str, default="/home/allenbi/projects_25/tactile_fusion/checkpoints/pretrain_joint_alignscratch/" "pretrain_multimodal_20260521_005341/checkpoint_step_34000.pth") p.add_argument("--tactile-encoder-pretrain-config", type=str, default="/home/allenbi/projects_25/tactile_fusion/config/model/pretrain.yaml") p.add_argument("--tactile-encoder-feature-dim", type=int, default=192, help="Per-finger feature dim emitted by encoder + shared trunk.") p.add_argument("--tactile-encoder-target-hw", nargs=2, type=int, default=[224, 224], help="Input H W the encoder expects after GPU-side resize.") p.add_argument("--tactile-encoder-lr-scale", type=float, default=0.1, help="LR scale for the tactile-fusion encoder + shared trunk (vs --lr).") p.add_argument("--freeze-tactile-encoder", action="store_true", help="Set requires_grad=False on the tactile-fusion branch (frozen in-graph variant).") p.add_argument("--state-history", type=int, default=1, help="Number of past qpos frames to concatenate as state (1 = current only).") p.add_argument("--tactile-history", type=int, default=1, help="Number of past tactile frames to concatenate as tactile.") p.add_argument("--delta-actions", action="store_true", help="Train on per-step deltas (qpos[t+1]-qpos[t]) instead of absolute " "qpos[t+1] targets. Requires state_dim == action_dim. Inference still " "returns absolute target poses (delta is integrated back inside the policy).") p.add_argument("--state-indices", nargs="*", type=int, default=None, help="Subset of qpos dims to use as state (space-separated ints). When set, " "--state-dim is overridden to len(indices). Useful for training on a " "subset of the hand (e.g. wrist + thumb + index = 13 dims).") p.add_argument("--action-indices", nargs="*", type=int, default=None, help="Subset of action dims to predict. When set, --action-dim is overridden " "to len(indices). For --delta-actions, must equal --state-indices.") p.add_argument("--chunk-size", type=int, default=30) p.add_argument("--n-action-steps", type=int, default=30) p.add_argument("--no-vae", action="store_true") p.add_argument("--use-flow-matching", action="store_true", help="Replace the L1 action regression head with a per-step flow-matching " "MLP head. Usually paired with --no-vae.") p.add_argument("--fm-hidden-dim", type=int, default=512) p.add_argument("--fm-n-mlp-layers", type=int, default=3) p.add_argument("--fm-time-embed-dim", type=int, default=128) p.add_argument("--fm-inference-steps", type=int, default=10, help="ODE-solve steps at inference; baked into the checkpoint config.") p.add_argument("--lr", type=float, default=1e-5) p.add_argument("--lr-backbone", type=float, default=1e-5) p.add_argument("--weight-decay", type=float, default=1e-4) p.add_argument("--lr-warmup-steps", type=int, default=0, help="Linear warmup over the first N steps from --lr-warmup-init to --lr. " "0 disables warmup.") p.add_argument("--lr-warmup-init", type=float, default=0.0, help="Starting LR at step 0 when warmup is enabled. Linear ramp to --lr at " "step --lr-warmup-steps, then flat.") p.add_argument("--lr-cosine", action="store_true", help="After warmup, cosine-decay LR from --lr to --lr-min over the remaining " "(args.steps - args.lr_warmup_steps) steps.") p.add_argument("--lr-min", type=float, default=0.0, help="Floor LR for cosine decay (only used when --lr-cosine).") p.add_argument("--batch-size", type=int, default=8) p.add_argument("--num-workers", type=int, default=4) p.add_argument("--bf16", action="store_true", help="Run the training forward/backward under torch.autocast(bfloat16). " "Weights, optimizer, checkpoints and validation stay fp32.") p.add_argument("--steps", type=int, default=100_000) p.add_argument("--log-every", type=int, default=50) p.add_argument("--save-every", type=int, default=5000) p.add_argument("--out", type=str, default="./checkpoints") p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--resume", type=str, default=None, help="Path to a previous checkpoint (.pt) whose model state_dict is loaded " "before training begins. Optimizer / LR scheduler / step counter start " "fresh, so this acts as a 'fine-tune' from existing weights.") p.add_argument("--run-name", type=str, default=None, help="Subfolder name under --out and --tb-dir. Defaults to task_.") p.add_argument("--tb-dir", type=str, default="./runs", help="Tensorboard log root. Each run goes under //.") p.add_argument("--no-tb", action="store_true", help="Disable tensorboard logging.") p.add_argument("--no-stats", action="store_true", help="Skip the upfront stats pass (e.g. for smoke runs).") p.add_argument("--stats-batch-size", type=int, default=32) p.add_argument("--stats-num-workers", type=int, default=4) p.add_argument("--val-episodes", type=int, default=0, help="Hold out this many trailing episodes (sorted) as a validation set.") p.add_argument("--val-every", type=int, default=200, help="Run validation every N steps. Ignored if --val-episodes == 0.") p.add_argument("--val-batches", type=int, default=0, help="Cap validation to N minibatches (0 = full val pass).") p.add_argument("--patience", type=int, default=10, help="Early stop after this many consecutive val checks with no improvement. " "0 disables early stopping.") return p.parse_args() def _resolve_dims(args) -> tuple[int, int, list[int] | None, list[int] | None]: """Return (state_dim_per_frame, action_dim, state_indices, action_indices), overriding the CLI ``--state-dim`` / ``--action-dim`` with the slice length when ``--state-indices`` / ``--action-indices`` is provided.""" state_indices = list(args.state_indices) if args.state_indices is not None else None action_indices = list(args.action_indices) if args.action_indices is not None else None state_dim = len(state_indices) if state_indices is not None else args.state_dim action_dim = len(action_indices) if action_indices is not None else args.action_dim if args.delta_actions and state_indices is not None and action_indices is not None: if state_indices != action_indices: raise ValueError( "--delta-actions requires --state-indices == --action-indices " f"(got state={state_indices} vs action={action_indices})." ) return state_dim, action_dim, state_indices, action_indices def _resolve_tactile_dim(args) -> int: """If --tactile-image-keys is set, auto-derive tactile_dim from the encoder's feature_dim × n_fingers (overrides --tactile-dim). Otherwise return --tactile-dim as-is.""" if args.tactile_image_keys: return len(args.tactile_image_keys) * int(args.tactile_encoder_feature_dim) return args.tactile_dim def make_spec(args) -> DatasetSpec: state_dim, action_dim, state_indices, action_indices = _resolve_dims(args) return DatasetSpec( state_dim=state_dim, action_dim=action_dim, image_keys=list(args.image_keys), image_shape=tuple(args.image_shape), tactile_dim=_resolve_tactile_dim(args), tactile_keys=list(args.tactile_keys), chunk_size=args.chunk_size, state_history=args.state_history, tactile_history=args.tactile_history, delta_actions=args.delta_actions, state_indices=state_indices, action_indices=action_indices, tactile_image_keys=list(args.tactile_image_keys), tactile_image_pair=args.tactile_image_pair, tactile_image_group=args.tactile_image_group, tactile_bg_group=args.tactile_bg_group, ) def make_policy(args) -> ACTPolicy: # The dataset returns state/tactile concatenated across history frames, so the # policy's input dims are scaled by the history depth. state_dim, action_dim, state_indices, action_indices = _resolve_dims(args) tactile_dim_per_frame = _resolve_tactile_dim(args) cfg = ACTConfig( state_dim=state_dim * max(1, args.state_history), action_dim=action_dim, image_keys=list(args.image_keys), image_shape=tuple(args.image_shape), tactile_dim=tactile_dim_per_frame * max(1, args.tactile_history), chunk_size=args.chunk_size, n_action_steps=args.n_action_steps, use_vae=not args.no_vae, optimizer_lr=args.lr, optimizer_lr_backbone=args.lr_backbone, optimizer_weight_decay=args.weight_decay, state_history=args.state_history, tactile_history=args.tactile_history, use_flow_matching=args.use_flow_matching, fm_hidden_dim=args.fm_hidden_dim, fm_n_mlp_layers=args.fm_n_mlp_layers, fm_time_embed_dim=args.fm_time_embed_dim, fm_inference_steps=args.fm_inference_steps, delta_actions=args.delta_actions, state_indices=state_indices, action_indices=action_indices, tactile_image_keys=list(args.tactile_image_keys), tactile_image_pair=args.tactile_image_pair, tactile_encoder_modality=args.tactile_encoder_modality if args.tactile_image_keys else None, tactile_encoder_ckpt=args.tactile_encoder_ckpt if args.tactile_image_keys else None, tactile_encoder_pretrain_config=args.tactile_encoder_pretrain_config if args.tactile_image_keys else None, tactile_encoder_target_hw=tuple(args.tactile_encoder_target_hw), tactile_encoder_feature_dim=args.tactile_encoder_feature_dim, tactile_encoder_lr_scale=args.tactile_encoder_lr_scale, tactile_encoder_freeze=args.freeze_tactile_encoder, ) return ACTPolicy(cfg) def make_dataset(args) -> torch.utils.data.Dataset: spec = make_spec(args) if args.smoke or args.data is None: return RandomDemoDataset(spec, length=max(args.batch_size * 4, 64)) return EpisodicHDF5Dataset(args.data, spec, episode_glob=args.episode_glob) def make_train_val_datasets(args) -> tuple[torch.utils.data.Dataset, torch.utils.data.Dataset | None]: """Build (train, val) datasets. When --val-episodes > 0, the last N episode files (sorted) are held out for validation; otherwise val is None.""" spec = make_spec(args) if args.smoke or args.data is None: return RandomDemoDataset(spec, length=max(args.batch_size * 4, 64)), None import glob, os all_files = sorted(glob.glob(os.path.join(args.data, args.episode_glob))) if not all_files: raise FileNotFoundError(f"No HDF5 episodes under {args.data!r} matching {args.episode_glob!r}.") n_val = max(0, args.val_episodes) if n_val == 0: return EpisodicHDF5Dataset(args.data, spec, episode_glob=args.episode_glob), None if n_val >= len(all_files): raise ValueError(f"--val-episodes={n_val} but only {len(all_files)} episodes available.") train_files = all_files[:-n_val] val_files = all_files[-n_val:] train_ds = EpisodicHDF5Dataset(root=None, spec=spec, files=train_files) val_ds = EpisodicHDF5Dataset(root=None, spec=spec, files=val_files) return train_ds, val_ds @torch.no_grad() def validate(policy, val_loader, device, max_batches: int = 0) -> dict[str, float]: """Run policy in eval mode over val_loader, return mean L1 (and KLD if present).""" was_training = policy.training policy.eval() total_l1 = total_kld = 0.0 n_l1 = n_kld = 0 for i, batch in enumerate(val_loader): if max_batches and i >= max_batches: break batch = move_to(batch, device) _, info = policy(batch) total_l1 += float(info["l1_loss"]); n_l1 += 1 if "kld_loss" in info: total_kld += float(info["kld_loss"]); n_kld += 1 if was_training: policy.train() out = {"val_l1": total_l1 / max(n_l1, 1)} if n_kld: out["val_kld"] = total_kld / n_kld return out def move_to(batch: dict[str, torch.Tensor], device: str) -> dict[str, torch.Tensor]: return {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()} def _default_run_name(args) -> str: if args.data: task = Path(args.data).name or "run" else: task = "smoke" return f"{task}_{time.strftime('%Y%m%d_%H%M%S')}" def main() -> None: args = parse_args() run_name = args.run_name or _default_run_name(args) out_dir = Path(args.out) / run_name out_dir.mkdir(parents=True, exist_ok=True) writer = None if not args.no_tb: tb_dir = Path(args.tb_dir) / run_name tb_dir.mkdir(parents=True, exist_ok=True) writer = SummaryWriter(log_dir=str(tb_dir)) print(f"[train] tensorboard: {tb_dir}") # Persist a small run manifest next to the checkpoints for reproducibility. with (out_dir / "args.json").open("w") as fh: json.dump(vars(args) | {"run_name": run_name}, fh, indent=2, default=str) policy = make_policy(args).to(args.device) resume_ckpt = None if args.resume: print(f"[train] resuming model weights from {args.resume}") resume_ckpt = torch.load(args.resume, map_location=args.device, weights_only=False) missing, unexpected = policy.load_state_dict(resume_ckpt["model"], strict=False) if missing or unexpected: print(f" missing={missing}\n unexpected={unexpected}") has_opt = "optimizer" in resume_ckpt has_sched = "scheduler" in resume_ckpt print(f" loaded at original step={resume_ckpt.get('step', '?')}; " f"optimizer state {'will be restored' if has_opt else 'NOT in checkpoint (fresh Adam moments)'}; " f"scheduler state {'will be restored' if has_sched else 'NOT in checkpoint'}.") policy.train() dataset, val_dataset = make_train_val_datasets(args) print(f"[train] run={run_name}") print(f"[train] train dataset={type(dataset).__name__} size={len(dataset)} device={args.device}") if val_dataset is not None: print(f"[train] val dataset={type(val_dataset).__name__} size={len(val_dataset)} " f"({args.val_episodes} episode(s) held out)") state_dim_eff, action_dim_eff, state_idx_eff, action_idx_eff = _resolve_dims(args) tactile_dim_eff = _resolve_tactile_dim(args) print(f"[train] obs: state_dim={state_dim_eff} tactile_dim={tactile_dim_eff} " f"image_keys={args.image_keys} tactile_keys={args.tactile_keys}") if args.tactile_image_keys: print(f"[train] tactile_image_keys={args.tactile_image_keys} " f"modality={args.tactile_encoder_modality} " f"freeze_encoder={args.freeze_tactile_encoder} " f"encoder_lr_scale={args.tactile_encoder_lr_scale}") print(f"[train] action_dim={action_dim_eff} chunk_size={args.chunk_size}") if state_idx_eff is not None: print(f"[train] state_indices={state_idx_eff}") if action_idx_eff is not None: print(f"[train] action_indices={action_idx_eff}") if not args.no_stats and not args.smoke: print("[stats] computing mean/std over the training set...") t_stats0 = time.time() # When tactile_image_keys is set, the encoder produces the tactile # vector inside ACTPolicy — there's no precomputed tactile in the # batch to normalize. Skip the tactile stat pass in that case. has_tactile_for_stats = _resolve_tactile_dim(args) > 0 and not args.tactile_image_keys # When using tactile_image_keys, the dataset would otherwise load big # image tensors per sample just to feed compute_stats — wastes ~40 min # on a 40K-frame dataset. Temporarily disable image loading on the # underlying spec for the duration of the stats pass. _saved_tik = None if args.tactile_image_keys and hasattr(dataset, "spec"): _saved_tik = list(dataset.spec.tactile_image_keys) dataset.spec.tactile_image_keys = [] try: stats = compute_stats( dataset, batch_size=args.stats_batch_size, num_workers=args.stats_num_workers, has_state=args.state_dim > 0, has_tactile=has_tactile_for_stats, ) finally: if _saved_tik is not None: dataset.spec.tactile_image_keys = _saved_tik if stats.state_mean is not None: policy.set_state_stats(stats.state_mean, stats.state_std) print(f"[stats] state mean[:4]={stats.state_mean[:4].tolist()} " f"std[:4]={stats.state_std[:4].tolist()}") if stats.tactile_mean is not None: policy.set_tactile_stats(stats.tactile_mean, stats.tactile_std) print(f"[stats] tactile mean={stats.tactile_mean.tolist()} std={stats.tactile_std.tolist()}") policy.set_action_stats(stats.action_mean, stats.action_std) print(f"[stats] action mean[:4]={stats.action_mean[:4].tolist()} " f"std[:4]={stats.action_std[:4].tolist()}") print(f"[stats] done in {time.time() - t_stats0:.1f}s") loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=args.device == "cuda", drop_last=True, persistent_workers=args.num_workers > 0, ) val_loader = None if val_dataset is not None: val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=max(0, args.num_workers // 2), pin_memory=args.device == "cuda", drop_last=False, persistent_workers=args.num_workers > 0, ) optimizer = torch.optim.AdamW(policy.get_param_groups(), weight_decay=args.weight_decay) # LR schedule: optional linear warmup followed by optional cosine decay. LambdaLR multiplies # each param group's base LR by lambda(step). if args.lr_warmup_steps > 0 or args.lr_cosine: import math target_lrs = [pg["lr"] for pg in optimizer.param_groups] decay_steps = max(args.steps - args.lr_warmup_steps, 1) def lr_factor(step_i: int, target_lr: float) -> float: # Warmup: linear from --lr-warmup-init to target_lr. if step_i < args.lr_warmup_steps: start = args.lr_warmup_init / target_lr if target_lr > 0 else 0.0 return start + (1.0 - start) * (step_i / max(args.lr_warmup_steps, 1)) # Post-warmup: cosine decay to --lr-min if requested, else flat at target_lr. if not args.lr_cosine: return 1.0 progress = (step_i - args.lr_warmup_steps) / decay_steps progress = min(max(progress, 0.0), 1.0) end_ratio = args.lr_min / target_lr if target_lr > 0 else 0.0 return end_ratio + (1.0 - end_ratio) * 0.5 * (1.0 + math.cos(math.pi * progress)) lambdas = [(lambda lr_target: (lambda s: lr_factor(s, lr_target)))(t) for t in target_lrs] scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambdas) decay_str = f" then cosine to {args.lr_min:.2e}" if args.lr_cosine else "" print(f"[train] LR: warmup {args.lr_warmup_init:.2e} → {target_lrs} over " f"{args.lr_warmup_steps} steps{decay_str}") else: scheduler = None # If resuming, restore optimizer / scheduler state if present in the checkpoint. if resume_ckpt is not None: if "optimizer" in resume_ckpt: optimizer.load_state_dict(resume_ckpt["optimizer"]) print(f"[train] restored optimizer state ({len(optimizer.param_groups)} groups)") if scheduler is not None and "scheduler" in resume_ckpt: scheduler.load_state_dict(resume_ckpt["scheduler"]) print(f"[train] restored scheduler state") best_val = float("inf") bad_evals = 0 early_stop = False def save_ckpt(name: str, step_i: int): ckpt = out_dir / name payload = { "step": step_i, "model": policy.state_dict(), "config": policy.config.__dict__, "optimizer": optimizer.state_dict(), } if scheduler is not None: payload["scheduler"] = scheduler.state_dict() torch.save(payload, ckpt) return ckpt step = 0 t0 = time.time() while step < args.steps and not early_stop: for batch in loader: if step >= args.steps or early_stop: break batch = move_to(batch, args.device) with torch.autocast(device_type="cuda" if "cuda" in args.device else "cpu", dtype=torch.bfloat16, enabled=args.bf16): loss, info = policy(batch) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() if scheduler is not None: scheduler.step() if writer is not None: writer.add_scalar("train/loss", float(info["loss"]), step) writer.add_scalar("train/lr", optimizer.param_groups[0]["lr"], step) writer.add_scalar("train/l1_loss", float(info["l1_loss"]), step) if "kld_loss" in info: writer.add_scalar("train/kld_loss", float(info["kld_loss"]), step) if step % args.log_every == 0: elapsed = time.time() - t0 it_per_s = step / max(elapsed, 1e-6) msg = f"step {step:>6d} | loss={info['loss']:.4f} l1={info['l1_loss']:.4f}" if "kld_loss" in info: msg += f" kld={info['kld_loss']:.4f}" msg += f" | {it_per_s:.2f} it/s" print(msg, flush=True) if writer is not None: writer.add_scalar("train/it_per_s", it_per_s, step) if args.save_every > 0 and step > 0 and step % args.save_every == 0: ckpt = save_ckpt(f"act_{step:07d}.pt", step) print(f"[train] saved {ckpt}") # Validation + early-stop bookkeeping if val_loader is not None and args.val_every > 0 and step > 0 and step % args.val_every == 0: vinfo = validate(policy, val_loader, args.device, max_batches=args.val_batches) if writer is not None: writer.add_scalar("val/l1_loss", vinfo["val_l1"], step) if "val_kld" in vinfo: writer.add_scalar("val/kld_loss", vinfo["val_kld"], step) improved = vinfo["val_l1"] < best_val - 1e-6 if improved: best_val = vinfo["val_l1"] bad_evals = 0 save_ckpt("act_best.pt", step) tag = " * BEST (saved act_best.pt)" else: bad_evals += 1 tag = f" no-improve {bad_evals}/{args.patience}" kld_str = f" val_kld={vinfo['val_kld']:.4f}" if "val_kld" in vinfo else "" print(f"[val ] step {step:>6d} | val_l1={vinfo['val_l1']:.4f}{kld_str} " f"| best={best_val:.4f}{tag}", flush=True) if args.patience > 0 and bad_evals >= args.patience: print(f"[train] early-stop triggered at step {step} " f"(no val improvement in {args.patience} checks)", flush=True) early_stop = True step += 1 final = save_ckpt("act_final.pt", step) print(f"[train] done. saved {final}") if val_loader is not None and best_val < float("inf"): print(f"[train] best val_l1 = {best_val:.4f} (act_best.pt)") if writer is not None: writer.flush() writer.close() if __name__ == "__main__": main()