train.py

"""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_<timestamp>.")
    p.add_argument("--tb-dir", type=str, default="./runs",
                   help="Tensorboard log root. Each run goes under <tb-dir>/<run-name>/.")
    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()