scripts/train_bridge_cte.py

"""Standalone eight-GPU Bridge Stage-1 CTE training, with no policy/release CTE.

The released effect-v3 objective is used, not an exact reproduction of the
paper's MSE objective. See docs/bridge_cte_stage1.md for the full contract.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import shutil
import signal
import sys
import time
import uuid
from dataclasses import asdict
from datetime import datetime, timedelta
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import numpy as np
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader

from cosmos_framework.model.zeva.causal_transition_encoder import CausalTransitionEncoder, CausalTransitionEncoderConfig
from cosmos_framework.model.zeva.cte_losses import CTELossConfig
from scripts.bridge_cte_data import BridgeCTEDataset, PairedEpisodeSampler, ValidationSampler, collate_episodes, load_metadata
from scripts.bridge_cte_loss import bridge_cte_loss
from scripts.widowx_contract import digest_file, fingerprint


def atomic_json(path, payload):
    path = Path(path)
    temporary = path.with_name(path.name + f".{os.getpid()}.tmp")
    temporary.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n")
    temporary.replace(path)


def model_digest(model):
    digest = hashlib.sha256()
    for name, tensor in model.state_dict().items():
        digest.update(name.encode())
        digest.update(tensor.detach().cpu().contiguous().numpy().tobytes())
    return digest.hexdigest()


class FrozenFrameEncoder:
    def __init__(self, path, device, microbatch=4):
        from cosmos_framework.model.generator.tokenizers.wan2pt2_vae_4x16x16 import WanVAE
        self.vae = WanVAE(vae_pth=str(path), device=device, dtype=torch.bfloat16,
                          temporal_window=4, encode_exact_durations=[1])
        self.device, self.microbatch = device, microbatch
        if any(p.requires_grad for p in self.vae.model.parameters()):
            raise AssertionError("VAE must be frozen")

    @torch.no_grad()
    def __call__(self, batch):
        valid = batch["valid"].to(self.device, non_blocking=True)
        latents = []
        # Each single-frame encode resets the VAE's temporal cache. This is
        # intentionally the same operation online and offline, with no future
        # frames; resizing is not needed for Bridge's native 256x256 images.
        for chunk in batch["frames"].split(self.microbatch):
            pixels = chunk.to(self.device, dtype=torch.bfloat16, non_blocking=True) / 127.5 - 1
            latent = self.vae.encode(pixels.unsqueeze(2)).squeeze(2)
            if tuple(latent.shape[1:]) != (48, 16, 16):
                raise ValueError(f"Unexpected VAE output: {latent.shape}")
            latents.append(latent)
        values = torch.cat(latents)
        if not torch.isfinite(values).all():
            raise FloatingPointError("Nonfinite frozen VAE features")
        padded = values.new_zeros((*valid.shape, 48, 16, 16))
        padded[valid] = values
        return padded


def init_wandb(args, config):
    if args.wandb_mode == "disabled":
        return None
    import wandb
    id_path = args.output / "wandb_run_id.json"
    run_id = json.loads(id_path.read_text())["id"] if id_path.exists() else uuid.uuid4().hex[:8]
    atomic_json(id_path, {"id": run_id})
    kwargs = dict(project=args.wandb_project, entity=os.environ.get("WANDB_ENTITY"),
                  name=args.output.name, group="bridge_cte_stage1", id=run_id,
                  dir=str(args.output), config=config, resume="allow",
                  settings=wandb.Settings(init_timeout=20, base_url=os.environ.get("WANDB_BASE_URL", "https://api.wandb.ai")))
    try:
        run = wandb.init(mode=args.wandb_mode, **kwargs)
    except wandb.errors.Error as error:
        if args.wandb_mode != "online":
            raise
        print(f"WANDB_OFFLINE_FALLBACK error_type={type(error).__name__}", flush=True)
        wandb.teardown()
        run = wandb.init(mode="offline", **kwargs)
    atomic_json(args.output / "wandb_status.json", dict(mode=run.settings.mode, project=run.project,
                                                       entity=run.entity, id=run.id, directory=run.dir))
    return run


def build_loader(dataset, sampler, args, rank):
    # Private loader RNG: constructing/restarting a loader never consumes the
    # model's random stream. There are no stochastic data augmentations.
    generator = torch.Generator().manual_seed(args.seed + 10000 + rank)
    kwargs = dict(dataset=dataset, batch_sampler=sampler, collate_fn=collate_episodes,
                  num_workers=args.workers, pin_memory=True, generator=generator)
    if args.workers:
        kwargs.update(prefetch_factor=2, multiprocessing_context="spawn")
    return DataLoader(**kwargs)


def move_batch(batch, device):
    return {key: value.to(device, non_blocking=True) for key, value in batch.items() if key != "frames"}


def reduce_metrics(metrics, device):
    keys = sorted(metrics)
    values = torch.stack([metrics[k].detach().float() for k in keys]).to(device)
    if dist.is_initialized():
        dist.all_reduce(values)
        values /= dist.get_world_size()
    return dict(zip(keys, values.cpu().tolist(), strict=True))


@torch.no_grad()
def validate(model, encoder, dataset, args, world, rank, device):
    model.eval()
    sampler = ValidationSampler(dataset.episodes, args.batch_size, world, rank, args.seed + 70000)
    loader = build_loader(dataset, sampler, args, rank)
    totals, batches = {}, 0
    for raw in loader:
        inputs = encoder(raw)
        batch = move_batch(raw, device)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            output = model(inputs, batch["actions"], batch["valid"], batch["transition_valid"])
        losses = bridge_cte_loss(output, batch["actions"], batch["valid"], batch["semantic_ids"], batch["episode_ids"])
        for key, value in losses.items():
            totals[key] = totals.get(key, 0) + value.detach()
        batches += 1
    metrics = reduce_metrics({key: value / batches for key, value in totals.items()}, device)
    model.train()
    return metrics


def rng_state():
    return dict(python=random.getstate(), numpy=np.random.get_state(),
                torch=torch.get_rng_state(), cuda=torch.cuda.get_rng_state())


def restore_rng(state):
    random.setstate(state["python"])
    np.random.set_state(state["numpy"])
    torch.set_rng_state(state["torch"])
    torch.cuda.set_rng_state(state["cuda"])


def save_checkpoint(model, optimizer, args, contract, step, epoch, batch_cursor, world, rank):
    states = [None] * world
    if dist.is_initialized():
        dist.all_gather_object(states, rng_state())
    else:
        states[0] = rng_state()
    if rank == 0:
        directory = args.output / "checkpoints"
        directory.mkdir(exist_ok=True)
        path = directory / f"cte_step_{step:09d}.pt"
        if not path.exists():
            if shutil.disk_usage(directory).free < args.min_free_gb * 2**30:
                raise OSError("Disk safety floor reached; no checkpoint is automatically deleted")
            payload = dict(format="bridge_cte_stage1_v1", model_config=model.cfg.to_dict(),
                           model=model.state_dict(), optimizer=optimizer.state_dict(),
                           step=step, epoch=epoch, batch_cursor=batch_cursor, rng_states=states,
                           contract=contract, contract_fingerprint=fingerprint(contract),
                           training_args=vars(args) | {"output": str(args.output)})
            temporary = path.with_suffix(f".{os.getpid()}.tmp")
            torch.save(payload, temporary)
            with temporary.open("rb") as stream:
                os.fsync(stream.fileno())
            temporary.replace(path)
        atomic_json(args.output / "latest_checkpoint.json", dict(path=str(path), step=step,
                                                                epoch=epoch, batch_cursor=batch_cursor))
        print(f"CHECKPOINT_SAVED step={step} path={path} bytes={path.stat().st_size}", flush=True)
    if dist.is_initialized():
        dist.barrier()


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset", default="/home/gpu4/tianyi/TTTVLA/datasets/bridge_orig_1.0.0_lerobot")
    parser.add_argument("--vae", default="/home/gpu4/tianyi/models/wan22_vae/Wan2.2_VAE.pth")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--resume", help="Explicit Stage-1 checkpoint, never a RoboCasa CTE")
    parser.add_argument("--batch-size", type=int, default=8, help="Episodes per GPU, including paired partners")
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--vae-microbatch", type=int, default=4)
    parser.add_argument("--lr", type=float, default=0.0001)
    parser.add_argument("--warmup-steps", type=int, default=500)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--max-steps", type=int, default=2_147_483_647)
    parser.add_argument("--save-every", type=int, default=10000)
    parser.add_argument("--validate-every", type=int, default=2000)
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--val-fraction", type=float, default=0.02)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--memory-fraction", type=float, default=0.125, help="PyTorch allocator budget as a fraction of total GPU RAM")
    parser.add_argument("--min-free-gb", type=float, default=2.0)
    parser.add_argument("--wandb-mode", choices=("online", "offline", "disabled"), default=os.environ.get("WANDB_MODE", "online"))
    parser.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "Zeva"))
    parser.add_argument("--skip-initial-validation", action="store_true", help="Smoke checks only")
    args = parser.parse_args()
    if min(args.batch_size, args.vae_microbatch, args.warmup_steps, args.max_steps,
           args.save_every, args.validate_every, args.log_every) < 1 or args.batch_size % 2:
        parser.error("Positive intervals and an even per-GPU batch size are required")
    if not 0 < args.memory_fraction <= 1 or not 0 < args.val_fraction < 1:
        parser.error("Invalid memory or validation fraction")
    return args


def main():
    args = parse_args()
    rank, world, local_rank = (int(os.environ.get(k, v)) for k, v in (("RANK", 0), ("WORLD_SIZE", 1), ("LOCAL_RANK", 0)))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    torch.cuda.set_per_process_memory_fraction(args.memory_fraction, device)
    free, total = torch.cuda.mem_get_info(device)
    if free < total * args.memory_fraction + 2 * 2**30:
        raise RuntimeError(f"GPU {local_rank} has insufficient free VRAM for budget + 2 GiB headroom")
    if world > 1:
        dist.init_process_group("nccl", device_id=device, timeout=timedelta(minutes=10))
    torch.set_num_threads(2)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.benchmark = False
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    info, train, val = load_metadata(args.dataset, args.seed, args.val_fraction)
    cfg = CausalTransitionEncoderConfig(action_dim=7, image_channels=48)
    model = CausalTransitionEncoder(cfg).to(device)
    # Unused cross_norm modules exist in the public architecture. DDP detects
    # those without silently freezing any actually used CTE parameters.
    optimizer = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad),
                                  lr=args.lr, betas=(0.9, 0.95), weight_decay=args.weight_decay)
    dataset = BridgeCTEDataset(args.dataset, info, train)
    val_dataset = BridgeCTEDataset(args.dataset, info, val)
    base_sampler = PairedEpisodeSampler(train, args.batch_size, world, rank, args.seed)
    files = [ROOT / "scripts" / n for n in ("train_bridge_cte.py", "bridge_cte_data.py", "bridge_cte_loss.py")]
    files += [ROOT / "cosmos_framework/model/zeva" / n for n in ("causal_transition_encoder.py", "cte_losses.py")]
    contract = dict(version="bridge_cte_stage1_v1", dataset=str(Path(args.dataset).resolve()),
                    metadata={n: digest_file(Path(args.dataset) / "meta" / n) for n in ("info.json", "episodes.jsonl", "tasks.jsonl")},
                    vae_path=str(Path(args.vae).resolve()), vae_sha256=digest_file(args.vae),
                    code_sha256={str(p.relative_to(ROOT)): digest_file(p) for p in files},
                    action_contract="bridge_q01_q99_xyzrpy_clipped_gripper_open_positive_v1",
                    frame_encoding="single_frame_rgb_native256x256_wan22_bf16_v1", image_key="observation.images.image_0",
                    boundary_stride=4, effect_raw_controls=16, fps=5,
                    full_trajectories=True, cte_initialization="random_no_release_weights",
                    loss="release_effect_v3_mse_diagnostic_only_bridge_masked_global_task_v1",
                    model_config=cfg.to_dict(), loss_config=asdict(CTELossConfig()),
                    seed=args.seed, val_fraction=args.val_fraction, world_size=world,
                    per_rank_batch=args.batch_size, lr=args.lr, warmup_steps=args.warmup_steps,
                    weight_decay=args.weight_decay, sampler="exhaustive_anchors_paired_partner_v1")
    args.output = args.output.resolve()
    if rank == 0:
        args.output.mkdir(parents=True, exist_ok=True)
        if (args.output / "manifest.json").exists() and not args.resume:
            raise FileExistsError("Run already exists; explicitly --resume or choose a new output directory")
        if shutil.disk_usage(args.output).free < args.min_free_gb * 2**30 + 1024**3:
            raise OSError("Not enough disk headroom for a new CTE run")
    if dist.is_initialized():
        dist.barrier()
    step, epoch, cursor = 0, 0, 0
    resume_rng = None
    initialization_hash = model_digest(model) if rank == 0 else None
    if args.resume:
        checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False)
        if checkpoint.get("format") != "bridge_cte_stage1_v1" or checkpoint.get("contract_fingerprint") != fingerprint(contract):
            raise ValueError("Checkpoint/data/code/batch/optimizer contract differs; refusing ambiguous resume")
        model.load_state_dict(checkpoint["model"], strict=True)
        optimizer.load_state_dict(checkpoint["optimizer"])
        step, epoch, cursor = checkpoint["step"], checkpoint["epoch"], checkpoint["batch_cursor"]
        resume_rng = checkpoint["rng_states"][rank]
        del checkpoint
    config = dict(contract=contract, args=vars(args) | {"output": str(args.output)},
                  train_episodes=len(train), validation_episodes=len(val),
                  train_unknown_instruction=sum(e["semantic_id"] < 0 for e in train),
                  steps_per_anchor_epoch=base_sampler.total_batches, global_episode_batch=args.batch_size * world,
                  trainable_parameters=sum(p.numel() for p in model.parameters() if p.requires_grad),
                  initial_random_model_sha256=initialization_hash)
    if rank == 0:
        atomic_json(args.output / "manifest.json", config)
        atomic_json(args.output / "split.json", dict(train=[e["episode_index"] for e in train],
                                                    validation=[e["episode_index"] for e in val]))
        logger = init_wandb(args, config)
    else:
        logger = None
    if dist.is_initialized():
        dist.barrier()
    encoder = FrozenFrameEncoder(args.vae, device, args.vae_microbatch)
    ddp = DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True,
                                   broadcast_buffers=False) if world > 1 else model
    if resume_rng is not None:
        restore_rng(resume_rng)
    print(f"RANK_READY rank={rank} pid={os.getpid()} cuda={local_rank} trainable={config['trainable_parameters']} "
          f"resume_step={step} max_allocated_GiB={torch.cuda.max_memory_allocated()/2**30:.3f}", flush=True)

    def record(kind, values):
        if rank != 0:
            return
        record = dict(kind=kind, timestamp=datetime.now().astimezone().isoformat(), step=step,
                      anchor_epoch=epoch + cursor / base_sampler.total_batches, **values)
        with (args.output / "metrics.jsonl").open("a") as stream:
            stream.write(json.dumps(record, allow_nan=False) + "\n")
        atomic_json(args.output / "status.json", record)
        if logger is not None:
            logger.log({f"{kind}/{k}": v for k, v in values.items()} | {"global_step": step}, step=step)
        print(json.dumps(record, allow_nan=False), flush=True)

    if not args.resume:
        save_checkpoint(model, optimizer, args, contract, step, epoch, cursor, world, rank)
    if not args.skip_initial_validation:
        record("validation", validate(model, encoder, val_dataset, args, world, rank, device))
    requested_stop = False

    def handle_stop(signum, frame):
        nonlocal requested_stop
        requested_stop = True
        print(f"STOP_REQUESTED rank={rank} signal={signum}; finishing current update and saving", flush=True)

    signal.signal(signal.SIGTERM, handle_stop)
    signal.signal(signal.SIGINT, handle_stop)
    model.train()
    last_log_time, last_log_step = time.perf_counter(), step
    while step < args.max_steps:
        sampler = PairedEpisodeSampler(train, args.batch_size, world, rank, args.seed, epoch, cursor)
        loader = build_loader(dataset, sampler, args, rank)
        for raw in loader:
            optimizer.zero_grad(set_to_none=True)
            inputs = encoder(raw)
            batch = move_batch(raw, device)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                output = ddp(inputs, batch["actions"], batch["valid"], batch["transition_valid"])
            losses = bridge_cte_loss(output, batch["actions"], batch["valid"], batch["semantic_ids"], batch["episode_ids"])
            if not torch.isfinite(losses["total"]):
                raise FloatingPointError(f"Nonfinite loss at step {step + 1}, rank {rank}")
            losses["total"].backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0, error_if_nonfinite=True)
            lr = args.lr * min(1.0, (step + 1) / args.warmup_steps)
            for group in optimizer.param_groups:
                group["lr"] = lr
            optimizer.step()
            model.update_ema_target()
            step, cursor = step + 1, cursor + 1
            if step <= 3 or step % args.log_every == 0:
                values = reduce_metrics(losses | {"grad_norm": grad_norm}, device)
                now = time.perf_counter()
                values.update(lr=lr, seconds_per_step=(now-last_log_time)/(step-last_log_step),
                              gpu_max_allocated_gb=torch.cuda.max_memory_allocated()/2**30,
                              gpu_reserved_gb=torch.cuda.memory_reserved()/2**30)
                record("train", values)
                last_log_time, last_log_step = now, step
            stop_flag = torch.tensor(int(requested_stop), device=device)
            if dist.is_initialized():
                dist.all_reduce(stop_flag, op=dist.ReduceOp.MAX)
            stop = bool(stop_flag.item()) or step >= args.max_steps
            if step % args.save_every == 0 or stop:
                save_checkpoint(model, optimizer, args, contract, step, epoch, cursor, world, rank)
            if step % args.validate_every == 0 and not stop:
                record("validation", validate(model, encoder, val_dataset, args, world, rank, device))
                # Validation time is not reported as optimizer-step throughput.
                last_log_time, last_log_step = time.perf_counter(), step
            if stop:
                record("stopped", dict(checkpoint_step=step, reason="signal" if stop_flag.item() else "max_steps"))
                if logger is not None:
                    logger.finish()
                if dist.is_initialized():
                    dist.destroy_process_group()
                return
        epoch, cursor = epoch + 1, 0
        if rank == 0:
            print(f"ANCHOR_EPOCH_COMPLETE epoch={epoch} train_episodes={len(train)} all_anchors_seen=true", flush=True)


if __name__ == "__main__":
    main()