nano_WaveGen/train_text2wave.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from transformers import LongT5ForConditionalGeneration, T5ForConditionalGeneration, T5Tokenizer
from accelerate import Accelerator
from accelerate.utils import set_seed
from concurrent.futures import ThreadPoolExecutor
import numpy as np
from pathlib import Path
import yaml
from tqdm import tqdm
from typing import Dict, List, Tuple, Optional
import argparse
import os
import re
import warnings
from collections import defaultdict
import time
from datetime import datetime
import sys
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Add parent directory to path to find data and utils modules
SCRIPT_DIR = Path(__file__).resolve().parent
WAVEGEN_ROOT = SCRIPT_DIR.parent
if str(WAVEGEN_ROOT) not in sys.path:
    sys.path.insert(0, str(WAVEGEN_ROOT))

# Suppress the specific transformers warning about past_key_values
warnings.filterwarnings("ignore", message="Passing a tuple of `past_key_values` is deprecated")

from data.movi_dataset import create_dataloader
from utils.save_generation_results import save_generation_results


class Text2WaveModel(nn.Module):
    """Text to Superquadric Wave Parameters Model"""
    
    def __init__(
        self,
        model_name: str = "google/long-t5-tglobal-base",
        max_objects: int = 10,
        num_frames: int = 24,
        max_history_frames: int = 3,
        random_history_sampling: bool = True,
        decoder_noise_std: float = 0.0,
    ):
        super().__init__()
        
        self.max_objects = max_objects
        self.num_frames = num_frames
        self.max_history_frames = max_history_frames
        self.random_history_sampling = random_history_sampling
        self.decoder_noise_std = float(decoder_noise_std)
        # exists(1) + shape(2) + scale(3) + translation(3) + rotation(3) + velocity(3)
        self.object_param_dim = 15
        
        # Load appropriate T5-family model (LongT5 for large checkpoints, vanilla T5 for smaller variants)
        self.model_name = model_name
        self.is_longt5 = "long-t5" in model_name.lower()
        self.tokenizer = T5Tokenizer.from_pretrained(model_name)
        if self.is_longt5:
            self.t5_model = LongT5ForConditionalGeneration.from_pretrained(model_name)
        else:
            self.t5_model = T5ForConditionalGeneration.from_pretrained(model_name)
        
        # Resize model embeddings to match tokenizer if needed
        if self.tokenizer.vocab_size != self.t5_model.config.vocab_size:
            self.t5_model.resize_token_embeddings(self.tokenizer.vocab_size)
        
        # Get T5 hidden size
        self.hidden_size = self.t5_model.config.d_model
        
        # Output projection layers
        # Object parameters: exists + shape[2] + scale[3] + translation[3] + rotation[3] + velocity[3]
        self.object_proj = nn.Linear(self.hidden_size, max_objects * self.object_param_dim)
        
        # World parameters: camera_pos(3) + camera_quat(4) + scene_scale(1) = 8
        self.world_proj = nn.Linear(self.hidden_size, 8)
        
        # Physics parameters: mass(1) + friction(1) + restitution(1) = 3
        self.physics_proj = nn.Linear(self.hidden_size, max_objects * 3)
        
        # Relative time embedding
        self.time_embed = nn.Linear(1, self.hidden_size)

        # History embedding (autoregressive context up to max_history_frames)
        history_feature_dim = max_history_frames * (max_objects * self.object_param_dim + 8) + max_objects * 3
        self.history_feature_dim = history_feature_dim
        self.history_proj = nn.Linear(history_feature_dim, self.hidden_size)

        
        # Initialize weights with small values to prevent NaN
        self._init_weights()
        
    def _init_weights(self):
        """Initialize weights for stability"""
        # Very small initialization for output projections
        for module in [self.object_proj, self.world_proj, self.physics_proj]:
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            nn.init.zeros_(module.bias)
            
        # Time embedding initialization
        nn.init.normal_(self.time_embed.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.time_embed.bias)

        # History embedding initialization
        nn.init.normal_(self.history_proj.weight, mean=0.0, std=0.02)
        nn.init.zeros_(self.history_proj.bias)
        
    def _initialize_history_state(
        self,
        history_frames: Optional[Dict[str, torch.Tensor]],
        batch_size: int,
        device: torch.device,
    ) -> Tuple[List[Dict[str, torch.Tensor]], torch.Tensor]:
        """Prepare history buffer and physics state for autoregressive decoding."""
        history_buffer: List[Dict[str, torch.Tensor]] = []

        physics_state = torch.zeros(
            batch_size,
            self.max_objects,
            3,
            device=device,
            dtype=torch.float32,
        )

        if history_frames is not None:
            objects_hist = history_frames.get('objects')
            world_hist = history_frames.get('world')
            physics_hist = history_frames.get('physics')

            if physics_hist is not None:
                physics_state = physics_hist.to(device=device, dtype=torch.float32)

            if objects_hist is not None and world_hist is not None:
                history_len = objects_hist.shape[1]
                for idx in range(history_len):
                    history_buffer.append({
                        'objects': objects_hist[:, idx, :, :self.object_param_dim].to(device=device, dtype=torch.float32),
                        'world': world_hist[:, idx, :8].to(device=device, dtype=torch.float32),
                    })

        if len(history_buffer) == 0:
            history_buffer.append({
                'objects': torch.zeros(batch_size, self.max_objects, self.object_param_dim, device=device),
                'world': torch.zeros(batch_size, 8, device=device),
            })

        history_buffer = history_buffer[-self.max_history_frames:]

        return history_buffer, physics_state

    def sample_decoder_noise(self, batch_size: int, device: torch.device) -> Optional[torch.Tensor]:
        """Sample decoder noise embedding when noise std > 0."""
        if self.decoder_noise_std <= 0:
            return None
        noise = torch.randn(batch_size, self.hidden_size, device=device)
        return noise * self.decoder_noise_std

    def _build_history_embedding(
        self,
        history_buffer: List[Dict[str, torch.Tensor]],
        physics_state: torch.Tensor,
        use_frames: int,
    ) -> torch.Tensor:
        """Convert most recent history frames into conditioning embedding."""
        batch_size = physics_state.shape[0]
        device = physics_state.device

        frame_dim = self.max_objects * self.object_param_dim + 8
        history_tensor = torch.zeros(
            batch_size,
            self.max_history_frames * frame_dim,
            device=device,
        )

        use_frames = min(use_frames, self.max_history_frames)
        recent_frames = history_buffer[-use_frames:] if use_frames > 0 else []
        for slot, frame in enumerate(recent_frames):
            offset = slot * frame_dim
            obj_flat = frame['objects'].reshape(batch_size, -1)
            world_feat = frame['world']
            history_tensor[:, offset:offset + obj_flat.shape[1]] = obj_flat
            history_tensor[:, offset + obj_flat.shape[1]:offset + frame_dim] = world_feat

        physics_flat = physics_state.reshape(batch_size, -1)
        history_features = torch.cat([history_tensor, physics_flat], dim=-1)
        return self.history_proj(history_features)

    def forward(
        self,
        input_text: List[str],
        target_frames: torch.Tensor,  # [batch, num_target_frames, ...]
        history_frames: Optional[Dict[str, torch.Tensor]] = None,  # History context (objects/world/physics)
        relative_times: torch.Tensor = None,  # [batch, num_target_frames]
        static_object_params: Optional[torch.Tensor] = None,  # Optional static params to enforce (exists+shape+scale)
        noise: Optional[torch.Tensor] = None,  # Optional additive noise for decoder embeddings
    ):
        """
        Forward pass for text to wave parameter generation
        
        Args:
            input_text: List of text descriptions
            target_frames: Target frame indices to predict
            history_frames: Optional history frames for conditioning
            relative_times: Relative time positions [-1, 1] for each target frame
        """
        batch_size = len(input_text)
        num_target_frames = target_frames.shape[1]
        
        # Format input text for T5 task
        # Use standard T5 format for text-to-text generation
        formatted_text = [f"translate to wave: {text}" for text in input_text]
        
        # Tokenize input text
        text_inputs = self.tokenizer(
            formatted_text,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(target_frames.device)
        

        # Encode text with T5
        try:
            # First, let's try a simple forward pass with dummy decoder input
            # T5 expects decoder_input_ids starting with pad token
            decoder_start_token_id = self.t5_model.config.pad_token_id
            decoder_input_ids = torch.full(
                (batch_size, 1),
                decoder_start_token_id,
                dtype=torch.long,
                device=text_inputs.input_ids.device
            )

            # Try using the full model forward pass
            outputs = self.t5_model(
                input_ids=text_inputs.input_ids,
                attention_mask=text_inputs.attention_mask,
                decoder_input_ids=decoder_input_ids,
                return_dict=True,
                output_hidden_states=True
            )

            encoder_outputs = outputs.encoder_last_hidden_state
        except Exception as e:
            if 'log_message' in globals():
                log_message(f"ERROR in encoder: {e}")
            else:
                print(f"ERROR in encoder: {e}")
            raise

        # Autoregressive decoding with history conditioning
        history_buffer, physics_state = self._initialize_history_state(
            history_frames,
            batch_size,
            target_frames.device,
        )

        if static_object_params is not None:
            static_object_params = static_object_params.to(
                device=target_frames.device,
                dtype=torch.float32,
            )

        if noise is not None:
            noise = noise.to(device=encoder_outputs.device, dtype=encoder_outputs.dtype)

        outputs = []

        for f in range(num_target_frames):
            if self.random_history_sampling:
                max_available = min(len(history_buffer), self.max_history_frames)
                if max_available > 0:
                    use_history = int(torch.randint(
                        low=0,
                        high=max_available + 1,
                        size=(1,),
                        device=encoder_outputs.device,
                    ).item())
                else:
                    use_history = 0
            else:
                use_history = min(len(history_buffer), self.max_history_frames)

            if relative_times is not None:
                time_input = relative_times[:, f:f+1].unsqueeze(-1)
                time_embed = self.time_embed(time_input).squeeze(1)
            else:
                time_embed = torch.zeros(
                    batch_size,
                    self.hidden_size,
                    device=encoder_outputs.device,
                )

            history_embed = self._build_history_embedding(history_buffer, physics_state, use_history)
            decoder_embed = time_embed + history_embed
            if noise is not None:
                decoder_embed = decoder_embed + noise

            decoder_output = self.t5_model.decoder(
                inputs_embeds=decoder_embed.unsqueeze(1),  # [batch, 1, hidden_size]
                encoder_hidden_states=encoder_outputs,
                encoder_attention_mask=text_inputs.attention_mask,
            )

            hidden = decoder_output.last_hidden_state[:, 0]  # [batch, hidden_size]

            object_params = self.object_proj(hidden).view(batch_size, self.max_objects, self.object_param_dim)
            if static_object_params is not None:
                # Preserve the first 6 dimensions (exists + shape + scale) from provided static parameters
                static_slice = static_object_params[:, :, :6]
                if static_slice.shape[-1] < 6:
                    pad_width = 6 - static_slice.shape[-1]
                    pad = torch.zeros(*static_slice.shape[:-1], pad_width, device=object_params.device)
                    static_slice = torch.cat([static_slice, pad], dim=-1)
                object_params = object_params.clone()
                object_params[:, :, :6] = static_slice
            world_params = self.world_proj(hidden)
            physics_params = self.physics_proj(hidden).view(batch_size, self.max_objects, 3)

            outputs.append({
                'objects': object_params,
                'world': world_params,
                'physics': physics_params,
            })

            history_buffer.append({
                'objects': object_params,
                'world': world_params,
            })
            if len(history_buffer) > self.max_history_frames:
                history_buffer = history_buffer[-self.max_history_frames:]

            physics_state = physics_params

        return outputs


class BidirectionalTrainer:
    """Trainer for bidirectional prediction from middle frame"""
    
    def __init__(
        self,
        model: Text2WaveModel,
        config: Dict,
        accelerator: Accelerator,
    ):
        self.model = model
        self.config = config
        self.accelerator = accelerator
        base_model = accelerator.unwrap_model(model) if hasattr(accelerator, "unwrap_model") else model
        self.object_param_dim = getattr(base_model, "object_param_dim", 12)
        self.freeze_static_params = bool(config['training'].get('freeze_static_from_anchor', True))
        self.base_model = base_model
        self.sample_attempts = int(config['training'].get('multi_sample_attempts', 1))
        self.sample_attempts = max(1, self.sample_attempts)
        
        # Loss functions
        self.world_loss_fn = nn.MSELoss()
        self.physics_loss_fn = nn.MSELoss()
        
        # Loss weights from config
        loss_weights_config = config.get('loss', {}).get('weights', {})
        self.loss_weights = {
            'wave_loss(superquadric)': loss_weights_config.get('wave_loss', 1.0),
            'wave_contrastive_loss': loss_weights_config.get('wave_contrastive_loss', 2.0),
            'world_info_loss(camera,scale,time)': loss_weights_config.get('world_info_loss', 0.5),
            'controllable_info_loss(mass,friction,restitution)': loss_weights_config.get('controllable_info_loss', 0.1),
            'pla_loss': loss_weights_config.get('pla_loss', 3.0),
        }

        physics_cfg = config.get('physics', {})
        self.gravity = float(physics_cfg.get('gravity', 9.81))
        self.collision_buffer = float(physics_cfg.get('collision_buffer', 1.05))

        # Temporal configuration (dataset cached at 8 fps by default)
        self.frame_rate = float(config['training'].get('frame_rate', 8.0))
        self.frame_rate = max(self.frame_rate, 1e-6)

        presence_cfg = config.get('loss', {}).get('wave_presence', {})
        self.wave_count_weight = float(presence_cfg.get('count_weight', 0.2))
        self.wave_presence_threshold = float(presence_cfg.get('scale_threshold', 0.1))
        self.wave_presence_temperature = float(presence_cfg.get('temperature', 0.1))
        contrastive_cfg = config.get('loss', {}).get('wave_contrastive', {})
        self.wave_contrastive_temperature = float(contrastive_cfg.get('temperature', 0.2))

        # By convention the last three learnable slots before the inlier ratio store velocity
        self.velocity_slice = slice(max(self.object_param_dim - 3, 0), self.object_param_dim)

    def compute_loss(
        self,
        predictions: List[Dict],
        targets: Dict[str, torch.Tensor],
        frame_indices: List[int],
    ) -> Dict[str, torch.Tensor]:
        """Compute losses for predicted frames"""
        losses = {
            'wave_loss(superquadric)': 0.0,  # Wave loss (superquadric parameters)
            'wave_contrastive_loss': 0.0,  # Sequence-level contrastive alignment
            'world_info_loss(camera,scale,time)': 0.0,  # World info loss (camera, scale, relative time)
            'controllable_info_loss(mass,friction,restitution)': 0.0,  # Controllable info loss (mass, friction, restitution)
            'pla_loss': 0.0,  # Physical plausibility regularizer
            'wave_count_mse': 0.0,  # Count alignment between predicted and target waves
            'total': 0.0,
        }

        pla_entries = []
        pred_summaries: List[torch.Tensor] = []
        target_summaries: List[torch.Tensor] = []

        for i, (pred, frame_idx) in enumerate(zip(predictions, frame_indices)):
            # Object loss (only for existing objects)
            target_objects = targets['objects'][:, frame_idx]  # [batch, max_objects, 16]
            if target_objects.shape[-1] < self.object_param_dim:
                pad_width = self.object_param_dim - target_objects.shape[-1]
                pad = target_objects.new_zeros(*target_objects.shape[:-1], pad_width)
                target_objects = torch.cat([target_objects, pad], dim=-1)
            pred_objects = pred['objects']  # [batch, max_objects, self.object_param_dim]

            # Extract existence mask from target
            exists_mask = target_objects[:, :, 0] > 0.5  # [batch, max_objects]

            target_core = target_objects[:, :, :self.object_param_dim]

            # Sequence-level reconstruction with velocity-aware weighting
            object_loss = self._wave_reconstruction_loss(pred_objects, target_core, exists_mask)
            losses['wave_loss(superquadric)'] += object_loss

            # Soft count alignment using scale magnitude as presence proxy
            target_presence = target_objects[:, :, 0].float()
            pred_scale_norm = torch.linalg.norm(pred_objects[:, :, 3:6], dim=-1)
            presence_input = (pred_scale_norm - self.wave_presence_threshold) / max(self.wave_presence_temperature, 1e-6)
            pred_presence = torch.sigmoid(presence_input)
            pred_count = pred_presence.sum(dim=-1)
            target_count = target_presence.sum(dim=-1)
            count_mse = F.mse_loss(pred_count, target_count)
            losses['wave_count_mse'] += count_mse
            losses['wave_loss(superquadric)'] += self.wave_count_weight * count_mse

            pla_entries.append({
                'frame_idx': frame_idx,
                'pred_objects': pred_objects,
                'exists_mask': exists_mask,
            })

            # Aggregate summaries for contrastive objective
            mask = exists_mask.float().unsqueeze(-1)
            # Avoid division by zero by clamping the counts before inversion
            denom = mask.sum(dim=1).clamp_min(1.0)
            pred_summary = (pred_objects * mask).sum(dim=1) / denom
            target_summary = (target_core * mask).sum(dim=1) / denom
            pred_summaries.append(pred_summary)
            target_summaries.append(target_summary)

            # World loss
            target_world = targets['world'][:, frame_idx]  # [batch, 11]
            pred_world = pred['world']  # [batch, 8]

            # Compare only first 8 dimensions
            world_loss = self.world_loss_fn(
                pred_world,
                target_world[:, :8]
            )
            losses['world_info_loss(camera,scale,time)'] += world_loss
            
            # Physics loss (constant across frames, use frame 0)
            if i == 0:
                target_physics = targets['physics']  # [batch, max_objects, 3]
                pred_physics = pred['physics']  # [batch, max_objects, 3]
                
                physics_loss = self.physics_loss_fn(
                    pred_physics[exists_mask],
                    target_physics[exists_mask]
                )
                losses['controllable_info_loss(mass,friction,restitution)'] = physics_loss

        # Average over frames
        num_frames = len(predictions)
        losses['wave_loss(superquadric)'] /= num_frames
        losses['world_info_loss(camera,scale,time)'] /= num_frames
        losses['wave_count_mse'] /= num_frames

        # Anchor PLA loss around the observed middle frame to provide a reference state
        total_frames = targets['objects'].shape[1]
        middle_idx = total_frames // 2
        anchor_objects = targets['objects'][:, middle_idx]
        anchor_exists = anchor_objects[:, :, 0] > 0.5
        pla_entries.append({
            'frame_idx': middle_idx,
            'pred_objects': anchor_objects[:, :, :self.object_param_dim].detach(),
            'exists_mask': anchor_exists,
        })

        # Physical regularizer
        pla_loss = self._compute_pla_regularizer(pla_entries)
        losses['pla_loss'] = pla_loss

        # Contrastive alignment between predicted and target trajectories
        if pred_summaries:
            pred_stack = torch.stack(pred_summaries, dim=0).mean(dim=0)
            target_stack = torch.stack(target_summaries, dim=0).mean(dim=0)
            losses['wave_contrastive_loss'] = self._contrastive_clip_loss(pred_stack, target_stack)
        else:
            device = targets['objects'].device
            losses['wave_contrastive_loss'] = torch.zeros((), device=device)

        # Compute total loss
        for key, weight in self.loss_weights.items():
            if key in losses:
                losses['total'] += weight * losses[key]

        return losses

    def _wave_reconstruction_loss(
        self,
        pred_objects: torch.Tensor,
        target_objects: torch.Tensor,
        exists_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Velocity-aware reconstruction loss combining position L1 and velocity L1."""
        device = pred_objects.device
        dtype = pred_objects.dtype
        if not exists_mask.any():
            return torch.zeros((), device=device, dtype=dtype)

        pred_active = pred_objects[exists_mask]
        target_active = target_objects[exists_mask]

        base_l1 = F.l1_loss(pred_active, target_active, reduction='mean')

        if self.velocity_slice.start >= self.velocity_slice.stop:  # degenerate slice when dim < 3
            velocity_l1 = torch.zeros((), device=device, dtype=dtype)
        else:
            pred_velocity = pred_active[..., self.velocity_slice]
            target_velocity = target_active[..., self.velocity_slice]
            velocity_l1 = F.l1_loss(pred_velocity, target_velocity, reduction='mean')

        return 0.5 * base_l1 + 0.5 * velocity_l1

    def _contrastive_clip_loss(
        self,
        pred_summary: torch.Tensor,
        target_summary: torch.Tensor,
    ) -> torch.Tensor:
        """InfoNCE-style contrastive loss between predicted and target clip summaries."""
        device = pred_summary.device
        dtype = pred_summary.dtype
        batch = pred_summary.size(0)
        if batch <= 1:
            return torch.zeros((), device=device, dtype=dtype)

        dim = min(pred_summary.size(-1), target_summary.size(-1))
        if dim == 0:
            return torch.zeros((), device=device, dtype=dtype)
        if pred_summary.size(-1) != dim:
            pred_summary = pred_summary[..., :dim]
        if target_summary.size(-1) != dim:
            target_summary = target_summary[..., :dim]

        temperature = max(self.wave_contrastive_temperature, 1e-6)
        pred_norm = F.normalize(pred_summary, dim=-1)
        target_norm = F.normalize(target_summary, dim=-1)
        dim_post = min(pred_norm.size(-1), target_norm.size(-1))
        if dim_post == 0:
            return torch.zeros((), device=device, dtype=dtype)
        if pred_norm.size(-1) != dim_post:
            pred_norm = pred_norm[..., :dim_post]
        if target_norm.size(-1) != dim_post:
            target_norm = target_norm[..., :dim_post]
        logits = pred_norm @ target_norm.transpose(0, 1)
        logits = logits / temperature

        labels = torch.arange(batch, device=device)
        loss_forward = F.cross_entropy(logits, labels)
        loss_backward = F.cross_entropy(logits.transpose(0, 1), labels)

        return 0.5 * (loss_forward + loss_backward)

    def _compute_pla_regularizer(self, entries: List[Dict[str, torch.Tensor]]) -> torch.Tensor:
        """Encourage rigid-body consistency, free-fall dynamics, and collision plausibility."""
        model_device = next(self.model.parameters()).device
        if not entries:
            return torch.tensor(0.0, device=model_device)

        # Sort by frame index to obtain temporal order
        sorted_entries = sorted(entries, key=lambda x: x['frame_idx'])

        device = sorted_entries[0]['pred_objects'].device
        dtype = sorted_entries[0]['pred_objects'].dtype

        preds = torch.stack([item['pred_objects'] for item in sorted_entries], dim=0)  # [F, B, O, 12]
        exists = torch.stack([item['exists_mask'].float() for item in sorted_entries], dim=0)  # [F, B, O]

        frame_count, batch_size, max_objects, _ = preds.shape

        if frame_count <= 1:
            return torch.tensor(0.0, device=device, dtype=dtype)

        exists_expanded = exists.unsqueeze(-1)
        exists_total = exists_expanded.sum()
        if exists_total.item() == 0:
            return torch.tensor(0.0, device=device, dtype=dtype)

        # 1. Shape and scale invariance for rigid bodies
        shape_params = preds[..., 1:3]
        scale_params = preds[..., 3:6]

        shape_mean = (shape_params * exists_expanded).sum(dim=0) / exists_expanded.sum(dim=0).clamp_min(1.0)
        scale_mean = (scale_params * exists_expanded).sum(dim=0) / exists_expanded.sum(dim=0).clamp_min(1.0)

        shape_loss = ((shape_params - shape_mean) ** 2 * exists_expanded).sum() / exists_expanded.sum().clamp_min(1.0)
        scale_loss = ((scale_params - scale_mean) ** 2 * exists_expanded).sum() / exists_expanded.sum().clamp_min(1.0)

        # 2. Free-fall consistency via discrete Euler-Lagrange residuals
        freefall_loss = torch.tensor(0.0, device=device, dtype=dtype)
        rotation_loss = torch.tensor(0.0, device=device, dtype=dtype)
        collision_penalty = torch.tensor(0.0, device=device, dtype=dtype)
        velocity_loss = torch.tensor(0.0, device=device, dtype=dtype)

        positions = preds[..., 6:9]

        if frame_count >= 3:
            radii = torch.linalg.norm(preds[..., 3:6], dim=-1)

            accel = positions[2:] - 2 * positions[1:-1] + positions[:-2]

            exists_triplet = exists[1:-1] * exists[:-2] * exists[2:]
            exists_triplet_expanded = exists_triplet.unsqueeze(-1)

            # Collision detection to gate free-fall prior
            center_positions = positions[1:-1].reshape(-1, max_objects, 3)
            center_exists = exists[1:-1].reshape(-1, max_objects)
            center_radii = radii[1:-1].reshape(-1, max_objects)

            if center_positions.numel() > 0:
                dist = torch.cdist(center_positions, center_positions, p=2)  # [N, O, O]
                radius_sum = (center_radii.unsqueeze(-1) + center_radii.unsqueeze(-2)) * self.collision_buffer
                exists_pair = center_exists.unsqueeze(-1) * center_exists.unsqueeze(-2)

                eye = torch.eye(max_objects, device=device).unsqueeze(0)
                non_diag = (1 - eye)

                penetration = torch.relu((radius_sum - dist) * non_diag) * exists_pair
                collision_penalty = penetration.pow(2).sum() / (non_diag * exists_pair).sum().clamp_min(1.0)

                contact_any = (penetration > 0).any(dim=-1).view(frame_count - 2, batch_size, max_objects)
            else:
                contact_any = torch.zeros(frame_count - 2, batch_size, max_objects, device=device, dtype=torch.bool)

            contact_mask = contact_any.float()

            gravity_vec = torch.tensor([0.0, 0.0, -self.gravity], device=device, dtype=dtype).view(1, 1, 1, 3)
            residual = accel + gravity_vec

            freefall_mask = exists_triplet_expanded * (1.0 - contact_mask.unsqueeze(-1))
            valid_count = freefall_mask.sum().clamp_min(1.0)
            freefall_loss = (residual.pow(2) * freefall_mask).sum() / valid_count

            rotations = preds[..., 9:12]
            rot_sin = torch.sin(rotations)
            rot_cos = torch.cos(rotations)
            rot_features = torch.cat([rot_sin, rot_cos], dim=-1)
            rot_acc = rot_features[2:] - 2 * rot_features[1:-1] + rot_features[:-2]

            rot_mask = exists_triplet_expanded * (1.0 - contact_mask.unsqueeze(-1))
            rot_valid = rot_mask.sum().clamp_min(1.0)
            rotation_loss = (rot_acc.pow(2) * rot_mask).sum() / rot_valid

        if frame_count >= 2:
            velocities = preds[..., 12:15]
            diff = (positions[1:] - positions[:-1]) * self.frame_rate
            exists_pair = exists[1:] * exists[:-1]
            diff_expanded = exists_pair.unsqueeze(-1)

            velocity_residual = (velocities[1:] - diff).pow(2) * diff_expanded
            valid_velocity = diff_expanded.sum()
            velocity_loss = velocity_residual.sum()

            first_pair = (exists[0] * exists[1]).unsqueeze(-1)
            velocity_loss += ((velocities[0] - diff[0]) ** 2 * first_pair).sum()
            valid_velocity += first_pair.sum()

            velocity_loss = velocity_loss / valid_velocity.clamp_min(1.0)

        pla_loss = (
            shape_loss
            + scale_loss
            + freefall_loss
            + rotation_loss
            + collision_penalty
            + velocity_loss
        )
        return pla_loss

    def _select_anchor_frame(self, num_frames: int) -> int:
        """Determine which frame should serve as the initial anchor."""
        cfg = self.config['training'].get('initial_frame', {})
        strategy = cfg.get('strategy', 'middle')

        if strategy == 'random':
            base_idx = int(torch.randint(low=0, high=num_frames, size=(1,), device=torch.device('cpu')).item())
        elif strategy == 'fixed':
            base_idx = int(cfg.get('index', num_frames // 2))
        else:
            base_idx = num_frames // 2

        offset = int(cfg.get('offset', 0))
        anchor_idx = base_idx + offset
        anchor_idx = max(0, min(num_frames - 1, anchor_idx))
        return anchor_idx

    def _generate_full_sequence(
        self,
        text: List[str],
        objects: torch.Tensor,
        world: torch.Tensor,
        physics: torch.Tensor,
        teacher_prob: float,
        anchor_idx: Optional[int] = None,
        use_noise: bool = False,
    ) -> Tuple[List[Dict[str, torch.Tensor]], List[int], float]:
        """Generate a full sequence of predictions given an anchor frame."""
        batch_size, num_frames = objects.shape[:2]
        if anchor_idx is None:
            anchor_idx = self._select_anchor_frame(num_frames)

        static_object_params = None
        if self.freeze_static_params:
            anchor_static = objects[:, anchor_idx, :, :6]
            static_object_params = anchor_static

        if teacher_prob > 0.0:
            teacher_mask = (torch.rand(batch_size, device=objects.device) < teacher_prob).float()
        else:
            teacher_mask = torch.zeros(batch_size, device=objects.device, dtype=torch.float32)

        def sample_noise():
            return self.base_model.sample_decoder_noise(batch_size, objects.device) if use_noise else None

        half_span = max(num_frames - 1, 1) / 2.0
        inference_time = 0.0
        predictions_by_idx: Dict[int, Dict[str, torch.Tensor]] = {}

        anchor_rel_times = torch.zeros(
            (batch_size, 1), dtype=torch.float32, device=objects.device
        )
        anchor_targets = torch.full(
            (batch_size, 1), anchor_idx, dtype=torch.long, device=objects.device
        )

        start = time.time()
        anchor_preds = self.model(
            input_text=text,
            target_frames=anchor_targets,
            history_frames=None,
            relative_times=anchor_rel_times,
            static_object_params=static_object_params,
            noise=sample_noise(),
        )
        inference_time += time.time() - start
        anchor_pred = anchor_preds[0]
        predictions_by_idx[anchor_idx] = anchor_pred

        anchor_gt_objects = objects[:, anchor_idx, :, :self.object_param_dim]
        if anchor_gt_objects.shape[-1] < self.object_param_dim:
            pad_width = self.object_param_dim - anchor_gt_objects.shape[-1]
            pad = anchor_gt_objects.new_zeros(*anchor_gt_objects.shape[:-1], pad_width)
            anchor_gt_objects = torch.cat([anchor_gt_objects, pad], dim=-1)
        anchor_gt_world = world[:, anchor_idx, :8]
        anchor_pred_objects = anchor_pred['objects']
        if static_object_params is not None:
            anchor_pred_objects[:, :, :6] = static_object_params[:, :, :6]
        anchor_pred_world = anchor_pred['world']

        teacher_mask_objs = teacher_mask.view(batch_size, 1, 1)
        teacher_mask_world = teacher_mask.view(batch_size, 1)

        blended_objects = anchor_pred_objects * (1.0 - teacher_mask_objs) + anchor_gt_objects * teacher_mask_objs
        blended_world = anchor_pred_world * (1.0 - teacher_mask_world) + anchor_gt_world * teacher_mask_world

        history_objects = blended_objects.unsqueeze(1)
        history_world = blended_world.unsqueeze(1)
        history_physics = physics.clone()

        def make_history_seed():
            return {
                'objects': history_objects.clone(),
                'world': history_world.clone(),
                'physics': history_physics.clone(),
            }

        backward_indices = list(range(anchor_idx - 1, -1, -1))
        forward_indices = list(range(anchor_idx + 1, num_frames))

        def run_direction(target_indices: List[int]):
            nonlocal inference_time
            if not target_indices:
                return

            rel_times = torch.tensor(
                [(idx - anchor_idx) / half_span for idx in target_indices],
                dtype=torch.float32,
                device=objects.device,
            ).unsqueeze(0).repeat(batch_size, 1)

            target_tensor = torch.tensor(
                target_indices,
                dtype=torch.long,
                device=objects.device,
            ).unsqueeze(0).repeat(batch_size, 1)

            history_frames = make_history_seed()

            start_time = time.time()
            preds = self.model(
                input_text=text,
                target_frames=target_tensor,
                history_frames=history_frames,
                relative_times=rel_times,
                static_object_params=static_object_params,
                noise=sample_noise(),
            )
            inference_time += time.time() - start_time

            for idx, pred in zip(target_indices, preds):
                if static_object_params is not None:
                    pred['objects'][:, :, :6] = static_object_params[:, :, :6]
                predictions_by_idx[idx] = pred

        run_direction(backward_indices)
        run_direction(forward_indices)

        ordered_indices = list(range(num_frames))
        predictions = [predictions_by_idx[idx] for idx in ordered_indices]
        return predictions, ordered_indices, inference_time

    def _compute_losses(
        self,
        batch: Dict[str, torch.Tensor],
    ) -> Tuple[Dict[str, torch.Tensor], float, int]:
        """Shared logic for computing losses and metadata."""
        text = batch['text']
        objects = batch['objects']  # [batch, num_frames, max_objects, 16]
        world = batch['world']  # [batch, num_frames, 11]
        physics = batch['physics']  # [batch, max_objects, 3]

        batch_size, num_frames = objects.shape[:2]
        anchor_idx = self._select_anchor_frame(num_frames)
        teacher_prob = float(self.config['training'].get('initial_teacher_forcing_prob', 0.5))

        targets = {
            'objects': objects,
            'world': world,
            'physics': physics,
        }

        attempts = self.sample_attempts if self.model.training else 1
        use_noise = attempts > 1
        best_losses: Optional[Dict[str, torch.Tensor]] = None
        best_predictions: Optional[List[Dict[str, torch.Tensor]]] = None
        best_frame_indices: Optional[List[int]] = None
        best_inference_time: float = 0.0
        best_total_value: Optional[float] = None

        for attempt in range(attempts):
            predictions, frame_indices, inference_time = self._generate_full_sequence(
                text=text,
                objects=objects,
                world=world,
                physics=physics,
                teacher_prob=teacher_prob,
                anchor_idx=anchor_idx,
                use_noise=use_noise,
            )

            losses = self.compute_loss(predictions, targets, frame_indices)
            total_value = float(losses['total'].detach())
            if best_total_value is None or total_value < best_total_value:
                if best_losses is not None:
                    del best_losses
                if best_predictions is not None:
                    del best_predictions
                best_total_value = total_value
                best_losses = losses
                best_predictions = predictions
                best_frame_indices = frame_indices
                best_inference_time = inference_time
            else:
                del losses
                del predictions
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()

        assert best_losses is not None and best_predictions is not None and best_frame_indices is not None
        num_predicted_frames = len(best_predictions)
        frames_per_second = num_predicted_frames / best_inference_time if best_inference_time > 0 else 0.0

        return best_losses, frames_per_second, num_predicted_frames
    
    def train_step(
        self,
        batch: Dict[str, torch.Tensor],
        step: int,
    ) -> Dict[str, float]:
        """Single training step with bidirectional prediction"""
        self.model.train()

        losses, frames_per_second, num_predicted_frames = self._compute_losses(batch)

        self.accelerator.backward(losses['total'])

        loss_dict = {k: v.item() if torch.is_tensor(v) else float(v) for k, v in losses.items()}
        loss_dict['inference_fps'] = frames_per_second
        loss_dict['frames_predicted'] = num_predicted_frames

        return loss_dict

    def evaluate_batch(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
        """Compute losses without gradient updates."""
        was_training = self.model.training
        self.model.eval()
        with torch.no_grad():
            losses, frames_per_second, num_predicted_frames = self._compute_losses(batch)
        if was_training:
            self.model.train()

        loss_dict = {k: v.item() if torch.is_tensor(v) else float(v) for k, v in losses.items()}
        loss_dict['inference_fps'] = frames_per_second
        loss_dict['frames_predicted'] = num_predicted_frames
        return loss_dict





def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--train_config', type=str, default='configs/default.yaml',
                       help='Training configuration file')
    parser.add_argument('--data_root', type=str,
                       default='../data/movi_a_128x128',
                       help='Root directory of MOVi dataset')
    parser.add_argument('--output_dir', type=str, default='core_space',
                       help='Directory to save checkpoints and generation results')
    parser.add_argument('--resume_step', type=int, default=None,
                       help='Resume training from specific step')
    args = parser.parse_args()
    
    # Load training config
    with open(args.train_config, 'r') as f:
        config = yaml.safe_load(f)
        
    # Initialize accelerator with DDP configuration
    from accelerate import DistributedDataParallelKwargs
    
    ddp_kwargs = DistributedDataParallelKwargs(
        find_unused_parameters=True,
        broadcast_buffers=False
    )
    
    # 注意：混合精度通过launch_text2wave_training.sh中的--mixed_precision参数控制
    # 如果遇到NaN问题，请确保shell脚本中没有启用mixed_precision
    accelerator = Accelerator(
        gradient_accumulation_steps=1,
        kwargs_handlers=[ddp_kwargs]
    )
    
    # Set seed
    set_seed(42)
    
    # Create model
    model_name = config.get('text2wave_model', {}).get('model_name', "google/t5-v1_1-small")
    model = Text2WaveModel(
        model_name=model_name,
        max_objects=10,
        num_frames=24,
        max_history_frames=config['training']['max_history_frames'],
        random_history_sampling=config['training'].get('random_history_sampling', True),
        decoder_noise_std=config['training'].get('decoder_noise_std', 0.0),
    )
    
    # Create optimizer
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=config['training']['learning_rate'],
        weight_decay=0.01,
    )
    
    # Create dataloaders
    train_dataloader = create_dataloader(
        data_root=args.data_root,
        split='train',
        batch_size=config['training']['batch_size'],
        num_workers=config['data']['num_workers'],
        shuffle=True,
        max_samples=config['data'].get('max_sequences', -1),
    )
    
    val_dataloader = create_dataloader(
        data_root=args.data_root,
        split='validation',
        batch_size=config['training']['batch_size'],
        num_workers=config['data']['num_workers'],
        shuffle=False,
        max_samples=10,  # Use only 10 validation samples
    )
    
    # Prepare for distributed training
    model, optimizer, train_dataloader, val_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, val_dataloader
    )
    
    checkpoint_dir = Path("checkpoints_text2wave")
    if accelerator.is_main_process:
        checkpoint_dir.mkdir(parents=True, exist_ok=True)

    log_file_path = checkpoint_dir / "training_log.txt"

    def log_message(message: str):
        """Log to stdout and append to training_log.txt from main process."""
        if not accelerator.is_main_process:
            return
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        formatted = f"{timestamp} {message}"
        print(formatted)
        try:
            with open(log_file_path, 'a') as fp:
                fp.write(formatted + "\n")
        except Exception:
            pass

    best_metrics_path = checkpoint_dir / "best_metrics.json"
    if best_metrics_path.exists():
        try:
            best_metrics_path.unlink()
        except OSError as exc:
            log_message(f"Warning: failed to remove legacy best_metrics.json due to {exc}")

    best_train_loss = float('inf')
    best_val_loss = float('inf')

    evaluation_cfg = config['training'].get('evaluation', {})
    eval_max_batches = evaluation_cfg.get('max_batches', 5)

    training_stats_path = checkpoint_dir / "training_stats.npz"
    loaded_step_history: Optional[List[int]] = None
    loaded_loss_history: Dict[str, List[float]] = {}
    if training_stats_path.exists():
        try:
            stats = np.load(training_stats_path, allow_pickle=True)
            best_train_loss = float(stats.get('best_train_loss', best_train_loss))
            best_val_loss = float(stats.get('best_val_loss', best_val_loss))
            if 'step_history' in stats:
                loaded_step_history = stats['step_history'].tolist()
            if 'loss_history_keys' in stats and 'loss_history_values' in stats:
                keys = stats['loss_history_keys'].tolist()
                values = stats['loss_history_values'].tolist()
                for key, value in zip(keys, values):
                    loaded_loss_history[str(key)] = list(np.asarray(value, dtype=float))
        except Exception as exc:
            log_message(f"Warning: failed to load training_stats.npz due to {exc}")

    executor = ThreadPoolExecutor(max_workers=1)
    pending_futures: List = []

    def cleanup_futures():
        pending_futures[:] = [f for f in pending_futures if not f.done()]

    def submit_task(fn, *args, **kwargs):
        cleanup_futures()
        future = executor.submit(fn, *args, **kwargs)
        pending_futures.append(future)
        return future

    def recursive_to_cpu(obj):
        if isinstance(obj, torch.Tensor):
            return obj.detach().cpu()
        if isinstance(obj, dict):
            return {k: recursive_to_cpu(v) for k, v in obj.items()}
        if isinstance(obj, list):
            return [recursive_to_cpu(v) for v in obj]
        if isinstance(obj, tuple):
            return tuple(recursive_to_cpu(v) for v in obj)
        return obj

    def save_checkpoint_async(path: Path, payload: Dict):
        def _task():
            torch.save(payload, path)
        submit_task(_task)

    def save_generation_async(predictions: List[Dict], targets: Dict[str, torch.Tensor], texts: List[str], step: int, save_config: Dict, metadata: Dict, batch_data: Dict, data_root: str, data_split: str):
        def _task():
            save_generation_results(
                predictions=predictions,
                targets=targets,
                texts=texts,
                step=step,
                output_dir=args.output_dir,
                save_config=save_config,
                metadata=metadata,
                batch_data=batch_data,
                data_root=data_root,
                data_split=data_split
            )
        submit_task(_task)

    def compute_validation_loss(max_batches: Optional[int]) -> Optional[float]:
        limit = -1 if max_batches is None else max_batches
        if limit == 0:
            return None
        total = 0.0
        count = 0
        for batch_idx, val_batch in enumerate(val_dataloader):
            val_losses = trainer.evaluate_batch(val_batch)
            total += val_losses['total']
            count += 1
            if limit > 0 and (batch_idx + 1) >= limit:
                break
        if count == 0:
            return None
        return total / count

    # Create trainer
    trainer = BidirectionalTrainer(model, config, accelerator)
    
    # Get max_steps from config
    max_steps = config['training']['max_steps']
    
    # Calculate and display dataset traversal information
    if accelerator.is_main_process:
        steps_per_epoch = len(train_dataloader)
        total_epochs = max_steps / steps_per_epoch
        log_message("=" * 60)
        log_message("Dataset Information:")
        log_message(f"- Training samples: {len(train_dataloader.dataset) if hasattr(train_dataloader, 'dataset') else 'N/A'}")
        log_message(f"- Batch size: {config['training']['batch_size']}")
        log_message(f"- Steps per epoch (full dataset): {steps_per_epoch}")
        log_message(f"- Total training steps: {max_steps}")
        log_message(f"- Will traverse dataset: {total_epochs:.2f} times")
        log_message("=" * 60)
    
    # Resume from checkpoint if specified
    start_step = 0
    resumed_from = None
    if args.resume_step is not None:
        checkpoint_path = checkpoint_dir / f"step{args.resume_step}.pt"
        if checkpoint_path.exists():
            log_message(f"Resuming from checkpoint step {args.resume_step}")
            checkpoint = torch.load(checkpoint_path, map_location='cpu')
            accelerator.unwrap_model(model).load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            start_step = checkpoint.get('step', args.resume_step)
            resumed_from = checkpoint_path
        else:
            log_message(f"Warning: Checkpoint for step {args.resume_step} not found, starting from scratch")
    else:
        latest_checkpoint_path = checkpoint_dir / "latest.pt"
        if latest_checkpoint_path.exists():
            try:
                log_message("Resuming from latest checkpoint")
                checkpoint = torch.load(latest_checkpoint_path, map_location='cpu')
                accelerator.unwrap_model(model).load_state_dict(checkpoint['model_state_dict'])
                optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
                start_step = checkpoint.get('step', 0)
                resumed_from = latest_checkpoint_path
            except Exception as exc:
                log_message(f"Warning: failed to load latest checkpoint due to {exc}; attempting best checkpoint")
                try:
                    corrupt_path = latest_checkpoint_path.with_suffix(latest_checkpoint_path.suffix + ".corrupt")
                    latest_checkpoint_path.rename(corrupt_path)
                    log_message(f"Renamed corrupt latest checkpoint to {corrupt_path.name}")
                except Exception as rename_exc:
                    log_message(f"Warning: could not rename corrupt latest checkpoint: {rename_exc}")
        if resumed_from is None:
            best_checkpoint_path = checkpoint_dir / "best.pt"
            if best_checkpoint_path.exists():
                try:
                    log_message("Resuming from best checkpoint")
                    checkpoint = torch.load(best_checkpoint_path, map_location='cpu')
                    accelerator.unwrap_model(model).load_state_dict(checkpoint['model_state_dict'])
                    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
                    start_step = checkpoint.get('step', 0)
                    resumed_from = best_checkpoint_path
                except Exception as exc:
                    log_message(f"Warning: failed to load best checkpoint due to {exc}; starting from scratch")

    # Setup local logging and plotting
    log_dir = checkpoint_dir
    loss_history = defaultdict(list)
    step_history: List[int] = []
    if loaded_step_history:
        step_history.extend(int(s) for s in loaded_step_history)
    if loaded_loss_history:
        for key, values in loaded_loss_history.items():
            loss_history[key].extend(values)
    last_plot_time = time.time()
    plot_path = log_dir / "losses.png"

    def save_training_stats():
        if not accelerator.is_main_process:
            return
        keys = sorted(loss_history.keys())
        loss_arrays = [np.array(loss_history[k], dtype=np.float32) for k in keys]
        np.savez(
            training_stats_path,
            best_train_loss=best_train_loss,
            best_val_loss=best_val_loss,
            step_history=np.array(step_history, dtype=np.int64),
            loss_history_keys=np.array(keys, dtype=object),
            loss_history_values=np.array(loss_arrays, dtype=object),
        )

    def update_loss_plot():
        if not accelerator.is_main_process or not step_history:
            return
        x_values = np.array(step_history, dtype=np.int64)
        keys = [k for k, v in sorted(loss_history.items()) if v]
        if not keys:
            return

        def align_series(series: List[float]) -> np.ndarray:
            y_vals = np.array(series, dtype=np.float32)
            if len(y_vals) > len(x_values):
                y_vals = y_vals[-len(x_values):]
            elif len(y_vals) < len(x_values):
                pad = np.full(len(x_values) - len(y_vals), np.nan, dtype=np.float32)
                y_vals = np.concatenate([pad, y_vals])
            return y_vals

        fig_height = 3 * (len(keys) + 1)
        fig, axes = plt.subplots(len(keys) + 1, 1, figsize=(10, fig_height), sharex=True)
        if not isinstance(axes, np.ndarray):
            axes = np.array([axes])

        cmap = plt.get_cmap('tab10', len(keys))

        aggregated_ax = axes[0]
        aggregated_ax.set_title("Training Losses (all)")
        aggregated_ax.set_ylabel("Loss")
        aggregated_ax.grid(True, alpha=0.3)

        for idx, key in enumerate(keys):
            y_aligned = align_series(loss_history[key])
            if np.all(np.isnan(y_aligned)):
                continue
            color = cmap(idx % cmap.N)
            aggregated_ax.plot(x_values, y_aligned, label=key, color=color)
            ax = axes[idx + 1]
            ax.plot(x_values, y_aligned, color=color)
            ax.set_ylabel(key)
            ax.grid(True, alpha=0.3)

        axes[-1].set_xlabel("Step")
        aggregated_ax.legend()
        fig.tight_layout()
        fig.savefig(plot_path)
        plt.close(fig)
        save_training_stats()

    if accelerator.is_main_process and step_history:
        update_loss_plot()

    # Training loop
    global_step = start_step
    
    with tqdm(total=max_steps, initial=start_step, disable=not accelerator.is_local_main_process, position=0, leave=True) as pbar:
        while global_step < max_steps:
            for batch in train_dataloader:
                # Training step
                losses = trainer.train_step(batch, global_step)
                
                # Update progress
                if accelerator.is_local_main_process:
                    pbar.update(1)
                    # Add fps info to losses for display
                    display_losses = losses.copy()
                    display_losses['fps'] = losses['inference_fps']
                    pbar.set_postfix(display_losses)
                    
                    # Print step info to create a log history
                    loss_str = f"Step {global_step}: "
                    for k, v in losses.items():
                        if k not in ['inference_fps', 'frames_predicted']:
                            loss_str += f"{k}={v:.4f} "
                    loss_str += f"| {losses['frames_predicted']} frames @ {losses['inference_fps']:.1f} fps (training speed, inference faster)"
                    tqdm.write(loss_str)
                    
                    
                if accelerator.is_main_process:
                    step_history.append(global_step)
                    for k, v in losses.items():
                        if k in ['inference_fps', 'frames_predicted']:
                            continue
                        loss_history[k].append(v)
                    current_time = time.time()
                    if current_time - last_plot_time >= 10:
                        update_loss_plot()
                        last_plot_time = current_time

                # Save checkpoint and generation results
                # Save at step 5 for testing, then at regular intervals
                save_condition = (global_step == 5) or (global_step > 0 and global_step % config['training']['save_generation']['save_interval'] == 0)
                if save_condition:
                    if accelerator.is_main_process:
                        generation_save_dir = Path(args.output_dir)
                        generation_save_dir.mkdir(parents=True, exist_ok=True)

                        current_train_loss = losses['total']
                        val_loss = compute_validation_loss(eval_max_batches)

                        model_state = recursive_to_cpu(accelerator.get_state_dict(model))
                        optimizer_state = recursive_to_cpu(optimizer.state_dict())
                        payload = {
                            'step': global_step,
                            'model_state_dict': model_state,
                            'optimizer_state_dict': optimizer_state,
                            'config': config,
                        }

                        latest_checkpoint_path = checkpoint_dir / "latest.pt"
                        save_checkpoint_async(latest_checkpoint_path, dict(payload))
                        save_training_stats()

                        is_new_best = False
                        if val_loss is not None:
                            if val_loss < best_val_loss:
                                best_val_loss = val_loss
                                best_train_loss = min(best_train_loss, current_train_loss)
                                is_new_best = True
                        else:
                            if current_train_loss < best_train_loss:
                                best_train_loss = current_train_loss
                                is_new_best = True

                        if is_new_best:
                            best_checkpoint_path = checkpoint_dir / "best.pt"
                            save_checkpoint_async(best_checkpoint_path, dict(payload))
                            save_training_stats()
                            if val_loss is not None:
                                log_message(f"New best checkpoint at step {global_step}: train_loss={current_train_loss:.6f}, val_loss={val_loss:.6f}")
                            else:
                                log_message(f"New best checkpoint at step {global_step}: train_loss={current_train_loss:.6f}")

                            if config['training']['save_generation']['enabled']:
                                with torch.no_grad():
                                    val_batch = next(iter(val_dataloader))
                                    texts = val_batch['text'][:5]
                                val_objects = val_batch['objects'][:5]
                                val_world = val_batch['world'][:5]
                                val_physics = val_batch.get('physics')
                                if val_physics is not None:
                                    val_physics = val_physics[:5]
                                else:
                                    val_physics = torch.zeros_like(val_objects[:, 0, :, :3])
                                val_device = val_objects.device
                                val_batch_size, val_num_frames = val_objects.shape[:2]
                                anchor_idx = trainer._select_anchor_frame(val_num_frames)
                                predictions, generated_indices, _ = trainer._generate_full_sequence(
                                    text=texts,
                                    objects=val_objects,
                                    world=val_world,
                                    physics=val_physics,
                                    teacher_prob=0.0,
                                    anchor_idx=anchor_idx,
                                )

                            val_objects_cpu = val_objects.detach().cpu()
                            val_world_cpu = val_world.detach().cpu()
                            val_physics_cpu = val_physics.detach().cpu()
                            val_batch_cpu = recursive_to_cpu(val_batch)
                            predictions_cpu = [{
                                'objects': pred['objects'].detach().cpu(),
                                'world': pred['world'].detach().cpu(),
                                'physics': pred['physics'].detach().cpu(),
                            } for pred in predictions]
                            targets_cpu = {
                                'objects': val_objects_cpu,
                                'world': val_world_cpu,
                                'physics': val_physics_cpu,
                            }
                            metadata = {
                                'sequence_names': val_batch.get('sequence_names', None)[:5] if 'sequence_names' in val_batch else None,
                                'generated_indices': generated_indices,
                            }
                            save_generation_async(
                                predictions=predictions_cpu,
                                targets=targets_cpu,
                                texts=list(texts),
                                step=global_step,
                                save_config=config['training']['save_generation'],
                                metadata=metadata,
                                batch_data=val_batch_cpu,
                                data_root=args.data_root,
                                data_split='validation'
                            )
                        else:
                            msg = f"No improvement at step {global_step}: train_loss={current_train_loss:.6f}"
                            if val_loss is not None:
                                msg += f", val_loss={val_loss:.6f}"
                            log_message(msg)

                # Gradient clipping before optimizer step
                if accelerator.sync_gradients:
                    clip_val = config['training'].get('gradient_clip_val', 1.0)
                    accelerator.clip_grad_norm_(model.parameters(), max_norm=clip_val)
                
                optimizer.step()
                optimizer.zero_grad()
                
                global_step += 1
                
                if global_step >= max_steps:
                    break
                    
    # Ensure latest plot is written
    if accelerator.is_main_process:
        update_loss_plot()

    # Ensure asynchronous tasks complete before final save
    executor.shutdown(wait=True)

    # Final save
    if accelerator.is_main_process:
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
        final_checkpoint_path = checkpoint_dir / f"step{global_step}_final.pt"
        
        torch.save({
            'step': global_step,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'config': config,
        }, final_checkpoint_path)
        
        # Update best.pt to point to final checkpoint
        best_path = checkpoint_dir / "best.pt"
        if best_path.exists() or best_path.is_symlink():
            best_path.unlink()
        best_path.symlink_to(final_checkpoint_path.name)

        log_message(f"Saved final checkpoint: {final_checkpoint_path}")
        
        

if __name__ == "__main__":
    main()