training_code/train.py

# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
import argparse
import os
import sys
from pathlib import Path
import wandb

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.utils.data import DataLoader, DistributedSampler, Subset
from tqdm import tqdm
from torch.distributed.checkpoint.state_dict import (
    get_model_state_dict,
    get_optimizer_state_dict,
    set_optimizer_state_dict,
    StateDictOptions,
)
from safetensors.torch import save_file, load_file
import json

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from configs import VA_CONFIGS
from distributed.fsdp import shard_model, apply_ac
from distributed.util import (
    _configure_model,
    init_distributed,
    dist_mean,
    dist_max,
    dist_sum
)
from einops import rearrange
from modules.utils import (
    load_transformer,
)
from utils import (
    init_logger, 
    logger, 
    get_mesh_id, 
    sample_timestep_id,
    data_seq_to_patch,
    warmup_constant_lambda,
    FlowMatchScheduler
)

from dataset import MultiLatentLeRobotDataset
import gc


class Trainer:
    def __init__(self, config):
        if config.enable_wandb and config.rank == 0:
            wandb.login(host=os.environ['WANDB_BASE_URL'], key=os.environ['WANDB_API_KEY'])
            self.wandb = wandb
            # Was hardcoded to the literal string 'test_lln' -- every run collided under the
            # same name in the dashboard. run_name (set by run() from --config-name) plus a
            # timestamp keeps runs distinguishable and re-runs of the same config from colliding.
            import time as _time
            run_name = getattr(config, 'run_name', None) or config.get('__name__', 'run')
            run_name = f"{run_name}_{_time.strftime('%Y%m%d_%H%M%S')}"
            self.wandb.init(
                entity=os.environ["WANDB_TEAM_NAME"],
                project=os.getenv("WANDB_PROJECT", "va_robotwin"),
                # dir=log_dir,
                config=config,
                mode="online",
                name=run_name,
            )
            logger.info(f"WandB logging enabled, run name: {run_name}")
        self.step = 0
        self.config = config
        self.device = torch.device(f"cuda:{config.local_rank}")
        self.dtype = config.param_dtype
        self.patch_size = config.patch_size

        # Load models
        logger.info("Loading models...")

        # Load and shard transformer with FSDP
        logger.info("Loading transformer...")

        if hasattr(config, 'resume_from') and config.resume_from:
            transformer_path = os.path.join(config.resume_from, 'transformer')
            if config.rank == 0:
                logger.info(f"Resuming from checkpoint: {transformer_path}")
        else:
            transformer_path = os.path.join(config.wan22_pretrained_model_name_or_path, 'transformer')

        self.transformer = load_transformer(
            transformer_path,
            torch_dtype=torch.float32,
            torch_device='cpu',
            attn_mode="flex"
        )

        logger.info("Setting up activation checkpointing ...")
        apply_ac(self.transformer)

        logger.info("Setting up FSDP...")
        shard_fn = shard_model
        self.transformer = _configure_model(
            model=self.transformer,
            shard_fn=shard_fn,
            param_dtype=self.dtype,
            device=self.device,
            eval_mode=False,
        )
        self.transformer.train()
        self.transformer.requires_grad_(True)

        # Optimizer
        self.optimizer = torch.optim.AdamW(
            [p for p in self.transformer.parameters() if p.requires_grad],
            lr=config.learning_rate,
            betas=(config.beta1, config.beta2),
            eps=1e-8,
            weight_decay=config.weight_decay,
            fused=True,
            foreach=False,
        )

        self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(self.optimizer, 
            lr_lambda=lambda step: warmup_constant_lambda(step, warmup_steps=config.warmup_steps))

        # Setup dataloaders
        logger.info("Setting up datasets...")
        full_dataset = MultiLatentLeRobotDataset(config=config)

        # Held-out validation split. Two mutually exclusive modes, both absent for every config
        # outside lift_new_*/place_cube_* -- so this stays a no-op for libero/robotwin/franka/demo.
        #
        #   config.num_train_episodes -- ordered split: episodes [0, N) train, everything from N
        #     on validates. place_cube_bowl_new needs this because its dataset is a chronological
        #     slice of George's recordings (first 100 train, next 20 val) and the val tail comes
        #     from later teleop sessions on purpose. A random split would put clips from the same
        #     session on both sides and report an optimistically low val loss.
        #   config.val_split -- seeded random fraction (lift_new_*).
        num_train_episodes = getattr(config, 'num_train_episodes', 0)
        val_split = getattr(config, 'val_split', 0.0)
        self.val_loader = None
        if num_train_episodes > 0:
            # episode_index is only unique within one LatentLeRobotDataset, so an ordered split
            # is ambiguous once several are concatenated; these configs load exactly one.
            if len(full_dataset._datasets) != 1:
                raise ValueError(
                    f"num_train_episodes requires a single underlying dataset, got "
                    f"{len(full_dataset._datasets)} -- use val_split instead"
                )
            # new_metas[i] is the sample at dataset index i (LatentLeRobotDataset.__getitem__),
            # so split on its episode_index rather than on i: that stays correct if an episode
            # ever contributes more than one clip.
            ep_of = [m["episode_index"] for m in full_dataset._datasets[0].new_metas]
            train_indices = [i for i, e in enumerate(ep_of) if e < num_train_episodes]
            val_indices = [i for i, e in enumerate(ep_of) if e >= num_train_episodes]
            if not train_indices or not val_indices:
                raise ValueError(
                    f"num_train_episodes={num_train_episodes} gives "
                    f"{len(train_indices)} train / {len(val_indices)} val samples over "
                    f"episodes {min(ep_of)}..{max(ep_of)}"
                )
            train_dataset = Subset(full_dataset, train_indices)
            val_dataset = Subset(full_dataset, val_indices)
            if config.rank == 0:
                logger.info(f"train/val split: {len(train_dataset)} train / {len(val_dataset)} "
                            f"val (ordered, episodes <{num_train_episodes} train)")
        elif val_split > 0:
            n = len(full_dataset)
            n_val = max(1, round(n * val_split))
            g = torch.Generator().manual_seed(getattr(config, 'val_seed', 42))
            perm = torch.randperm(n, generator=g).tolist()
            val_indices, train_indices = perm[:n_val], perm[n_val:]
            train_dataset = Subset(full_dataset, train_indices)
            val_dataset = Subset(full_dataset, val_indices)
            if config.rank == 0:
                logger.info(f"train/val split: {len(train_dataset)} train / {len(val_dataset)} "
                            f"val (val_split={val_split}, seed={getattr(config, 'val_seed', 42)})")
        else:
            train_dataset = full_dataset
            val_dataset = None

        train_sampler = DistributedSampler(
            train_dataset,
            num_replicas=config.world_size,
            rank=config.rank,
            shuffle=True,
            seed=42
        ) if config.world_size > 1 else None
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=config.batch_size,
            shuffle=(train_sampler is None),
            num_workers=config.load_worker,
            sampler=train_sampler,
        )

        if val_dataset is not None:
            # shuffle=False: deterministic order is preferable for validation. Every rank must
            # still see the same NUMBER of batches -- FSDP forward is a collective op, so a rank
            # that runs out of batches early would leave the others hanging on a collective that
            # never comes. DistributedSampler guarantees equal counts (via padding) same as it
            # does for train_sampler above.
            val_sampler = DistributedSampler(
                val_dataset,
                num_replicas=config.world_size,
                rank=config.rank,
                shuffle=False,
                seed=42,
            ) if config.world_size > 1 else None
            self.val_loader = DataLoader(
                val_dataset,
                batch_size=config.batch_size,
                shuffle=False,
                num_workers=min(config.load_worker, 4),
                sampler=val_sampler,
            )
        self.val_interval = getattr(config, 'val_interval', config.save_interval)

        self.train_scheduler_latent = FlowMatchScheduler(shift=self.config.snr_shift, sigma_min=0.0, extra_one_step=True)
        self.train_scheduler_latent.set_timesteps(1000, training=True)
        self.train_scheduler_action = FlowMatchScheduler(shift=self.config.action_snr_shift, sigma_min=0.0, extra_one_step=True)
        self.train_scheduler_action.set_timesteps(1000, training=True)

        # Per-run subdirectory so runs NEVER overwrite each other's checkpoints. Previously
        # every run wrote to save_root/checkpoints/checkpoint_step_N, keyed only by step number,
        # so a second run (e.g. place_cube_bowl_base) silently clobbered the first
        # (lift_new_base) at each shared step (200/400/...). The tag must be identical across all
        # ranks AND unique per launch: SLURM_JOB_ID (this repo's Slurm path) and
        # TORCHELASTIC_RUN_ID (bare torchrun) are both shared across a launch's ranks; the
        # timestamp fallback is only reached single-process (world_size 1), where ranks can't
        # disagree. run_name is the config name, so dirs read e.g. place_cube_bowl_base_job542/.
        run_name = getattr(config, 'run_name', None) or config.get('__name__', 'run')
        _slurm_id = os.environ.get('SLURM_JOB_ID')
        if _slurm_id:
            run_subdir = f"{run_name}_job{_slurm_id}"
        elif os.environ.get('TORCHELASTIC_RUN_ID'):
            run_subdir = f"{run_name}_{os.environ['TORCHELASTIC_RUN_ID']}"
        else:
            import time as _time
            run_subdir = f"{run_name}_{_time.strftime('%Y%m%d_%H%M%S')}"
        self.save_dir = Path(config.save_root) / run_subdir / "checkpoints"
        self.save_dir.mkdir(parents=True, exist_ok=True)
        if config.rank == 0:
            logger.info(f"Checkpoints for this run -> {self.save_dir}")

        self.gradient_accumulation_steps = getattr(config, 'gradient_accumulation_steps', 1)
        self.train_loader_iter = None
        # if hasattr(config, 'resume_from') and config.resume_from:
        #     self._load_training_state(config.resume_from)
    
    def _get_next_batch(self):
        """Get next batch from iterator, reset if epoch is finished."""
        if self.train_loader_iter is None:
            self.train_loader_iter = iter(self.train_loader)
        
        try:
            batch = next(self.train_loader_iter)
        except StopIteration:
            # Reset sampler and iterator when epoch finishes
            if hasattr(self.train_loader.sampler, 'set_epoch'):
                self.train_loader.sampler.set_epoch(self.train_loader.sampler.epoch + 1)
            self.train_loader_iter = iter(self.train_loader)
            batch = next(self.train_loader_iter)
        
        return batch

    @torch.no_grad()
    def _add_noise(self, latent, train_scheduler, action_mask=False, action_mode=False, noisy_cond_prob=0.):
        B, C, F, H, W = latent.shape

        timestep_ids = sample_timestep_id(batch_size=F, num_train_timesteps=train_scheduler.num_train_timesteps)
        noise = torch.zeros_like(latent).normal_()
        timesteps = train_scheduler.timesteps[timestep_ids].to(device=self.device)
        noisy_latents =train_scheduler.add_noise(latent, noise, timesteps, t_dim=2)
        targets =train_scheduler.training_target(latent, noise, timesteps)

        patch_f, patch_h, patch_w = self.patch_size
        if action_mode:
            patch_f = patch_h = patch_w = 1
        
        latent_grid_id = get_mesh_id(
            latent.shape[-3] // patch_f,  # F
            latent.shape[-2] // patch_h,  # H
            latent.shape[-1] // patch_w,  # W
            t=1 if action_mode else 0,  # 1 for action mode (0 for latent), not used
            f_w=1,
            f_shift=0,
            action=action_mode
        ).to(self.device)  # shape: [4, seq_len]
        latent_grid_id = latent_grid_id[None].repeat(B, 1, 1)

        if torch.rand(1).item() < noisy_cond_prob:
            cond_timestep_ids = sample_timestep_id(
                    batch_size=F,
                    min_timestep_bd=0.5, 
                    max_timestep_bd=1.0, 
                    num_train_timesteps=train_scheduler.num_train_timesteps,
                )
            noise = torch.zeros_like(latent).normal_()
            cond_timesteps = train_scheduler.timesteps[cond_timestep_ids].to(device=self.device)
            latent = train_scheduler.add_noise(latent, noise, cond_timesteps, t_dim=2)
        else:
            cond_timesteps = torch.zeros_like(timesteps)

        if action_mask is not None:
            noisy_latents *= action_mask.float()
            targets *= action_mask.float()
            latent *= action_mask.float()

        return dict(
            timesteps=timesteps[None].repeat(B, 1),
            noisy_latents=noisy_latents,
            targets=targets,
            latent=latent,
            cond_timesteps=cond_timesteps[None].repeat(B, 1),
            grid_id=latent_grid_id,
        )

    @torch.no_grad()
    def _prepare_input_dict(self, batch_dict):
        """Prepare input dict following infer code pattern from wan_va_server.py."""
        # Generate grid_id following infer code (no batch dimension yet)
        # For action mode: get_mesh_id(shape[-3], shape[-2], shape[-1], t=1, f_w=1, f_shift, action=True)
        latent_dict = self._add_noise(
            latent=batch_dict['latents'], 
            train_scheduler=self.train_scheduler_latent, 
            action_mask=None, 
            action_mode=False,
            noisy_cond_prob=0.5)
        
        action_dict = self._add_noise(
            latent=batch_dict['actions'], 
            train_scheduler=self.train_scheduler_action, 
            action_mask=batch_dict['actions_mask'], 
            action_mode=True,
            noisy_cond_prob=0.0)

        latent_dict['text_emb'] = batch_dict['text_emb']
        action_dict['text_emb'] = batch_dict['text_emb']
        action_dict['actions_mask'] = batch_dict['actions_mask']

        input_dict = {
            'latent_dict': latent_dict,
            'action_dict': action_dict,
            'chunk_size': torch.randint(1, 5, (1,)).item(),
            'window_size': torch.randint(4, 65, (1,)).item(),
        }
        return input_dict

    def convert_input_format(self, input_dict):
        """Convert input dict to match transformer input format if needed."""
        for key, value in input_dict.items():
            input_dict[key] = value.to(self.device)#.to(self.dtype)
        return input_dict

    def compute_loss(self,
        input_dict,
        pred
    ):
        latent_pred, action_pred = pred
        action_pred = rearrange(action_pred, 'b (f n) c -> b c f n 1', f=input_dict['action_dict']['targets'].shape[-3])
        latent_pred = data_seq_to_patch(
                        self.patch_size, latent_pred,
                        input_dict['latent_dict']['targets'].shape[-3], input_dict['latent_dict']['targets'].shape[-2],
                        input_dict['latent_dict']['targets'].shape[-1], batch_size=latent_pred.shape[0])
        Bn, Fn = input_dict['latent_dict']['timesteps'].shape
        latent_loss_weight = self.train_scheduler_latent.training_weight(input_dict['latent_dict']['timesteps'].flatten()).reshape(Bn, Fn)
        action_loss_weight = self.train_scheduler_action.training_weight(input_dict['action_dict']['timesteps'].flatten()).reshape(Bn, Fn)

        # Frame-wise video loss calculation
        latent_loss = F.mse_loss(latent_pred.float(), input_dict['latent_dict']['targets'].float().detach(), reduction='none')
        latent_loss = latent_loss * latent_loss_weight[:, None, :, None, None]
        # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C)
        latent_loss = latent_loss.permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
        latent_loss = latent_loss.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
        # Sum per frame and compute mask per frame
        latent_loss_per_frame = latent_loss.sum(dim=1)  # (B*F,)
        latent_mask_per_frame = torch.ones_like(latent_loss).sum(dim=1)  # (B*F,)
        latent_loss = (latent_loss_per_frame / (latent_mask_per_frame + 1e-6)).mean()

        # Frame-wise action loss calculation
        action_loss = F.mse_loss(action_pred.float(), input_dict['action_dict']['targets'].float().detach(), reduction='none')
        action_loss = action_loss * action_loss_weight[:, None, :, None, None]
        action_loss = action_loss * input_dict['action_dict']['actions_mask'].float()
        # Permute to (B, F, H, W, C) and flatten to (B*F, H*W*C)
        action_loss = action_loss.permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
        action_mask = input_dict['action_dict']['actions_mask'].float().permute(0, 2, 3, 4, 1)  # (B, C, F, H, W) -> (B, F, H, W, C)
        action_loss = action_loss.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
        action_mask = action_mask.flatten(0, 1).flatten(1)  # (B, F, H, W, C) -> (B*F, H*W*C)
        # Sum per frame and normalize by mask per frame
        action_loss_per_frame = action_loss.sum(dim=1)  # (B*F,)
        action_mask_per_frame = action_mask.sum(dim=1)  # (B*F,)
        action_loss = (action_loss_per_frame / (action_mask_per_frame + 1e-6)).mean()

        return latent_loss / self.gradient_accumulation_steps, action_loss / self.gradient_accumulation_steps

    @torch.no_grad()
    def _action_l2_metrics(self, pred, input_dict):
        """De-normalized L2 error between a one-step flow-match reconstruction of the clean
        action and ground truth, restricted to the real (non-padded) action channels.

        One-step reconstruction: FlowMatchScheduler defines noisy = (1-sigma)*clean + sigma*noise
        and target = noise - clean (see wan_va/utils/scheduler.py), so a model that predicts the
        true target exactly satisfies clean = noisy - sigma*target -- same identity
        FlowMatchScheduler.step(..., to_final=True) uses, just applied here with per-frame sigmas
        (scheduler.step() itself only handles a single scalar timestep, not our per-frame ones,
        so this reimplements the same broadcast pattern _add_noise/scheduler.add_noise use).
        This is an approximation (real inference denoises over many steps, not one) -- cheap
        enough to log every training step, not a substitute for real rollout eval.

        Fully generic / safe for every OTHER config in this repo: returns per-real-channel sums
        only; gripper-accuracy and named/grouped breakdowns are opt-in via config fields
        (gripper_channel_index, action_channel_groups / action_channel_names) that only
        lift_new_configs/va_lift_new_cfg.py sets. Absent those fields, this still returns valid
        per-channel-index sums, just without the semantic labels.
        """
        _, action_pred = pred
        action_dict = input_dict['action_dict']
        F_dim = action_dict['targets'].shape[-3]
        action_pred = rearrange(action_pred, 'b (f n) c -> b c f n 1', f=F_dim)

        noisy = action_dict['noisy_latents']
        timesteps = action_dict['timesteps'][0]  # [F] -- identical across the batch dim (see _add_noise)
        timestep_id = torch.argmin(
            (self.train_scheduler_action.timesteps[:, None].to(timesteps.device) - timesteps[None]).abs(),
            dim=0,
        )
        sigma = self.train_scheduler_action.sigmas.to(timesteps.device)[timestep_id]  # [F]
        shape = [1] * noisy.ndim
        shape[2] = sigma.shape[0]  # t_dim=2, matches _add_noise's convention
        sigma = sigma.view(shape)

        x0_pred = noisy - sigma * action_pred.float()
        x0_true = action_dict['latent']
        mask = action_dict['actions_mask'].float()

        n_real = len(self.config.used_action_channel_ids)
        q01 = torch.tensor(self.config.norm_stat['q01'][:n_real], device=noisy.device, dtype=torch.float32)
        q99 = torch.tensor(self.config.norm_stat['q99'][:n_real], device=noisy.device, dtype=torch.float32)
        denorm_shape = [1, n_real] + [1] * (noisy.ndim - 2)

        def denorm(x):
            return (x[:, :n_real].float() + 1) / 2 * (q99 - q01 + 1e-6).view(denorm_shape) + q01.view(denorm_shape)

        pred_real = denorm(x0_pred)
        true_real = denorm(x0_true)
        m = mask[:, :n_real]

        sq_err = (pred_real - true_real) ** 2 * m
        sum_dims = (0, 2, 3, 4)
        out = {
            'sq_err_sum': sq_err.sum(dim=sum_dims).detach(),   # [n_real]
            'mask_sum': (m.sum(dim=sum_dims) + 1e-6).detach(),  # [n_real]
        }

        gripper_idx = getattr(self.config, 'gripper_channel_index', None)
        if gripper_idx is not None:
            gripper_pred = (pred_real[:, gripper_idx:gripper_idx + 1] > 0.5).float()
            gripper_true = (true_real[:, gripper_idx:gripper_idx + 1] > 0.5).float()
            gripper_mask = m[:, gripper_idx:gripper_idx + 1]
            out['gripper_match_sum'] = ((gripper_pred == gripper_true).float() * gripper_mask).sum().detach()
            out['gripper_mask_sum'] = (gripper_mask.sum() + 1e-6).detach()
        return out

    def _rmse_report(self, sq_err_sum, mask_sum, gripper_match_sum=None, gripper_mask_sum=None):
        """Turn accumulated (already cross-rank-summed) sq_err_sum/mask_sum into final scalar
        metrics. Per-group/per-name breakdowns are opt-in via config.action_channel_groups /
        config.action_channel_names -- absent those, only 'overall' (and gripper_accuracy, if
        gripper_match_sum is given) is reported."""
        sq_err_sum = sq_err_sum.detach().cpu()
        mask_sum = mask_sum.detach().cpu()
        report = {'overall': (sq_err_sum.sum() / (mask_sum.sum() + 1e-6)).sqrt().item()}

        groups = getattr(self.config, 'action_channel_groups', None)
        names = getattr(self.config, 'action_channel_names', None)
        if groups:
            for name, idxs in groups.items():
                idxs_t = torch.tensor(idxs)
                report[f'rmse_{name}'] = (sq_err_sum[idxs_t].sum() / (mask_sum[idxs_t].sum() + 1e-6)).sqrt().item()
        elif names:
            for i, name in enumerate(names):
                report[f'rmse_{name}'] = (sq_err_sum[i] / (mask_sum[i] + 1e-6)).sqrt().item()

        if gripper_match_sum is not None:
            report['gripper_accuracy'] = (
                gripper_match_sum.detach().cpu() / (gripper_mask_sum.detach().cpu() + 1e-6)
            ).item()
        return report

    @torch.no_grad()
    def validate(self):
        """Run the held-out validation split (config.val_split). All ranks must call this
        together -- FSDP forward is a collective op, so skipping it on some ranks (e.g. "only
        rank 0 validates") would hang the others. Only rank 0 needs the returned report for
        logging, but every rank has to process its shard of val_loader for that report to be
        correct (and to avoid a hang)."""
        if self.val_loader is None:
            return {}
        self.transformer.eval()

        n_real = len(self.config.used_action_channel_ids)
        sq_err_sum = torch.zeros(n_real, device=self.device)
        mask_sum = torch.zeros(n_real, device=self.device)
        gripper_match_sum = torch.zeros((), device=self.device)
        gripper_mask_sum = torch.zeros((), device=self.device)
        has_gripper = getattr(self.config, 'gripper_channel_index', None) is not None
        latent_losses, action_losses = [], []

        for batch in self.val_loader:
            batch = self.convert_input_format(batch)
            input_dict = self._prepare_input_dict(batch)
            output = self.transformer(input_dict, train_mode=True)
            latent_loss, action_loss = self.compute_loss(input_dict, output)
            # compute_loss divides by gradient_accumulation_steps for training's backward
            # scaling; undo that here so validation loss is directly comparable to a fresh
            # per-batch loss, not scaled by an accumulation window that doesn't apply here.
            latent_losses.append(latent_loss.detach() * self.gradient_accumulation_steps)
            action_losses.append(action_loss.detach() * self.gradient_accumulation_steps)

            l2 = self._action_l2_metrics(output, input_dict)
            sq_err_sum += l2['sq_err_sum']
            mask_sum += l2['mask_sum']
            if has_gripper:
                gripper_match_sum += l2['gripper_match_sum']
                gripper_mask_sum += l2['gripper_mask_sum']

        self.transformer.train()

        n_batches = max(1, len(self.val_loader))
        latent_loss_show = dist_mean(torch.stack(latent_losses).sum() / n_batches)
        action_loss_show = dist_mean(torch.stack(action_losses).sum() / n_batches)

        sq_err_sum = dist_sum(sq_err_sum)
        mask_sum = dist_sum(mask_sum)
        if has_gripper:
            gripper_match_sum = dist_sum(gripper_match_sum)
            gripper_mask_sum = dist_sum(gripper_mask_sum)

        report = self._rmse_report(
            sq_err_sum, mask_sum,
            gripper_match_sum if has_gripper else None,
            gripper_mask_sum if has_gripper else None,
        )
        report['latent_loss'] = latent_loss_show.item()
        report['action_loss'] = action_loss_show.item()
        return report

    def _train_step(self, batch, batch_idx):
        """Train a single batch, returns losses for logging."""
        batch = self.convert_input_format(batch)
        input_dict = self._prepare_input_dict(batch)
        
        should_sync = (batch_idx + 1) % self.gradient_accumulation_steps == 0
        
        if not should_sync:
            self.transformer.set_requires_gradient_sync(False)
        else:
            self.transformer.set_requires_gradient_sync(True)

        output = self.transformer(input_dict, train_mode=True)
        latent_loss, action_loss = self.compute_loss(input_dict, output)
        loss = latent_loss + action_loss

        loss.backward()

        # Reuses the already-computed forward output -- no extra forward pass needed to also
        # report a physically-interpretable action error alongside the flow-matching loss.
        l2 = self._action_l2_metrics(output, input_dict)
        losses = {
            'latent_loss': latent_loss.detach(),
            'action_loss': action_loss.detach(),
            'sq_err_sum': l2['sq_err_sum'],
            'mask_sum': l2['mask_sum'],
        }
        if 'gripper_match_sum' in l2:
            losses['gripper_match_sum'] = l2['gripper_match_sum']
            losses['gripper_mask_sum'] = l2['gripper_mask_sum']

        # Only update weights after accumulating gradients
        if should_sync:
            total_norm = torch.nn.utils.clip_grad_norm_(self.transformer.parameters(), 2.0)
            self.optimizer.step()
            self.lr_scheduler.step()
            self.optimizer.zero_grad()
            
            losses['total_norm'] = total_norm
            losses['should_log'] = True
        else:
            losses['should_log'] = False

        return losses

    def save_checkpoint(self,):
        """Save model checkpoint in the same format as pretrained model."""
        try:
            state_dict = get_model_state_dict(
                self.transformer,
                options=StateDictOptions(full_state_dict=True, cpu_offload=True),
            )
            state_dict_bf16 = {k: v.to(torch.bfloat16) for k, v in state_dict.items()}
            # optim_state = get_optimizer_state_dict(
            #         self.transformer, self.optimizer,
            #         options=StateDictOptions(full_state_dict=True, cpu_offload=True),
            #     )

            # Only rank 0 saves the checkpoint
            if self.config.rank == 0:
                checkpoint_dir = self.save_dir / f"checkpoint_step_{self.step}"
                checkpoint_dir.mkdir(parents=True, exist_ok=True)

                # Save transformer in the same format as pretrained model
                transformer_dir = checkpoint_dir / "transformer"
                transformer_dir.mkdir(parents=True, exist_ok=True)

                logger.info(f"Saving transformer to {transformer_dir}")

                # Manually save in diffusers format (outside FSDP context to avoid deadlock)
                # Save model weights
                model_file = transformer_dir / "diffusion_pytorch_model.safetensors"
                save_file(state_dict_bf16, model_file)

                # Save config (copy from original transformer config and update _name_or_path)
                config_file = transformer_dir / "config.json"
                config_dict = dict(self.transformer.config)
                config_dict.pop('_name_or_path', None)
                with open(config_file, 'w') as f:
                    json.dump(config_dict, f, indent=2)

                # # Save optimizer state and training metadata in PyTorch format
                # training_state_path = checkpoint_dir / "training_state.pt"
                # logger.info(f"Saving training state to {training_state_path}")
                # torch.save({
                #     'step': self.step,
                #     'optimizer_state_dict': optim_state,
                #     'config': vars(self.config),
                # }, training_state_path)

                logger.info(f"Checkpoint saved successfully at step {self.step}")

            # Synchronize all processes after saving
            if dist.is_initialized():
                dist.barrier()

        except Exception as e:
            if self.config.rank == 0:
                logger.error(f"Failed to save checkpoint: {e}")
                import traceback
                logger.error(traceback.format_exc())
            # Ensure all processes stay synchronized even on error
            if dist.is_initialized():
                dist.barrier()

    def _load_training_state(self, checkpoint_path):
        """Load training state (optimizer + step) after FSDP and optimizer creation."""
        checkpoint_dir = Path(checkpoint_path)
        training_state_path = checkpoint_dir / "training_state.pt"

        if not training_state_path.exists():
            if self.config.rank == 0:
                logger.warning(f"Training state not found: {training_state_path}, starting from step 0")
            return

        if self.config.rank == 0:
            logger.info(f"Loading training state from {training_state_path}")

        # All ranks load the training state directly
        training_state = torch.load(training_state_path, map_location='cpu', weights_only=False)

        # All ranks load optimizer state (required for FSDP)
        set_optimizer_state_dict(
            self.transformer, self.optimizer,
            optim_state_dict=training_state['optimizer_state_dict'],
            options=StateDictOptions(full_state_dict=True, strict=False)
        )
        self.step = training_state.get('step', 0)

        if self.config.rank == 0:
            logger.info(f"Training state loaded, resuming from step {self.step}")

        # Synchronize all ranks
        if dist.is_initialized():
            dist.barrier()

    def train(self):
        """Main training loop - train by steps instead of epochs."""
        logger.info(f"Starting training for {self.config.num_steps} steps...")
        self.transformer.train()

        progress_bar = tqdm(
            total=self.config.num_steps,
            desc="Training",
            disable=(self.config.rank != 0),
            leave=True,
            dynamic_ncols=True,
            initial=self.step
        )

        self.optimizer.zero_grad()
        accumulated_latent_losses = []
        accumulated_action_losses = []
        accumulated_sq_err_sum = None
        accumulated_mask_sum = None
        accumulated_gripper_match_sum = None
        accumulated_gripper_mask_sum = None
        step_in_accumulation = 0

        while self.step < self.config.num_steps:
            # Get next batch (handles epoch reset automatically)
            batch = self._get_next_batch()

            losses = self._train_step(batch, step_in_accumulation)

            # Accumulate losses for logging
            accumulated_latent_losses.append(losses['latent_loss'])
            accumulated_action_losses.append(losses['action_loss'])
            accumulated_sq_err_sum = (losses['sq_err_sum'] if accumulated_sq_err_sum is None
                                       else accumulated_sq_err_sum + losses['sq_err_sum'])
            accumulated_mask_sum = (losses['mask_sum'] if accumulated_mask_sum is None
                                     else accumulated_mask_sum + losses['mask_sum'])
            if 'gripper_match_sum' in losses:
                accumulated_gripper_match_sum = (losses['gripper_match_sum'] if accumulated_gripper_match_sum is None
                                                  else accumulated_gripper_match_sum + losses['gripper_match_sum'])
                accumulated_gripper_mask_sum = (losses['gripper_mask_sum'] if accumulated_gripper_mask_sum is None
                                                 else accumulated_gripper_mask_sum + losses['gripper_mask_sum'])
            step_in_accumulation += 1

            # Log and checkpoint when optimizer steps
            if losses['should_log']:
                lr = self.lr_scheduler.get_last_lr()[0]

                # Average accumulated losses
                latent_loss_show = dist_mean(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item()
                action_loss_show = dist_mean(torch.stack(accumulated_action_losses).sum()).detach().cpu().item()
                max_latent_loss_show = dist_max(torch.stack(accumulated_latent_losses).sum()).detach().cpu().item()
                max_action_loss_show = dist_max(torch.stack(accumulated_action_losses).sum()).detach().cpu().item()

                train_action_report = self._rmse_report(
                    dist_sum(accumulated_sq_err_sum),
                    dist_sum(accumulated_mask_sum),
                    dist_sum(accumulated_gripper_match_sum) if accumulated_gripper_match_sum is not None else None,
                    dist_sum(accumulated_gripper_mask_sum) if accumulated_gripper_mask_sum is not None else None,
                )

                # Clear accumulated losses
                accumulated_latent_losses = []
                accumulated_action_losses = []
                accumulated_sq_err_sum = None
                accumulated_mask_sum = None
                accumulated_gripper_match_sum = None
                accumulated_gripper_mask_sum = None
                step_in_accumulation = 0

                torch.cuda.synchronize()
                if self.step % self.config.gc_interval == 0:
                    torch.cuda.empty_cache()
                    gc.collect()

                if self.config.rank == 0:
                    total_norm = losses['total_norm']
                    progress_bar.n += 1
                    progress_bar.set_postfix({
                        'latent_loss': f'{latent_loss_show:.4f}',
                        'action_loss': f'{action_loss_show:.4f}',
                        'act_rmse': f"{train_action_report['overall']:.4f}",
                        'step': self.step,
                        'grad_norm': f'{total_norm.item():.2f}',
                        'lr': f'{lr:.2e}'
                    })
                    logger.info(f"step {self.step} train_action_report: {train_action_report}")
                    if self.config.enable_wandb:
                        wandb_log = {
                            'loss_metrics/global_avg_video_loss': latent_loss_show,
                            'loss_metrics/global_avg_action_loss': action_loss_show,
                            'loss_metrics/global_max_video_loss': max_latent_loss_show,
                            'loss_metrics/global_max_action_loss': max_action_loss_show,
                            'grad_norm': total_norm.item(),
                            'lr': lr,
                        }
                        wandb_log.update({f'train_action/{k}': v for k, v in train_action_report.items()})
                        self.wandb.log(wandb_log, step=self.step)

                self.step += 1

                if self.step % self.config.save_interval == 0:
                    if self.config.rank == 0:
                        logger.info(f"Starting save model at step {self.step}")
                    self.save_checkpoint()

                if self.val_loader is not None and self.step % self.val_interval == 0:
                    if self.config.rank == 0:
                        logger.info(f"Running validation at step {self.step}")
                    val_report = self.validate()
                    if self.config.rank == 0:
                        logger.info(f"step {self.step} validation: {val_report}")
                        if self.config.enable_wandb:
                            self.wandb.log({f'val/{k}': v for k, v in val_report.items()}, step=self.step)

            if dist.is_initialized():
                dist.barrier()

        progress_bar.close()
        logger.info("Training completed!")


def run(args):
    """Main entry point."""
    config = VA_CONFIGS[args.config_name]
    config.run_name = args.config_name

    rank = int(os.getenv("RANK", 0))
    local_rank = int(os.environ.get('LOCAL_RANK', 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))

    init_distributed(world_size, local_rank, rank)

    config.rank = rank
    config.local_rank = local_rank
    config.world_size = world_size

    if args.save_root is not None:
        config.save_root = args.save_root

    if rank == 0:
        logger.info(f"Using config: {args.config_name}")
        logger.info(f"World size: {world_size}, Local rank: {local_rank}")

    trainer = Trainer(config)
    trainer.train()


def main():
    """Parse arguments and run training."""
    parser = argparse.ArgumentParser(description="Train WAN model for robotics")
    parser.add_argument(
        "--config-name",
        type=str,
        default='robotwin_train',
        help="Config name",
    )
    parser.add_argument(
        "--save-root",
        type=str,
        default=None,
        help="Root directory for saving checkpoints",
    )

    args = parser.parse_args()
    run(args)


if __name__ == "__main__":
    init_logger()
    main()