wan_va/train.py

# Copyright 2024-2025 The Robbyant Team Authors. All rights reserved.
import argparse
import os
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
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

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
)
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, dataset_indexes_ready
from .mcp import shift_latents_for_mcp, validate_mcp_settings
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
            self.wandb.init(
                entity=os.environ["WANDB_TEAM_NAME"],
                project=os.getenv("WANDB_PROJECT", "va_robotwin"),
                # dir=log_dir,
                config=config,
                mode="online",
                name='test_lln'
                # name=os.path.basename(os.path.normpath(job_config.job.dump_folder))
            )
            logger.info("WandB logging enabled")
        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
        self.enable_mcp = getattr(config, 'enable_mcp', True)

        if self.enable_mcp:
            validate_mcp_settings(
                num_mcp_depths=config.num_mcp_depths,
                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
                mcp_loss_weights=config.mcp_loss_weights,
            )

        # 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",
            disable_mcp=not self.enable_mcp,
        )

        if self.enable_mcp:
            validate_mcp_settings(
                num_mcp_depths=config.num_mcp_depths,
                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
                mcp_loss_weights=config.mcp_loss_weights,
                num_layers=len(self.transformer.blocks),
            )
            initialized = self.transformer.enable_mcp_training(
                num_mcp_depths=config.num_mcp_depths,
                mcp_blocks_per_depth=config.mcp_blocks_per_depth,
                mcp_hidden_collect_layers=config.mcp_hidden_collect_layers,
                init_from_backbone=config.mcp_init_from_backbone,
            )
            if initialized:
                total_mcp_blocks = (
                    config.num_mcp_depths * config.mcp_blocks_per_depth)
                if config.mcp_init_from_backbone:
                    num_backbone_blocks = len(self.transformer.blocks)
                    source_start = (
                        num_backbone_blocks - config.mcp_blocks_per_depth)
                    logger.info(
                        f"Initializing {config.num_mcp_depths} MCP depths "
                        f"({config.mcp_blocks_per_depth} blocks per depth) "
                        f"from backbone blocks[{source_start}:"
                        f"{num_backbone_blocks}]"
                    )
                    for depth in range(config.num_mcp_depths):
                        logger.info(
                            f"  MCP depth {depth + 1}: initialized from "
                            f"backbone blocks[{source_start}:"
                            f"{num_backbone_blocks}]"
                        )
                    logger.info(
                        f"Initialized {config.num_mcp_depths} x "
                        f"{config.mcp_blocks_per_depth} = "
                        f"{total_mcp_blocks} MCP blocks"
                    )
                else:
                    logger.info(
                        f"Initialized {total_mcp_blocks} MCP blocks from "
                        "scratch"
                    )
                logger.info(
                    "MCP hidden-state collection uses backbone indices "
                    f"{list(config.mcp_hidden_collect_layers)}"
                )
                logger.info(
                    "Initialized MCP hidden fuser and input projections "
                    "from scratch"
                )
            else:
                logger.info(
                    f"Loaded MCP modules from checkpoint: "
                    f"{config.num_mcp_depths} depths x "
                    f"{config.mcp_blocks_per_depth} blocks per depth"
                )

        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...")
        cache_ready = False
        if (
            config.world_size > 1
            and getattr(config, 'enable_dataset_index_cache', True)
        ):
            cache_ready_flag = torch.zeros(1, device=self.device, dtype=torch.int)
            if config.rank == 0:
                cache_ready_flag.fill_(int(dataset_indexes_ready(config)))
            dist.broadcast(cache_ready_flag, src=0)
            cache_ready = bool(cache_ready_flag.item())

        use_rank_zero_indexing = (
            config.world_size > 1
            and getattr(config, 'enable_dataset_index_cache', True)
            and not cache_ready
        )
        if config.rank == 0 and cache_ready:
            logger.info(
                "Dataset index and Arrow caches are complete; "
                "loading all ranks concurrently"
            )
        if use_rank_zero_indexing and config.rank != 0:
            dist.barrier()
            # Rank 0 has completed any requested rebuild at this point.
            config.rebuild_dataset_index_cache = False

        train_dataset = MultiLatentLeRobotDataset(config=config)

        if use_rank_zero_indexing:
            if config.rank == 0:
                dist.barrier()
            dist.barrier()

        if config.rank == 0:
            logger.info(
                "Dataset ready: %d samples from %d datasets "
                "(%d index cache hits, %d direct Arrow loads, %d rebuilt)",
                len(train_dataset),
                len(train_dataset._datasets),
                train_dataset.index_cache_hits,
                train_dataset.hf_cache_hits,
                train_dataset.index_cache_misses,
            )
        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,
        )

        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)
        self.train_scheduler_mcp = None
        if self.enable_mcp:
            self.train_scheduler_mcp = FlowMatchScheduler(
                shift=self.config.mcp_snr_shift,
                sigma_min=0.0,
                extra_one_step=True,
            )
            self.train_scheduler_mcp.set_timesteps(1000, training=True)

        self.save_dir = Path(config.save_root) / "checkpoints"
        self.save_dir.mkdir(parents=True, exist_ok=True)

        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., frame_shift=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=frame_shift,
            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."""
        chunk_size = torch.randint(1, 5, (1,)).item()
        # 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': chunk_size,
            'window_size': torch.randint(4, 65, (1,)).item(),
        }
        if self.enable_mcp:
            mcp_latent_dicts = []
            for depth in range(self.config.num_mcp_depths):
                frame_shift = (depth + 1) * chunk_size
                shifted_latents, valid_mask = shift_latents_for_mcp(
                    batch_dict['latents'], frame_shift)
                mcp_latent_dict = self._add_noise(
                    latent=shifted_latents,
                    train_scheduler=self.train_scheduler_mcp,
                    action_mask=None,
                    action_mode=False,
                    noisy_cond_prob=0.0,
                    frame_shift=frame_shift,
                )
                mcp_latent_dict.pop('latent')
                mcp_latent_dict.pop('cond_timesteps')
                mcp_latent_dict['valid_mask'] = valid_mask
                mcp_latent_dicts.append(mcp_latent_dict)
            input_dict['mcp_latent_dicts'] = mcp_latent_dicts
        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
    ):
        if self.enable_mcp:
            latent_pred, action_pred, mcp_pred_list = pred
        else:
            latent_pred, action_pred = pred
            mcp_pred_list = []
        if self.enable_mcp and len(mcp_pred_list) != self.config.num_mcp_depths:
            raise RuntimeError(
                "MCP output depth count must match num_mcp_depths")
        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()

        mcp_losses = []
        for mcp_pred, mcp_latent_dict in zip(
                mcp_pred_list, input_dict.get('mcp_latent_dicts', [])):
            mcp_pred = data_seq_to_patch(
                self.patch_size,
                mcp_pred,
                mcp_latent_dict['targets'].shape[-3],
                mcp_latent_dict['targets'].shape[-2],
                mcp_latent_dict['targets'].shape[-1],
                batch_size=mcp_pred.shape[0],
            )
            mcp_batch_size, mcp_num_frames = mcp_latent_dict[
                'timesteps'].shape
            mcp_loss_weight = self.train_scheduler_mcp.training_weight(
                mcp_latent_dict['timesteps'].flatten()).reshape(
                    mcp_batch_size, mcp_num_frames)
            mcp_loss = F.mse_loss(
                mcp_pred.float(),
                mcp_latent_dict['targets'].float().detach(),
                reduction='none',
            )
            mcp_loss = mcp_loss * mcp_loss_weight[:, None, :, None, None]
            valid_mask = mcp_latent_dict['valid_mask'].to(
                device=mcp_loss.device, dtype=mcp_loss.dtype)
            valid_count = valid_mask.expand_as(mcp_loss).sum()
            mcp_loss = (mcp_loss * valid_mask).sum() / valid_count.clamp_min(1.)
            mcp_losses.append(mcp_loss / self.gradient_accumulation_steps)

        return (
            latent_loss / self.gradient_accumulation_steps,
            action_loss / self.gradient_accumulation_steps,
            mcp_losses,
        )

    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, mcp_losses = self.compute_loss(
            input_dict, output)
        mcp_loss = sum(
            weight * depth_loss
            for weight, depth_loss in zip(
                self.config.mcp_loss_weights, mcp_losses)
        ) if mcp_losses else latent_loss.new_zeros(())
        loss = latent_loss + action_loss + mcp_loss

        loss.backward()

        losses = {
            'latent_loss': latent_loss.detach(),
            'action_loss': action_loss.detach(),
            'mcp_losses': [depth_loss.detach() for depth_loss in mcp_losses],
            'mcp_loss': mcp_loss.detach(),
        }
        
        # 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_mcp_losses = [
            [] for _ in range(self.config.num_mcp_depths)
        ] if self.enable_mcp else []
        accumulated_mcp_total_losses = []
        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'])
            for depth, mcp_loss in enumerate(losses['mcp_losses']):
                accumulated_mcp_losses[depth].append(mcp_loss)
            accumulated_mcp_total_losses.append(losses['mcp_loss'])
            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()
                mcp_loss_shows = [
                    dist_mean(torch.stack(depth_losses).sum()).detach().cpu().item()
                    for depth_losses in accumulated_mcp_losses
                ]
                mcp_total_loss_show = dist_mean(
                    torch.stack(accumulated_mcp_total_losses).sum()
                ).detach().cpu().item()

                # Clear accumulated losses
                accumulated_latent_losses = []
                accumulated_action_losses = []
                accumulated_mcp_losses = [
                    [] for _ in range(self.config.num_mcp_depths)
                ] if self.enable_mcp else []
                accumulated_mcp_total_losses = []
                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
                    postfix = {
                        'latent_loss': f'{latent_loss_show:.4f}',
                        'action_loss': f'{action_loss_show:.4f}',
                        'step': self.step,
                        'grad_norm': f'{total_norm.item():.2f}',
                        'lr': f'{lr:.2e}'
                    }
                    if self.enable_mcp:
                        postfix['mcp_loss'] = f'{mcp_total_loss_show:.4f}'
                    progress_bar.set_postfix(postfix)
                    if self.config.enable_wandb:
                        log_values = {
                            '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,
                        }
                        if self.enable_mcp:
                            log_values['loss_metrics/mcp_weighted_total'] = (
                                mcp_total_loss_show)
                            for depth, mcp_loss_show in enumerate(
                                    mcp_loss_shows):
                                log_values[
                                    f'loss_metrics/mcp_depth_{depth + 1}'] = (
                                        mcp_loss_show)
                        self.wandb.log(log_values, 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 dist.is_initialized():
                dist.barrier()

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


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

    overrides = {
        'wan22_pretrained_model_name_or_path': args.pretrained_model_path,
        'dataset_path': args.dataset_path,
        'empty_emb_path': args.empty_emb_path,
        'learning_rate': args.learning_rate,
        'cfg_prob': args.cfg_prob,
        'init_worker': args.init_worker,
        'load_worker': args.load_worker,
        'batch_size': args.batch_size,
        'gradient_accumulation_steps': args.gradient_accumulation_steps,
        'num_steps': args.num_steps,
        'save_interval': args.save_interval,
        'save_root': args.save_root,
    }
    for key, value in overrides.items():
        if value is not None:
            config[key] = value

    if args.dataset_path is not None and args.empty_emb_path is None:
        config.empty_emb_path = os.path.join(args.dataset_path, 'empty_emb.pt')
    if args.disable_wandb:
        config.enable_wandb = False

    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 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",
    )
    parser.add_argument(
        "--pretrained-model-path",
        type=str,
        default=None,
        help="Pretrained model root containing the transformer directory",
    )
    parser.add_argument(
        "--dataset-path",
        type=str,
        default=None,
        help="Root directory containing open-format LeRobot datasets",
    )
    parser.add_argument(
        "--empty-emb-path",
        type=str,
        default=None,
        help="Path to empty_emb.pt (defaults to DATASET_PATH/empty_emb.pt)",
    )
    parser.add_argument(
        "--disable-wandb",
        action="store_true",
        help="Disable Weights & Biases logging",
    )
    parser.add_argument("--learning-rate", type=float, default=None)
    parser.add_argument("--cfg-prob", type=float, default=None)
    parser.add_argument("--init-worker", type=int, default=None)
    parser.add_argument("--load-worker", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument(
        "--gradient-accumulation-steps", type=int, default=None)
    parser.add_argument("--num-steps", type=int, default=None)
    parser.add_argument("--save-interval", type=int, default=None)

    args = parser.parse_args()
    run(args)


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