Blame for src/train.py

"""
Main training script for RigidFormer, a multi-body rigid dynamics model.

The network predicts per-anchor accelerations, which are integrated with a
Verlet scheme and projected onto rigid motions via Kabsch alignment.
"""

import os
import sys
import argparse
import logging
from pathlib import Path
import yaml
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from tqdm import tqdm
import numpy as np
from datetime import datetime
import random
import subprocess
import json

try:
    import roma  # RoMa library for differentiable rigid registration
    ROMA_AVAILABLE = True
except ImportError:
    ROMA_AVAILABLE = False
    roma = None


# Add parent directory to path
sys.path.append(str(Path(__file__).parent.parent))

from src.data.multi_dataloader import create_multi_dataloaders
from src.models.encoder import RigidBodyEncoder, EncoderPack
from src.models.state import PersistentStateManager
from src.models.decoder import BidirectionalDecoder
from src.models.predictor import DynamicsPredictor
from src.models.transform_layer import create_transform_layer

from src.utils.intersection_handler import handle_intersections
from src.models.physformer_rope import rope_memo_reset
from src.utils.weight_ema import WeightEMA
from src.utils.anchor_rotation import centered_anchor_error

# Require CUDA KNN (no fallback). If unavailable, raise immediately.
try:
    from src.utils.cuda_knn import compute_nearest_distance_to_list_cuda, compute_knn_all_objects_vectorized, \
        knn_pack, compute_knn_packed
except Exception as e:
    raise RuntimeError(
        "CUDA KNN is required but not available. Please build/enable src.utils.cuda_knn."
    ) from e
# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Per-step NaN/Inf diagnostics inside the training forward (each costs host syncs). They only log, so they are
# opt-in: DEBUG_NAN=1. Non-finite losses are still caught by the training loop.
DEBUG_NAN = os.environ.get('DEBUG_NAN', '0') == '1'


def distance_2_plane(points, plane_point, plane_axis, body_dist_vec):
    """
    Check if points are closer to a plane than to another body, and update distance vectors.
    
    Args:
        points: Tensor of shape (B, M, 3) containing point positions
        plane_point: List [x, y, z] defining a point on the plane (e.g., [0,0,0])
        plane_axis: String "x", "y", or "z" defining the plane normal direction
        body_dist_vec: Tensor of shape (B, M, 3) containing vectors to nearest points on other body
    
    Returns:
        Updated distance vectors (B, M, 3) where vectors to plane replace body vectors when closer
    """
    B, M, _ = points.shape
    device = points.device
    
    # Get the axis index
    axis_map = {"x": 0, "y": 1, "z": 2}
    if plane_axis not in axis_map:
        raise ValueError(f"Invalid plane_axis: {plane_axis}. Must be 'x', 'y', or 'z'")
    axis_idx = axis_map[plane_axis]
    
    # Get plane position on the specified axis
    plane_pos = plane_point[axis_idx]
    
    # Calculate distance to plane for each point
    # Distance is just the absolute difference in the specified axis
    point_axis_coords = points[:, :, axis_idx]  # (B, M)
    dist_to_plane = torch.abs(point_axis_coords - plane_pos)  # (B, M)
    
    # Calculate distance from body vectors (norm of the vectors)
    dist_to_body = torch.norm(body_dist_vec, dim=2)  # (B, M)
    
    # Create vectors from points to plane
    # The vector points from the point to its projection on the plane
    plane_vectors = torch.zeros_like(points)  # (B, M, 3)
    plane_vectors[:, :, axis_idx] = plane_pos - point_axis_coords  # Only non-zero in the plane normal direction
    
    # Create mask for points closer to plane than to body
    closer_to_plane = dist_to_plane < dist_to_body  # (B, M)
    
    # Expand mask to match vector dimensions
    closer_to_plane_expanded = closer_to_plane.unsqueeze(-1).expand_as(body_dist_vec)  # (B, M, 3)
    
    # Update vectors: use plane vectors where closer to plane, otherwise keep body vectors
    updated_vectors = torch.where(closer_to_plane_expanded, plane_vectors, body_dist_vec)

    return updated_vectors


def _compute_3d_grid(n):
    """Compute 3D grid dimensions for n tokens (same logic as StateRepresentation)."""
    cube_root = int(round(n ** (1/3)))
    for x in range(cube_root, 0, -1):
        if n % x == 0:
            remaining = n // x
            sqrt_rem = int(round(remaining ** 0.5))
            for y in range(sqrt_rem, 0, -1):
                if remaining % y == 0:
                    z = remaining // y
                    return (x, y, z)
    return (n, 1, 1)


def _adapt_state_tokens_if_needed(state_dict, model, logger):
    """
    Adapt state tokens in checkpoint if size differs from current model.
    Uses 3D trilinear interpolation to preserve learned spatial patterns.

    Args:
        state_dict: Checkpoint state dict
        model: Current model (may be wrapped in DDP)
        logger: Logger instance

    Returns:
        Adapted state_dict
    """
    import torch.nn.functional as F
    from torch.nn.parallel import DistributedDataParallel as DDP

    # Get the actual model (unwrap DDP if needed)
    model_ref = model.module if isinstance(model, DDP) else model

    # If model has no state_repr (no_global_state=true), remove state keys from checkpoint
    if model_ref.state_repr is None:
        state_keys_to_remove = [k for k in state_dict if 'state_repr.' in k]
        if state_keys_to_remove:
            for k in state_keys_to_remove:
                del state_dict[k]
            logger.info(f"[Checkpoint Adapt] no_global_state=true: removed {len(state_keys_to_remove)} state_repr keys from checkpoint")
        return state_dict

    # Find state token keys (handle both DDP and non-DDP checkpoints)
    state_key_patterns = [
        'state_repr.register_tokens.weight',
        'module.state_repr.register_tokens.weight'
    ]

    # Determine the correct key based on what's in the state_dict
    ckpt_key = None
    for pattern in state_key_patterns:
        if pattern in state_dict:
            ckpt_key = pattern
            break

    if ckpt_key is None:
        logger.info("[Checkpoint Adapt] No state tokens found in checkpoint, skipping adaptation")
        return state_dict

    # Get checkpoint and model state token shapes
    ckpt_tokens = state_dict[ckpt_key]
    ckpt_num_tokens, ckpt_dim = ckpt_tokens.shape

    # Get model's expected shape
    model_num_tokens = model_ref.state_repr.num_tokens
    model_dim = model_ref.state_repr.token_dim

    # Check if adaptation is needed
    if ckpt_num_tokens == model_num_tokens and ckpt_dim == model_dim:
        logger.info(f"[Checkpoint Adapt] State tokens match: ({ckpt_num_tokens}, {ckpt_dim}), no adaptation needed")
        return state_dict

    if ckpt_dim != model_dim:
        logger.warning(f"[Checkpoint Adapt] State token_dim mismatch: checkpoint={ckpt_dim}, model={model_dim}")
        logger.warning("[Checkpoint Adapt] Cannot adapt token_dim, state tokens will be randomly initialized")
        # Remove the mismatched key so it will be randomly initialized
        del state_dict[ckpt_key]
        return state_dict

    # Adapt num_tokens using 3D interpolation
    logger.info(f"[Checkpoint Adapt] State tokens size changed: {ckpt_num_tokens} -> {model_num_tokens}")

    # Compute 3D grid dimensions
    ckpt_grid = _compute_3d_grid(ckpt_num_tokens)
    model_grid = model_ref.state_repr.grid_3d if hasattr(model_ref.state_repr, 'grid_3d') else _compute_3d_grid(model_num_tokens)

    logger.info(f"[Checkpoint Adapt] Grid: {ckpt_grid} -> {model_grid}")

    # Reshape to 3D grid: (N, D) -> (1, D, X, Y, Z)
    ckpt_3d = ckpt_tokens.T.reshape(1, ckpt_dim, *ckpt_grid)  # (1, D, X, Y, Z)

    # Trilinear interpolation to new grid size
    adapted_3d = F.interpolate(
        ckpt_3d.float(),
        size=model_grid,
        mode='trilinear',
        align_corners=True
    )  # (1, D, X', Y', Z')

    # Reshape back to (N', D)
    adapted_tokens = adapted_3d.reshape(ckpt_dim, model_num_tokens).T  # (N', D)

    # Convert back to original dtype
    adapted_tokens = adapted_tokens.to(ckpt_tokens.dtype)

    # Update state_dict
    state_dict[ckpt_key] = adapted_tokens

    logger.info(f"[Checkpoint Adapt] State tokens adapted: {ckpt_tokens.shape} -> {adapted_tokens.shape}")

    return state_dict



def _object_centered(ref_pose: torch.Tensor, vmask: torch.Tensor) -> torch.Tensor:
    """ref_pose (B, Nobj, Nv, 3) minus each object's x/y centroid over valid vertices vmask (B, Nobj, Nv).

    Only x and y are centred: horizontal translation is a symmetry of the dynamics, height above
    the ground plane (z = 0) is not, so z is kept absolute."""
    m = vmask.to(ref_pose.dtype).unsqueeze(-1)
    c = (ref_pose * m).sum(dim=2, keepdim=True) / m.sum(dim=2, keepdim=True).clamp_min(1.0)
    c = torch.cat([c[..., :2], torch.zeros_like(c[..., 2:3])], dim=-1)
    return ref_pose - c


class CUT3RDynamicsModel(nn.Module):
    """
    Complete model combining encoder, state management, decoder, and predictor.
    """
    
    def __init__(self, config):
        super().__init__()

        # Store config for later use
        self.config = config

        # Set device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Physics-parameter and time-step flags (read before submodule construction)
        # Generalization arm: anchor reference poses expressed relative to each object's
        # frame-0 centroid (removes absolute translation from the predictor query), and
        # rotary pairing that matches rotate_half so attention depends on displacement only.
        self._ref_pose_object_centered = bool(config['model'].get('ref_pose_object_centered', False))
        from src.models.physformer_rope import PhysFormerRoPE as _RopeCls
        _RopeCls.INTERLEAVED_PAIRING = bool(config['model'].get('rope_interleaved_pairing', False))

        self._if_phys_param = config.get('model', {}).get('if_phys_param', False)
        self._if_time_step = config.get('model', {}).get('if_time_step', False)
        self._film_conditioning = config.get('model', {}).get('film_conditioning', False)
        if self._if_phys_param:
            print(f"[Model] Physics parameters enabled: mass, friction, restitution (+3D to encoder/predictor)")
        if self._if_time_step:
            if self._film_conditioning:
                print(f"[Model] Time step enabled via FiLM only (NOT concatenated to encoder input)")
            else:
                print(f"[Model] Time step enabled: skip_step_size (s and s², +2D to encoder/predictor)")

        # Rollout state reset: reset state tokens at each step to avoid train-test mismatch
        # When training with short sequences (e.g., 4 frames), state tokens only see 2 updates.
        # During long rollouts (e.g., 400 frames), state drift causes OOD behavior.
        # Setting this to True resets state tokens at each rollout step to match training distribution.
        self._rollout_state_reset = config.get('model', {}).get('rollout_state_reset', False)
        if self._rollout_state_reset:
            print(f"[Model] Rollout state reset ENABLED: state tokens will be reset at each rollout step")

        # KNN acceleration: GPU-parallel grouped processing
        self._knn_num_groups = config.get('data', {}).get('knn_num_groups', 1)
        # Valid-vertex packing (model.knn_packed, model.encoder.packed): the distance features and the encoder's
        # pointwise layers run over the valid vertices only instead of every padded (object slot x vertex) entry.
        # Equal to the padded paths up to floating-point rounding; the packs are built once per batch / rollout.
        self._knn_packed = bool(config.get('model', {}).get('knn_packed', False))
        self._encoder_packed = bool(config.get('model', {}).get('encoder', {}).get('packed', False))
        if self._knn_packed or self._encoder_packed:
            print(f"[Model] valid-vertex packing: knn_packed={self._knn_packed}, encoder.packed={self._encoder_packed}")
        if self._knn_packed and self._knn_num_groups > 1:
            raise ValueError('model.knn_packed replaces data.knn_num_groups; set knn_num_groups to 1')
        # model.encoder.packed_pooling: 'segment' pools the packed rows directly (segment max / mean) instead of
        # scattering every level back to the padded layout; same values up to fp32 summation order in the mean.
        if (config.get('model', {}).get('encoder', {}).get('packed_pooling', 'scatter') != 'scatter'
                and not self._encoder_packed):
            raise ValueError('model.encoder.packed_pooling applies to the packed encoder path; set model.encoder.packed')
        if self._knn_num_groups > 1:
            print(f"[Model] KNN acceleration enabled: num_groups={self._knn_num_groups}")

        # model.kabsch_reference: the mesh the per-step Kabsch fit transforms. 'advancing' (default) = the
        # previous step's result, so every step re-transforms an already transformed mesh and rounding accumulates;
        # 'fixed' = one unchanged reference shape for the whole sequence / rollout (the first current frame), equal
        # in exact arithmetic (the fit is equivariant) without the accumulation.
        _kref = config.get('model', {}).get('kabsch_reference', 'advancing')
        if _kref not in ('fixed', 'advancing'):
            raise ValueError(f"model.kabsch_reference must be 'fixed' or 'advancing', got {_kref!r}")
        self._kabsch_reference_fixed = _kref == 'fixed'
        if self._kabsch_reference_fixed:
            print("[Model] Kabsch reference FIXED: every step fits the same reference mesh")

        # model.intersection_handling.after_kabsch_loss: which state the position-after-Kabsch loss supervises.
        # 'rigid' = the Kabsch output (the state the acceleration-after-Kabsch loss already uses), 'handled'
        # (default) = the state after the ground/object correction. The rollout state is the corrected one either way.
        _akl = config.get('model', {}).get('intersection_handling', {}).get('after_kabsch_loss', 'handled')
        if _akl not in ('rigid', 'handled'):
            raise ValueError(f"model.intersection_handling.after_kabsch_loss must be 'rigid' or 'handled', got {_akl!r}")
        self._after_kabsch_loss_rigid = _akl == 'rigid'
        if self._after_kabsch_loss_rigid:
            print("[Model] position-after-Kabsch loss supervises the Kabsch output (before intersection handling)")

        # Time-batched training forward (training.time_batched_forward): the network part of all teacher-forced
        # steps runs as one batch (see _forward_multi_object). Only valid when no step's network input depends on
        # an earlier step's prediction or state; the settings that would create such a dependency are refused.
        self._time_batched_forward = bool(config.get('training', {}).get('time_batched_forward', False))
        if self._time_batched_forward:
            _enc = config.get('model', {}).get('encoder', {})
            _bad = []
            if config.get('model', {}).get('if_penetration', config.get('model', {}).get('if_directional_distance', False)):
                _bad.append('model.if_penetration/if_directional_distance (uses the previous step\'s decoder output)')
            if not _enc.get('masked_pooling', False):
                _bad.append('model.encoder.masked_pooling=false (BatchNorm statistics would pool across steps)')
            if _enc.get('use_stn', False) or _enc.get('feature_transform', False):
                _bad.append('encoder STN / feature transform (their BatchNorms would pool across steps)')
            if _bad:
                raise ValueError("training.time_batched_forward is incompatible with: " + '; '.join(_bad))
            print("[Model] time-batched training forward ACTIVE (network part of all teacher-forced steps in one "
                  "batch; per-step BatchNorm statistics)", flush=True)

        # Debug mode: from data.random_step_size.if_debug or training.loss_weights.if_debug
        self._debug_mode = (
            config.get('data', {}).get('random_step_size', {}).get('if_debug', False) or
            config.get('training', {}).get('loss_weights', {}).get('if_debug', False)
        )

        # Config Consistency Check: Multi-timescale training requires time_step feature
        random_step_size_config = config.get('data', {}).get('random_step_size', {})
        if random_step_size_config.get('enable', False):
            if not self._if_time_step:
                raise ValueError(
                    "Config inconsistency detected!\n"
                    "  data.random_step_size.enable = True  (multi-timescale training enabled)\n"
                    "  model.if_time_step = False  (time step feature disabled)\n"
                    "\n"
                    "Multi-timescale training requires the network to receive step_size as input.\n"
                    "Please set model.if_time_step = true in your config file."
                )
            print(f"[Model] Multi-timescale training enabled: step_size will be used as condition feature")

        # Initialize components (multi-object only)

        # ============================================================
        # ARCHITECTURAL INVARIANTS — Translation Consistency & Permutation Equivariance
        # These properties are fundamental to the model. Enforce at init time.
        # ============================================================

        # 1. State tokens permanently removed — fixed 3D grid breaks translation consistency
        self._no_global_state = True
        self.state_repr = None
        assert config['model'].get('no_global_state', True) is True, \
            "no_global_state must be True: state tokens on a fixed 3D grid break translation consistency"

        # 2. rope_set_aggregation must be True — concat mode is anchor-order sensitive
        assert config['model']['decoder'].get('rope_set_aggregation', False) is True, \
            "rope_set_aggregation must be True: the concatenation mode is anchor-order sensitive, " \
            "breaking permutation invariance of ARoPE within each object"

        # 3. No sequence-index positional embeddings — would break object permutation equivariance
        assert not config['model']['decoder'].get('use_learnable_pe', False), \
            "use_learnable_pe must be False: sequence-index PE breaks object permutation equivariance"

        # ============================================================

        # Multi-scale predictor reads features from selected decoder layers (CUT3R-style)
        from src.models.multi_scale_predictor import QueryBasedMultiScalePredictor

        # Adjust layer indices based on actual number of decoder layers (needed before decoder init)
        num_decoder_layers = config['model']['decoder'].get('num_layers', 12)
        if num_decoder_layers >= 12:
            layer_indices = [0, 6, 9, 12]  # CUT3R's layers for 12-layer model
        elif num_decoder_layers >= 6:
            layer_indices = [0, 2, 4, num_decoder_layers]  # For 6+ layer model
        elif num_decoder_layers >= 4:
            # For 4-layer model
            layer_indices = [0, 1, 2, num_decoder_layers]  # Use all available layers
        else:
            # For 2-layer model (tiny), use all available layers
            layer_indices = list(range(num_decoder_layers + 1))  # [0, 1, 2] for 2 layers
        # model.predictor.layer_indices overrides the rule above, e.g. [0, 1, 2, 8] for a decoder grown from 4 to 8
        # layers (scripts/grow_decoder.py) that reads the same layers as the 4-layer model
        if config['model']['predictor'].get('layer_indices') is not None:
            layer_indices = [int(i) for i in config['model']['predictor']['layer_indices']]
            if sorted(set(layer_indices)) != layer_indices or layer_indices[-1] != num_decoder_layers:
                raise ValueError(f"model.predictor.layer_indices must be increasing and end with the last decoder layer "
                                 f"({num_decoder_layers}); got {layer_indices}")

        # Verify that all layer indices are valid for the decoder
        max_valid_index = num_decoder_layers  # 0 to num_decoder_layers inclusive (initial + N layers)
        for idx in layer_indices:
            assert 0 <= idx <= max_valid_index, \
                f"Invalid layer index {idx} for decoder with {num_decoder_layers} layers. " \
                f"Valid range: [0, {max_valid_index}]"

        # Store layer indices for decoder to only save needed layers
        self.predictor_layer_indices = set(layer_indices)  # Use set for O(1) lookup

        # Create mapping from original layer index to position in saved list
        # For example: if layer_indices = [0, 1, 2, 4], then:
        # layer_index_mapping = {0: 0, 1: 1, 2: 2, 4: 3}
        self.layer_index_mapping = {idx: i for i, idx in enumerate(sorted(layer_indices))}

        # Gated Attention configuration (Qwen3-style)
        gated_attention_config = config['model']['decoder'].get('gated_attention', {})
        gated_attention_enable = gated_attention_config.get('enable', False)
        gated_attention_type = gated_attention_config.get('type', 'elementwise')
        gated_attention_debug = gated_attention_config.get('if_debug', False)
        # Currently only support elementwise gate
        if gated_attention_enable and gated_attention_type != 'elementwise':
            raise ValueError(f"Unsupported gated_attention type: {gated_attention_type}. Only 'elementwise' is supported.")
        self._elementwise_gate = gated_attention_enable and gated_attention_type == 'elementwise'

        # Debug print for gated attention (print to console since FileHandler not yet set)
        if gated_attention_debug:
            msg = (f"[Gated Attention] enable={gated_attention_enable}, type={gated_attention_type}, "
                   f"elementwise_gate={self._elementwise_gate}")
            print(msg)  # Console output (FileHandler not ready at model init time)
            logger.info(msg)  # Also log (will appear in file after Trainer sets FileHandler)

        # State update mode: when False, state is just learnable embeddings (no bidirectional update)
        # Use with rollout_state_reset=True to avoid train-test mismatch
        self._state_update_enabled = config['model']['decoder'].get('state_update_enabled', True)

        # Objects Self-Attention: Enable direct object-to-object interaction
        # Useful for physics simulation where objects need to exchange information
        self._objects_self_attention = config['model']['decoder'].get('objects_self_attention', False)

        # PE type must be read before decoder init (use_se3 flag needs it)
        self._pe_type = config['model']['decoder'].get('pe_type', 'arope')

        # Initialize decoder with layer_indices_to_save for memory efficiency
        # Block AttnRes config (Kimi 2026)
        _block_attn_res_cfg = config['model']['decoder'].get('block_attn_res', {})
        _block_attn_res_enable = _block_attn_res_cfg.get('enable', False) if isinstance(_block_attn_res_cfg, dict) else bool(_block_attn_res_cfg)
        _block_attn_res_block_size = _block_attn_res_cfg.get('block_size', 4) if isinstance(_block_attn_res_cfg, dict) else 4
        _block_attn_res_final = bool(_block_attn_res_cfg.get('final_aggregate', False)) if isinstance(_block_attn_res_cfg, dict) else False

        self.decoder = BidirectionalDecoder(
            state_dim=config['model']['state'].get('token_dim', 768),
            num_layers=config['model']['decoder'].get('num_layers', 12),  # CUT3R uses 12
            num_heads=config['model']['decoder'].get('num_heads', 12),
            feedforward_dim=config['model']['decoder'].get('feedforward_dim', 3072),
            dropout=config['model']['decoder']['dropout'],
            return_all_layers=True,  # For multi-scale fusion
            use_self_attention=config['model']['decoder'].get('use_self_attention', False),  # Optional state SA
            center_rope_positions=config['model'].get('rope_center_positions', False),
            rope_set_aggregation=config['model']['decoder'].get('rope_set_aggregation', False),  # Set aggregation for RoPE
            rope_dim=config['model']['decoder'].get('rope_dim', None),  # Custom RoPE dimension (None = auto-compute)
            dim_per_coord_3d=config['model']['decoder'].get('dim_per_coord_3d', None),  # Dim per coordinate for single-point (3D) RoPE
            dim_per_coord_multi=config['model']['decoder'].get('dim_per_coord_multi', None),  # Dim per coordinate for multi-point RoPE
            layer_indices_to_save=self.predictor_layer_indices,  # Only save layers needed by predictor
            elementwise_gate=self._elementwise_gate,  # Qwen3: elementwise gated attention
            state_update_enabled=self._state_update_enabled,  # When False, state is learnable embeddings only
            objects_self_attention=self._objects_self_attention,  # Enable Objects Self-Attention for object-object interaction
            use_se3=(self._pe_type == 'se3'),  # SE(3) pairwise bias for obj self-attention
            block_attn_res=_block_attn_res_enable,  # Block Attention Residuals (Kimi 2026)
            block_attn_res_block_size=_block_attn_res_block_size,  # Layers per block
            block_attn_res_final_aggregate=_block_attn_res_final,  # final depth aggregation as the last features
            film_conditioning=(config['model']['decoder'].get('film_conditioning', False) and self._film_conditioning and self._if_time_step),  # Per-layer FiLM: requires decoder.film_conditioning + model.film_conditioning + model.if_time_step
        )

        # ============ Anchor-Vertex Pooling (AVP) ============
        avp_cfg = config.get('model', {}).get('avp', {})
        self._avp_enabled = avp_cfg.get('enable', False)
        self._avp_output_dim = avp_cfg.get('output_dim', 256)
        if self._avp_enabled:
            from src.models.anchor_vertex_pooling import AnchorVertexPooling
            _avp_vertex_dim = config.get('model', {}).get('encoder', {}).get(
                'backbone_dim', config.get('model', {}).get('encoder', {}).get('output_dim', 1024))
            self.avp = AnchorVertexPooling(
                vertex_dim=_avp_vertex_dim,
                output_dim=self._avp_output_dim,
                init_sigma=avp_cfg.get('init_sigma', 0.1)
            )
            print(f"[Model] AVP ENABLED: vertex_dim={_avp_vertex_dim}, output_dim={self._avp_output_dim}, "
                  f"init_sigma={avp_cfg.get('init_sigma', 0.1)}")
        else:
            self.avp = None

        # Dynamic predictor input_dim calculation:
        # Base: 12D (9D query: dist_vec+vel+rel_pos + 3D ref_pose)
        # + phys_params (if if_phys_param=true): +3D = 15D
        # + time_step (if if_time_step=true AND NOT film): +2D = 17D
        # + AVP (if avp.enable=true): +avp_output_dim
        predictor_input_dim = 12  # Base: 9D query + 3D ref_pose
        if self._if_phys_param:
            predictor_input_dim += 3  # Physics params
        if self._if_time_step and not self._film_conditioning:
            predictor_input_dim += 2  # Time step (s and s²) — skipped when FiLM handles it
        if self._avp_enabled:
            predictor_input_dim += self._avp_output_dim  # AVP features per anchor
            print(f"[Model] Predictor input_dim: {predictor_input_dim} "
                  f"(base 12 + phys {3 if self._if_phys_param else 0} + AVP {self._avp_output_dim})")

        self.predictor = QueryBasedMultiScalePredictor(
            state_dim=config['model']['state'].get('token_dim', 768),
            query_dim=predictor_input_dim,  # Dynamically computed based on flags
            hidden_dim=config['model']['predictor'].get('hidden_dims', [512, 256])[0],
            output_dim=config['model']['predictor']['output_dim'],
            layer_indices=layer_indices,
            layer_index_mapping=self.layer_index_mapping,  # Pass the mapping for memory-optimized decoder
            num_heads=config['model']['predictor'].get('num_heads', config['model']['decoder'].get('num_heads', 8)),  # Use predictor's num_heads if specified
            dropout=config['model']['predictor']['dropout'],
            use_ref_pose=config['model']['predictor'].get('use_ref_pose', True),  # Enable ref pose
            norm_type=config['model']['predictor'].get('norm_type', 'layer_norm'),
            qk_norm=config['model']['predictor'].get('qk_norm', False),
            fusion_activation='swiglu',
            use_rope=config['model']['predictor'].get('use_rope', True),  # Enable RoPE for position-aware attention
            rope_dim=config['model']['decoder'].get('rope_dim', None),  # Use same RoPE dim as decoder
            dim_per_coord_3d=config['model']['decoder'].get('dim_per_coord_3d', None),  # Single-Point RoPE config
            dim_per_coord_multi=config['model']['decoder'].get('dim_per_coord_multi', None),  # Multi-Point RoPE config
            elementwise_gate=self._elementwise_gate,  # Pass elementwise_gate config
            cross_object_attention=config['model']['predictor'].get('cross_object_attention', False),  # Cross-object anchor attention
            cross_object_kv=config['model']['predictor'].get('cross_object_kv', False),  # Cross-object key-value attention
            no_global_state=config['model'].get('no_global_state', False)  # Removes unused state input channels
        )
        # model.predictor.rmsnorm_eps: training_dtype -> the predictor's RMSNorm modules keep the eps they have in
        # bf16-autocast training also when evaluated in fp32 (their default eps follows the input dtype; see
        # pin_rmsnorm_eps). Applied here, so every construction of the model computes the configured function.
        _rms_eps_mode = config['model']['predictor'].get('rmsnorm_eps')
        if _rms_eps_mode == 'training_dtype':
            _n_pinned = self.predictor.pin_rmsnorm_eps()
            logger.info(f"[Predictor] RMSNorm eps pinned to the bf16-training values on {_n_pinned} modules")
        elif _rms_eps_mode not in (None, 'default'):
            raise ValueError("model.predictor.rmsnorm_eps must be 'training_dtype' or 'default'")
        # model.predictor.fp32_output: the acceleration read-out layer in fp32 under bf16 training autocast (evaluation
        # runs it in fp32 already, so evaluation is unchanged)
        self.predictor.fp32_output = bool(config['model']['predictor'].get('fp32_output', False))
        if self.predictor.fp32_output:
            logger.info("[Predictor] acceleration read-out (fusion_out) in fp32 during training")

        # Store attention flags for reference
        self._cross_object_attention = config['model']['predictor'].get('cross_object_attention', False)
        self._cross_object_kv = config['model']['predictor'].get('cross_object_kv', False)

        # Log attention configuration summary
        logger.info("=" * 60)
        logger.info("[Attention Config Summary]")
        logger.info(f"  Decoder Objects Self-Attention: {'ENABLED' if self._objects_self_attention else 'DISABLED'}")
        logger.info(f"  Predictor Cross-Object Anchor Attention: {'ENABLED' if self._cross_object_attention else 'DISABLED'}")
        logger.info(f"  Predictor Cross-Object Key-Value: {'ENABLED' if self._cross_object_kv else 'DISABLED'}")
        logger.info("=" * 60)

        # Store actual dt for video/visualization
        self.dt = config['data'].get('dt', 0.01)  # Actual timestep (0.01 = 100 FPS)
        # Verlet integration always uses dt=1; the network learns accelerations in this unit system.
        self.dt_verlet = 1.0
        


        # Create transform layer (configurable)
        transform_cfg = config.get('model', {}).get('transform', {})
        transform_method = transform_cfg.get('method', 'kabsch')
        self.learnable_transform = (transform_method == 'learnable')

        # RoMa differentiable Kabsch (if enabled, no STE needed)
        self.if_roma = bool(transform_cfg.get('if_roma', False))
        if self.if_roma and not ROMA_AVAILABLE:
            raise ImportError(
                "RoMa is required when if_roma=true. Install with: pip install roma"
            )

        # Build kwargs per method to avoid passing unsupported args
        if self.learnable_transform:
            tl_kwargs = {
                'use_6d_rotation': transform_cfg.get('use_6d_rotation', True),
                'learnable_mode': transform_cfg.get('learnable_mode', 'absolute')
            }
        else:
            tl_kwargs = {
                'regularize_rotation': transform_cfg.get('regularize_rotation', True)
            }
        self.transform_layer = create_transform_layer(
            method=transform_method,
            if_roma=self.if_roma,
            **tl_kwargs
        )
        
        # ============== Multi-object (MOVi) specific modules ==============
        self.enable_multi_object = bool(self.config.get('data', {}).get('multi_object', False))
        if self.enable_multi_object:
            # Cache for mesh faces (dataset_name:obj_idx -> faces array)
            self._faces_cache = {}

            # Shared PointNet encoder for arbitrary number of objects (weights shared)
            enc_cfg = config.get('model', {}).get('encoder', {})
            # Dynamic encoder input_dim calculation:
            # Base: 9D (dist_vec + velocity + rel_pos)
            # + phys_params (if if_phys_param=true): +3D
            # + time_step (if if_time_step=true and no FiLM): +2D (s and s²)
            encoder_input_dim = 9  # Base
            if self._if_phys_param:
                encoder_input_dim += 3  # Physics params
            if self._if_time_step and not self._film_conditioning:
                encoder_input_dim += 2  # Time step (s and s²) — skipped when FiLM handles it
            if enc_cfg.get('height_input', False):
                encoder_input_dim += 1  # Height above the ground plane (absolute z)
            self._encoder_input_dim = encoder_input_dim  # Save for logging

            self.shared_obj_encoder = RigidBodyEncoder(
                input_dim=encoder_input_dim,
                hidden_dims=enc_cfg.get('hidden_dims', [512, 1024]),
                output_dim=enc_cfg.get('backbone_dim', enc_cfg.get('output_dim', 896)),
                backbone_dim=enc_cfg.get('backbone_dim', enc_cfg.get('output_dim', 896)),
                use_stn=enc_cfg.get('use_stn', False),
                use_batch_norm=enc_cfg.get('use_batch_norm', True),
                masked_pooling=enc_cfg.get('masked_pooling', False),
                freeze_bn_stats=enc_cfg.get('freeze_bn_stats', False),
                height_input=enc_cfg.get('height_input', False),
                hierarchical_subsets=enc_cfg.get('hierarchical_subsets', False),
                packed_pooling=enc_cfg.get('packed_pooling', 'scatter'),
                subset_small_objects=enc_cfg.get('subset_small_objects', 'all_below_8'),
                feature_transform=enc_cfg.get('feature_transform', False),
                feature_transform_dim=enc_cfg.get('feature_transform_dim', 64),
                pooling_type=enc_cfg.get('pooling_type', 'max'),  # New: enhanced pooling
                num_hierarchical_levels=enc_cfg.get('num_hierarchical_levels', 2)  # New: hierarchical levels
            )
            # Project per-object latent to model token dim
            token_dim = config['model']['state'].get('token_dim', 768)
            enc_out_dim = enc_cfg.get('backbone_dim', enc_cfg.get('output_dim', 896))
            # Project PointNet encoder output (per-object) to token dimension for Decoder
            self.obj_feature_proj = nn.Linear(enc_out_dim, token_dim)
            # Anchor query dimension is always 9D (dist_vec + vel + rel_pos)
            anchor_query_dim = 9
            # Projection for anchor query features (for Predictor)
            self.anchor_query_proj = nn.Linear(anchor_query_dim, token_dim)

            # FiLM conditioning: modulate encoder output with [s, s²] instead of
            # concatenating step size to the input features
            if self._film_conditioning and self._if_time_step:
                from src.models.encoder import StepSizeFiLM
                self.encoder_film = StepSizeFiLM(time_dim=2, feature_dim=token_dim)
                print(f"[Model] FiLM conditioning ENABLED: step-size modulates encoder output (dim={token_dim})")
                print(f"[Model]   Encoder input: {encoder_input_dim}D (time NOT concatenated, handled by FiLM)")
                print(f"[Model]   Predictor query: {predictor_input_dim}D (time NOT in query)")
            else:
                self.encoder_film = None
                if self._film_conditioning and not self._if_time_step:
                    print(f"[Model] WARNING: film_conditioning=True but if_time_step=False; FiLM disabled")

            # ============== PE Type for ARoPE Ablation ==============
            # pe_type controls what 12D positional descriptor is fed to the decoder's RoPE path.
            # "arope" (default): original 4 anchors × 3D = 12D
            # "obb": OBB center + 3 half-axis vectors = 12D
            # "pca": PCA centroid + √λ·eigenvectors = 12D
            self._pe_type = config['model']['decoder'].get('pe_type', 'arope')
            _VALID_PE_TYPES = ('arope', 'obb', 'pca', 'se3')
            if self._pe_type not in _VALID_PE_TYPES:
                raise ValueError(
                    f"Invalid pe_type='{self._pe_type}'. Must be one of {_VALID_PE_TYPES}."
                )
            _pe_msg = f"[PE Type] Using pe_type='{self._pe_type}' for decoder/predictor positional encoding."
            print(_pe_msg)
            logger.info(_pe_msg)

            self._geometry_encoder = None
            if self._pe_type in ('obb', 'pca'):
                from src.models.geometry_encoders import OBBEncoder, PCAFrameEncoder
                if self._pe_type == 'obb':
                    self._geometry_encoder = OBBEncoder()
                    logger.info("[PE Type] Initialized OBBEncoder (center + 3 half-axis vectors -> 12D)")
                elif self._pe_type == 'pca':
                    self._geometry_encoder = PCAFrameEncoder()
                    logger.info("[PE Type] Initialized PCAFrameEncoder (centroid + scaled eigenvectors -> 12D)")

            # SE(3) pairwise bias module
            self._se3_bias_module = None
            if self._pe_type == 'se3':
                from src.models.se3_attention import SE3PairwiseBias
                _num_heads = config['model']['decoder'].get('num_heads', 6)
                _se3_hidden = config['model']['decoder'].get('se3_hidden_dim', 64)
                self._se3_bias_module = SE3PairwiseBias(
                    num_heads=_num_heads, hidden_dim=_se3_hidden, use_scale_ratio=True
                )
                logger.info(f"[PE Type] Initialized SE3PairwiseBias (7D -> {_se3_hidden} -> {_num_heads} heads, zero-init, scale-norm)")

            # ============== RoPE frequency configuration summary ==============
            _rope_dim = config['model']['decoder'].get('rope_dim', None)
            _dim_per_coord_3d = config['model']['decoder'].get('dim_per_coord_3d', None)
            _dim_per_coord_multi = config['model']['decoder'].get('dim_per_coord_multi', None)
            _hidden_dim = config['model']['decoder'].get('hidden_dim', 768)
            _num_heads = config['model']['decoder'].get('num_heads', 6)
            _head_dim = _hidden_dim // _num_heads if _num_heads > 0 else 0
            _anchor_k = config.get('data', {}).get('anchor_k', 4)
            _n_coords = _anchor_k * 3
            _rope_msg = (
                f"[RoPE Freq] rope_dim={_rope_dim} | "
                f"dim_per_coord_multi={_dim_per_coord_multi} (freq/coord) | "
                f"dim_per_coord_3d={_dim_per_coord_3d} | "
                f"hidden_dim={_hidden_dim} | num_heads={_num_heads} | head_dim={_head_dim} | "
                f"anchor_k={_anchor_k} | n_coords={_n_coords} | "
                f"check: {_n_coords}x{_dim_per_coord_multi}={_n_coords * _dim_per_coord_multi if _dim_per_coord_multi else '?'} "
                f"{'== ' + str(_rope_dim) if _dim_per_coord_multi and _n_coords * _dim_per_coord_multi == _rope_dim else '!= ' + str(_rope_dim) + ' (MISMATCH)'}"
            )
            print(_rope_msg)
            logger.info(_rope_msg)

    def forward(self, batch, return_all=False, teacher_forcing_ratio=1.0):
        """
        Forward pass through the model.

        Must remain consistent with rollout():
        1. Reference pose is frame 0 of the full sequence (not batch[:, 0])
        2. Input: current position only (no velocity to encoder)
        3. Output: acceleration, integrated via Verlet to get the next position
        4. Predictor maps (current_pos, ref_pose_frame0) -> acceleration
        5. Verlet: x_{t+1} = a*dt² + 2*x_t - x_{t-1}

        Args:
            batch: Dictionary with body positions and velocities
            return_all: Whether to return intermediate results
            teacher_forcing_ratio: Probability of using ground truth (1.0 = always teacher, 0.0 = always student)
            
        Returns:
            predictions: Dictionary with predicted positions, velocities, etc.
        """
        # Multi-object path only
        if 'obj_positions' in batch and getattr(self, 'enable_multi_object', False):
            return self._forward_multi_object(batch, return_all=return_all, teacher_forcing_ratio=teacher_forcing_ratio)
        else:
            raise ValueError("Expected multi-object batch with 'obj_positions'.")

    # ================= Multi-object forward (MOVi) =================
    def _forward_multi_object(self, batch, return_all=False, teacher_forcing_ratio=1.0):
        """
        Full multi-object forward path using shared per-object PointNet encoder and
        mask-aware per-vertex queries across all objects.
        """
        time_step_id = batch['time_step_id']  # (B,)
        # network inputs: the rigid-noise augmentation's perturbed frames when present (the losses keep comparing
        # against batch['obj_positions'])
        obj_pos = batch.get('obj_positions_input', batch['obj_positions'])  # (B,T,Nobj,Nv,3)
        obj_mask = batch['valid_obj_mask']  # (B,Nobj)
        vtx_mask = batch['valid_vtx_mask']  # (B,T,Nobj,Nv)
        ref_obj_positions = batch['ref_obj_positions']  # (B?,Nobj,Nv,3) or (Nobj,Nv,3)

        # Read physics parameters and time step if enabled
        phys_params = None  # (B, Nobj, 3): [mass, friction, restitution]
        skip_step_size_batch = None  # (B,)
        if self._if_phys_param:
            if 'phys_params' not in batch:
                raise KeyError("Batch is missing 'phys_params' required when if_phys_param=true")
            phys_params = batch['phys_params'].to(obj_pos.device)  # (B, Nobj, 3)
        if self._if_time_step:
            if 'skip_step_size' not in batch:
                raise KeyError("Batch is missing 'skip_step_size' required when if_time_step=true")
            skip_step_size_batch = batch['skip_step_size'].to(obj_pos.device)  # (B,)

        B, T, Nobj, Nv, _ = obj_pos.shape
        device = obj_pos.device

        # Normalize reference pose to (B,Nobj,Nv,3)
        if ref_obj_positions.dim() == 3:
            ref_pose = ref_obj_positions.unsqueeze(0).expand(B, -1, -1, -1).to(device)
        else:
            ref_pose = ref_obj_positions.to(device)  # (B,Nobj,Nv,3)

        # ========== Center scene by first frame ==========
        # If enabled, subtract the global center of the first frame from all coordinates;
        # relative distances and object shapes are unaffected.
        if self.config.get('data', {}).get('center_scene_by_t0', False):
            for b in range(B):
                vmask_b0 = vtx_mask[b, 0]  # (Nobj, Nv)
                v_t0 = obj_pos[b, 0][vmask_b0]  # Valid vertices at t=0
                if v_t0.shape[0] > 0:
                    center_t0 = v_t0.mean(dim=0)  # (3,)
                    
                    # Log centering (only once)
                    if not hasattr(self, '_center_t0_logged_train'):
                        logger.info(f"[Center-T0 TRAIN] Enabled. Sample b={b}: center_t0={center_t0.cpu().numpy().tolist()}, "
                                   f"obj_pos range before: [{obj_pos[b].min().item():.4f}, {obj_pos[b].max().item():.4f}]")
                    
                    # Shift positions
                    obj_pos[b] = obj_pos[b] - center_t0
                    ref_pose[b] = ref_pose[b] - center_t0
                    
                    if not hasattr(self, '_center_t0_logged_train'):
                        logger.info(f"[Center-T0 TRAIN] After centering b={b}: "
                                   f"obj_pos range: [{obj_pos[b].min().item():.4f}, {obj_pos[b].max().item():.4f}], "
                                   f"ref_pose range: [{ref_pose[b].min().item():.4f}, {ref_pose[b].max().item():.4f}]")
            if not hasattr(self, '_center_t0_logged_train'):
                self._center_t0_logged_train = True
                logger.info(f"[Center-T0 TRAIN] Scene centering applied to obj_pos and ref_pose for batch of {B} samples")

        # Require anchors for query and supervision (no fallback)
        if 'anchor_idx' not in batch:
            raise KeyError("Batch is missing 'anchor_idx' required for anchor-based querying")
        anchor_idx = batch['anchor_idx'].to(device).long()  # (B,Nobj,Nq)
        if anchor_idx.dim() != 3 or anchor_idx.shape[0] != B or anchor_idx.shape[1] != Nobj:
            raise ValueError(f"anchor_idx must be (B,Nobj,Nq); got {tuple(anchor_idx.shape)} vs (B={B},Nobj={Nobj})")
        Nq = int(anchor_idx.shape[2])

        # Initialize persistent state tokens and 3D positions for RoPE
        # When no_global_state=True, pass None to decoder (skips all state cross-attention)
        if self._no_global_state:
            state_tokens = None
            state_pos = None
        else:
            state_tokens, state_pos = self.state_repr.get_state_and_positions(B, device)

        # Storage across time steps
        all_pred_pos_for_rollout = []  # STE positions for autoregressive input (B, Nobj, Nv, 3)
        all_pred_pos_raw = []          # Raw positions for loss supervision (B, Nobj, Nv, 3)
        all_pred_pos_raw_norm = []     # Raw normalized positions for loss (before Kabsch)
        all_pred_pos_icp_norm = []     # ICP/Kabsch-transformed normalized positions for loss (after Kabsch)
        all_accel_no_icp_norm_anchor_seq = []  # (B, Nobj, Nq, 3)
        all_accel_icp_norm_seq = []
        all_pen_pred_vertices_seq = []  # (B, Nobj, Nv, 1) per step if enabled (ALL vertices)

        # Registration refs per object (initialized with GT at the first step)
        registration_refs = None

        # Previous step's features for penetration prediction (t>1)
        prev_obj_layers = None
        prev_state_layers = None

        # ========== ANCHOR VALIDATION (whole window, vectorized, one host sync) ==========
        # Every anchor of a valid object must sit on a valid vertex at every frame the loop reads
        # (vtx_mask[:, t] for t in 1..T-2). Checked once per batch instead of per (t, object).
        if T > 2:
            if anchor_idx.dim() != 3 or anchor_idx.shape[0] != B or anchor_idx.shape[2] != Nq:
                raise RuntimeError(
                    f"anchor_idx shape mismatch: got {tuple(anchor_idx.shape)}, expected (B={B},Nobj,Nq={Nq})")
            _vmask_win = vtx_mask[:, 1:T - 1]                                # (B,Tw,Nobj,Nv)
            _numv_win = _vmask_win.sum(dim=3)                                # (B,Tw,Nobj)
            _idx_win = anchor_idx.unsqueeze(1).expand(-1, T - 2, -1, -1)     # (B,Tw,Nobj,Nq)
            _at_anchor = torch.gather(_vmask_win, 3, _idx_win)               # (B,Tw,Nobj,Nq)
            _ovalid = obj_mask.unsqueeze(1)                                  # (B,1,Nobj)
            _zero_bad = _ovalid & (_numv_win == 0)
            _short_bad = _ovalid & (_numv_win < Nq)
            _anch_bad = (~_at_anchor) & _ovalid.unsqueeze(-1)
            _flags = torch.stack([_zero_bad.any(), _short_bad.any(), _anch_bad.any()])
            if bool(_flags.any().item()):  # the single sync; everything below is the error path
                zb, sb, ab = [bool(x.item()) for x in _flags]
                _where = (_zero_bad if zb else (_short_bad if sb else _anch_bad.any(dim=-1))).nonzero(as_tuple=False)
                _b0, _tw0, _i0 = [int(v) for v in _where[0][:3]]
                _nv0 = int(_numv_win[_b0, _tw0, _i0].item())
                ds_name = batch.get('dataset_name', None)
                ds_name_str = ds_name if isinstance(ds_name, str) else (ds_name[0] if hasattr(ds_name, '__getitem__') else 'unknown')
                logger.error(
                    f"[Anchor Debug] dataset={ds_name_str} t={_tw0 + 1} obj={_i0} sample={_b0} Nv={Nv} anchor_k={Nq} "
                    f"valid_vertices={_nv0} zero_any={zb} short_any={sb} bad_anchor_any={ab}")
                if zb:
                    raise RuntimeError(f"Zero valid vertices encountered at t={_tw0 + 1}, obj={_i0} for some valid samples; dataset likely corrupted.")
                if sb:
                    raise RuntimeError(f"Not enough valid vertices for anchors at t={_tw0 + 1}, obj={_i0}: valid={_nv0} < anchor_k={Nq}")
                raise RuntimeError(
                    f"Invalid anchors: anchor indices refer to invalid/padded vertices at t={_tw0 + 1}, obj={_i0}. "
                    f"Consider reducing anchor_k or ensuring consistent Nv per frame.")

        # Object slots that hold a valid object in at least one sample (fixed for the batch). The per-object
        # KNN skips the others: they are padding everywhere and their distance vectors are zeroed anyway.
        _knn_active_slots = (torch.nonzero(obj_mask.any(dim=0)).flatten().tolist()
                             if self._knn_num_groups <= 1 else None)
        # Constant scalars of the decoder attention mask and the valid-object index set used by the Kabsch
        # fit: both are fixed for the batch, so they are built once instead of at every step (the per-step
        # versions cost two host-to-device copies and one nonzero sync each step).
        _attn_zero_fwd = torch.tensor(0.0, dtype=torch.float32, device=obj_mask.device)
        _attn_ninf_fwd = torch.tensor(float('-inf'), dtype=torch.float32, device=obj_mask.device)
        _kabsch_valid_idx_fwd = obj_mask.reshape(-1).nonzero(as_tuple=True)[0]
        # Valid-vertex packs (model.knn_packed / model.encoder.packed), fixed for the batch: every frame of a
        # sequence has the same vertex mask, which is checked once here.
        _knn_pack_fwd = _enc_pack_fwd = None
        if self._knn_packed or self._encoder_packed:
            _vm1 = vtx_mask[:, 1]
            if not bool((vtx_mask[:, 1:T - 1] == _vm1.unsqueeze(1)).all()):
                raise ValueError('valid-vertex packing needs the same vertex mask in every frame of a sequence')
            if self._knn_packed:
                _knn_pack_fwd = knn_pack(_vm1, obj_mask)
            if self._encoder_packed:
                _enc_m = self.shared_obj_encoder
                _vr0 = batch.get('vtx_rank')
                _enc_pack_fwd = EncoderPack(_vm1.reshape(B * Nobj, Nv),
                                            _vr0.reshape(B * Nobj, Nv) if _vr0 is not None else None,
                                            len(_enc_m.hierarchical_layers), _enc_m.hierarchical_subsets, 1,
                                            small=_enc_m.subset_small_objects)

        def _net_step(t, curr_pos_bt, prev_pos_bt, B, Nobj, Nv, Nq, obj_mask, vtx_mask_t, anchor_idx, ref_pose,
                      phys_params, skip_step_size_batch, vtx_rank_all, state_tokens, bn_groups, knn_chunks, knn_slots,
                      enc_pack=None, knn_pk=None):
            """Network part of one step (distance features -> encoder -> AVP -> decoder -> predictor):
            per-anchor accelerations for the batch given current/previous positions. Called once per step,
            or once for all teacher-forced steps stacked along the batch (bn_groups = knn_chunks = steps)."""
            # Velocities
            vel_bt = curr_pos_bt - prev_pos_bt  # (B,Nobj,Nv,3)

            # Build per-object features and per-object latents
            per_obj_latents = []  # each (B, D)
            # Build per-object normalized query features and ref poses (padded)
            per_obj_vel_norm = []
            per_obj_relpos_norm = []
            per_obj_distvec_norm = []

            # Prepare far-away filler for invalid vertices in reference lists
            FAR = 1e6

            # (Anchor validation runs once per batch before the time loop: it depends only on the
            #  ground-truth masks and anchor indices, never on predictions.)

            # ========== VECTORIZED KNN FOR ALL OBJECTS AT ONCE ==========
            # curr_pos_bt: (B, Nobj, Nv, 3), vtx_mask_t: (B, Nobj, Nv)
            vmask_t = vtx_mask_t  # (B, Nobj, Nv)

            # Compute KNN distance vectors for ALL objects simultaneously
            # Only computes for valid objects and valid vertices!
            # Returns: (B, Nobj, Nv, 3) - vector from each vertex to nearest point on OTHER objects
            if knn_pk is not None:
                # packed valid vertices; with batched time steps the step-sized chunks share one pack
                _cs = B // knn_chunks
                dist_vec_all = (compute_knn_packed(curr_pos_bt, knn_pk) if knn_chunks == 1 else torch.cat(
                    [compute_knn_packed(curr_pos_bt[_c * _cs:(_c + 1) * _cs], knn_pk) for _c in range(knn_chunks)], dim=0))
            elif knn_chunks == 1:
                dist_vec_all = compute_knn_all_objects_vectorized(
                    curr_pos_bt, vmask_t, FAR, valid_obj_mask=obj_mask,
                    num_groups=self._knn_num_groups, active_slots=knn_slots
                )
            else:
                # batched time steps: one call per step-sized chunk (memory of the per-object cdist stays
                # that of a single step)
                _cs = B // knn_chunks
                dist_vec_all = torch.cat([
                    compute_knn_all_objects_vectorized(
                        curr_pos_bt[_c * _cs:(_c + 1) * _cs], vmask_t[_c * _cs:(_c + 1) * _cs], FAR,
                        valid_obj_mask=obj_mask[_c * _cs:(_c + 1) * _cs],
                        num_groups=self._knn_num_groups, active_slots=knn_slots)
                    for _c in range(knn_chunks)], dim=0)

            # Apply plane constraint adjustment (vectorized)
            # Reshape for distance_2_plane: (B, Nobj*Nv, 3)
            # Use reshape() instead of view() to handle non-contiguous tensors
            curr_pos_flat = curr_pos_bt.reshape(B, Nobj * Nv, 3)
            dist_vec_flat = dist_vec_all.reshape(B, Nobj * Nv, 3)
            dist_vec2_flat = distance_2_plane(curr_pos_flat, [0, 0, 0], "z", dist_vec_flat)
            dist_vec2_stacked = dist_vec2_flat.reshape(B, Nobj, Nv, 3)

            # Relative position (vectorized): curr_pos - ref_pose
            relpos_stacked = curr_pos_bt - ref_pose  # (B, Nobj, Nv, 3)

            # Velocity (already in correct shape)
            vel_stacked = vel_bt  # (B, Nobj, Nv, 3)

            # Noise augmentation (vectorized)
            if self.training:
                aug_config = self.config['data'].get('augmentation', {})
                if aug_config.get('dist_vec_noise', 0) > 0:
                    std = aug_config['dist_vec_noise']
                    dist_vec2_stacked = dist_vec2_stacked + torch.randn_like(dist_vec2_stacked) * std
                if aug_config.get('velocity_noise', 0) > 0:
                    std = aug_config['velocity_noise']
                    vel_stacked = vel_stacked + torch.randn_like(vel_stacked) * std
                if aug_config.get('rel_pos_noise', 0) > 0:
                    std = aug_config['rel_pos_noise']
                    relpos_stacked = relpos_stacked + torch.randn_like(relpos_stacked) * std

            # Apply vertex masks BEFORE normalization (vectorized across all objects)
            vmask_all = vtx_mask_t.unsqueeze(-1)  # (B,Nobj,Nv,1)
            dist_vec2_stacked = torch.where(vmask_all, dist_vec2_stacked, torch.zeros_like(dist_vec2_stacked))
            vel_stacked = torch.where(vmask_all, vel_stacked, torch.zeros_like(vel_stacked))
            relpos_stacked = torch.where(vmask_all, relpos_stacked, torch.zeros_like(relpos_stacked))

            # Apply network-predicted directional sign to ALL vertices' dist_vec (before encoder)
            # Only for t>1 (may have penetration in student forcing with predicted positions)
            # For t=1, GT positions (no penetration, sign=+1 implicit)
            # Use if_penetration as master switch (fallback to if_directional_distance for backwards compatibility)
            penetration_enabled = self.config['model'].get('if_penetration', self.config['model'].get('if_directional_distance', False))
            if penetration_enabled and t > 1:
                # Predict penetration sign using PREVIOUS step's features and CURRENT positions
                # curr_pos_bt was already selected by teacher/student forcing above
                # In student forcing mode: curr_pos_bt = all_pred_pos_for_rollout[-1] = blended_next (after Kabsch transform)
                # GT supervision will use predictions['obj_positions_ste'] to match
                pen_pred_all_vertices = self.predictor.predict_penetration_vertices(
                    prev_obj_layers,  # Use previous step's decoder outputs
                    prev_state_layers,  # Use previous step's state features
                    pos_full=curr_pos_bt  # AFTER Kabsch transform (student forcing) or GT (teacher forcing/t=1)
                )  # (B, Nobj, Nv, 1) in [-1,1]

                # Apply predicted sign to dist_vec for ALL vertices
                pen_sign_all = torch.sign(pen_pred_all_vertices + 1e-6)  # (B, Nobj, Nv, 1) in {-1,+1}
                dist_vec2_stacked = dist_vec2_stacked * pen_sign_all  # (B, Nobj, Nv, 3)

                # Store predicted signs for loss supervision (will compute GT later in validation)
                all_pen_pred_vertices_seq.append(pen_pred_all_vertices)
            elif penetration_enabled and t == 1:
                # First step: GT positions, no penetration prediction
                # Append placeholder (all +1) to maintain time dimension alignment
                # This will NOT be used in loss (masked out), but keeps dimensions consistent
                B, Nobj, Nv = dist_vec2_stacked.shape[0], dist_vec2_stacked.shape[1], dist_vec2_stacked.shape[2]
                placeholder = torch.ones((B, Nobj, Nv, 1), dtype=torch.float32, device=self.device)
                all_pen_pred_vertices_seq.append(placeholder)

            # No normalization — use raw values directly
            dist_vec2_norm_stacked = dist_vec2_stacked
            vel_norm_stacked = vel_stacked
            relpos_norm_stacked = relpos_stacked

            # Broadcast physics parameters and time_step to vertex dimension if enabled
            phys_broadcast = None  # (B, Nobj, Nv, 3)
            time_broadcast = None  # (B, Nobj, Nv, 2) - contains [s, s²]

            if self._if_phys_param:
                # Broadcast per-object physics params (B,Nobj,3) to all vertices (B,Nobj,Nv,3)
                phys_broadcast = phys_params.unsqueeze(2).expand(B, Nobj, Nv, 3)  # (B,Nobj,Nv,3)
                # Mask with vertex mask to ensure padding vertices get zero physics params
                phys_broadcast = torch.where(vmask_all, phys_broadcast, torch.zeros_like(phys_broadcast))

            if self._if_time_step:
                # Broadcast both s and s² to all objects and vertices (B,Nobj,Nv,2)
                # Input both s and s² because:
                #   - s: original step_size for temporal context
                #   - s²: the scaling factor that will be applied to network output
                #   - Network can leverage both for better learning
                step_size_s = skip_step_size_batch.float()  # (B,)
                step_size_sq = (skip_step_size_batch ** 2).float()  # (B,)
                # Stack s and s² -> (B, 2)
                time_features = torch.stack([step_size_s, step_size_sq], dim=-1)  # (B, 2)
                time_broadcast = time_features.view(B, 1, 1, 2).expand(B, Nobj, Nv, 2)  # (B,Nobj,Nv,2)
                # Time step doesn't need masking (it's a scalar property of the batch, not vertex-specific)
                # But for consistency, we can zero it for padding vertices
                # vmask_all is (B,Nobj,Nv,1), expand to match (B,Nobj,Nv,2)
                vmask_time = vmask_all.expand(-1, -1, -1, 2)  # (B,Nobj,Nv,2)
                time_broadcast = torch.where(vmask_time, time_broadcast, torch.zeros_like(time_broadcast))

            # PointNet Encoder: process each object's Nv vertices -> per-object global feature
            # Input dimension depends on flags:
            #   Base: 9D (dist_vec + velocity + rel_pos)
            #   +if_phys_param=true: +3D
            #   +if_time_step=true (no FiLM): +2D (s and s²)
            # Split into components and reshape to (B*Nobj, Nv, D) for encoder
            dist_reshaped = dist_vec2_norm_stacked.view(B * Nobj, Nv, 3)
            vel_reshaped = vel_norm_stacked.view(B * Nobj, Nv, 3)
            relpos_reshaped = relpos_norm_stacked.view(B * Nobj, Nv, 3)

            phys_reshaped = None
            if self._if_phys_param:
                phys_reshaped = phys_broadcast.view(B * Nobj, Nv, 3)

            time_reshaped = None
            if self._if_time_step and not self._film_conditioning:
                time_reshaped = time_broadcast.view(B * Nobj, Nv, 2)  # (s, s²)

            # Prepare vertex mask for encoder pooling
            # vmask_all is (B,Nobj,Nv,1), need (B*Nobj, Nv)
            vtx_mask_for_encoder = vmask_all.squeeze(-1).reshape(B * Nobj, Nv)  # (B*Nobj, Nv)
            height_reshaped = None
            if self.shared_obj_encoder.height_input:
                _h = torch.where(vmask_all, curr_pos_bt[..., 2:3], torch.zeros_like(curr_pos_bt[..., 2:3]))
                height_reshaped = _h.reshape(B * Nobj, Nv, 1)
            rank_reshaped = None
            if self.shared_obj_encoder.hierarchical_subsets:
                if vtx_rank_all is None:
                    raise KeyError("model.encoder.hierarchical_subsets requires data.fps_order: true (batch lacks 'vtx_rank')")
                rank_reshaped = vtx_rank_all.reshape(B * Nobj, Nv)

            # Call encoder with all features (encoder will concatenate them internally)
            # AVP: capture per-point features if AVP is enabled
            obj_features_encoded, _per_vtx_feat, _ = self.shared_obj_encoder(
                dist_reshaped,
                vel_reshaped,
                relpos_reshaped,
                phys_params=phys_reshaped,
                time_step=time_reshaped,
                valid_mask=vtx_mask_for_encoder,
                height=height_reshaped,
                vtx_rank=rank_reshaped,
                bn_groups=bn_groups,
                pack=enc_pack
            )  # obj_features_encoded: (B*Nobj, enc_out_dim), _per_vtx_feat: (B*Nobj, backbone_dim, Nv)

            # AVP: Anchor-Vertex Pooling — enrich anchor queries with collision-aware features
            _avp_anchor_features = None  # Will be (B, Nobj, Nq, avp_output_dim) if AVP enabled
            if self.avp is not None and _per_vtx_feat is not None:
                # _per_vtx_feat is (B*Nobj, backbone_dim, Nv) — Conv1d channels-first format
                # Transpose to (B*Nobj, Nv, backbone_dim) for AVP
                _pvf = _per_vtx_feat.transpose(1, 2)  # (B*Nobj, Nv, backbone_dim)
                # Anchor positions in flat format
                _curr_pos_flat = curr_pos_bt.reshape(B * Nobj, Nv, 3)  # (B*Nobj, Nv, 3)
                _anchor_pos_flat = torch.gather(
                    _curr_pos_flat, dim=1,
                    index=anchor_idx.reshape(B * Nobj, -1).unsqueeze(-1).expand(-1, -1, 3)
                )  # (B*Nobj, Nq, 3)
                _avp_anchor_features = self.avp(
                    _pvf, _curr_pos_flat, _anchor_pos_flat, vtx_mask_for_encoder
                )  # (B*Nobj, Nq, avp_output_dim)
                _avp_anchor_features = _avp_anchor_features.view(B, Nobj, anchor_idx.shape[2], -1)  # (B, Nobj, Nq, avp_output_dim)
                # Zero out padding objects (proj bias makes padding non-zero; consistent with other features)
                _avp_anchor_features = _avp_anchor_features * obj_mask.unsqueeze(-1).unsqueeze(-1).float()

            obj_features_encoded = obj_features_encoded.view(B, Nobj, -1)  # (B, Nobj, enc_out_dim)

            # Project to token dimension for Decoder
            obj_features = self.obj_feature_proj(obj_features_encoded)  # (B, Nobj, D)

            # FiLM: modulate encoder features with step-size [s, s²]
            if self.encoder_film is not None and skip_step_size_batch is not None:
                time_feat_film = torch.stack([
                    skip_step_size_batch.float(),
                    (skip_step_size_batch ** 2).float()
                ], dim=-1)  # (B, 2)
                obj_features = self.encoder_film(obj_features, time_feat_film)  # (B, Nobj, D)

            # Object anchor positions for RoPE (like RenderFormer's triangle approach)
            # RenderFormer: Triangle = 3 vertices (9D)
            # Our approach: Object = Nq anchors (Nq*3 D)
            # Gather anchor positions: (B, Nobj, Nq, 3)
            # Replace padding objects' positions with zeros before the gather:
            # padding objects have anchor_idx=0, which would gather arbitrary values into RoPE inputs.
            obj_mask_pos = obj_mask.unsqueeze(-1).unsqueeze(-1)  # (B, Nobj, 1, 1)
            safe_curr_pos_bt = torch.where(obj_mask_pos, curr_pos_bt, torch.zeros_like(curr_pos_bt))

            idx_exp_3_for_decoder = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B, Nobj, Nq, 3)
            obj_anchor_positions = torch.gather(safe_curr_pos_bt, dim=2, index=idx_exp_3_for_decoder)  # (B, Nobj, Nq, 3)
            # Flatten to (B, Nobj, Nq*3) for multi-point RoPE
            Nq = anchor_idx.shape[2]  # Number of anchors per object

            # ============== PE Type Routing (forward) ==============
            if self._pe_type in ('arope', 'se3'):
                # SE(3) uses arope anchors for cross-attention RoPE; SE(3) bias computed separately
                obj_anchor_flat = obj_anchor_positions.view(B, Nobj, Nq * 3)  # (B, Nobj, 12)
            elif self._pe_type in ('obb', 'pca'):
                obj_anchor_flat = self._geometry_encoder(
                    safe_curr_pos_bt, vtx_mask_t, obj_mask=obj_mask
                )  # (B, Nobj, 12)
            else:
                raise ValueError(f"Unknown pe_type='{self._pe_type}' in forward().")

            # ============== SE(3) Pairwise Bias Computation ==============
            _se3_bias = None
            if self._pe_type == 'se3' and self._se3_bias_module is not None:
                from src.models.se3_attention import extract_object_pose_features
                _rotmats, _centers, _scales = extract_object_pose_features(
                    safe_curr_pos_bt, ref_pose, vtx_mask_t, obj_mask
                )
                _se3_bias = self._se3_bias_module(_rotmats, _centers, _scales, obj_mask)

            # Decoder attention mask: SDPA additive format (0=attend, -inf=mask)
            # Prevents padding objects from participating in Decoder self-attention.
            # Combined with feature zeroing to ensure clean inputs.
            obj_attn_mask = torch.where(
                obj_mask.unsqueeze(1).unsqueeze(1),  # (B, 1, 1, Nobj)
                _attn_zero_fwd,
                _attn_ninf_fwd
            )  # (B, 1, 1, Nobj) float32

            # Zero out padding object features and anchor positions BEFORE Decoder
            # to prevent garbage from padding objects from destabilizing attention
            pad_mask_float = obj_mask.unsqueeze(-1).to(dtype=obj_features.dtype)
            obj_features = obj_features * pad_mask_float  # (B, Nobj, D)
            obj_anchor_flat = obj_anchor_flat * obj_mask.unsqueeze(-1).to(dtype=obj_anchor_flat.dtype)

            # Decoder interaction with per-object features
            # Per-layer FiLM: pass time_features for step-size conditioning in each layer
            _decoder_time_feat = time_feat_film if (self._film_conditioning and self._if_time_step) else None
            all_state_layers, all_obj_layers = self.decoder(
                state_tokens,
                obj_features,
                state_pos,
                obj_anchor_flat,
                mask=obj_attn_mask,  # SDPA additive mask for padding objects
                se3_bias=_se3_bias,
                time_features=_decoder_time_feat  # Per-layer FiLM
            )

            # Update state for next step (skip when no_global_state)
            if not self._no_global_state:
                state_tokens = all_state_layers[-1]

            # Build anchor query for Predictor (sample at anchor points Nq)
            # Query dimension depends on flags:
            #   Base: 9D (dist_vec + vel + rel_pos)
            #   +if_phys_param=true: +3D = 12D
            #   +if_time_step=true (no FiLM): +2D (s and s²)
            idx_exp_3 = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
            distvec_at_anchors = torch.gather(dist_vec2_norm_stacked, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)
            vel_at_anchors = torch.gather(vel_norm_stacked, dim=2, index=idx_exp_3)            # (B,Nobj,Nq,3)
            relpos_at_anchors = torch.gather(relpos_norm_stacked, dim=2, index=idx_exp_3)      # (B,Nobj,Nq,3)

            # Build query features list
            query_features = [distvec_at_anchors, vel_at_anchors, relpos_at_anchors]

            # Add physics parameters at anchors if enabled
            if self._if_phys_param:
                phys_at_anchors = torch.gather(phys_broadcast, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)
                query_features.append(phys_at_anchors)

            # Add time step at anchors if enabled AND not using FiLM
            if self._if_time_step and not self._film_conditioning:
                idx_exp_2 = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 2)  # (B,Nobj,Nq,2)
                time_at_anchors = torch.gather(time_broadcast, dim=2, index=idx_exp_2)  # (B,Nobj,Nq,2)
                query_features.append(time_at_anchors)

            # AVP: Append collision-aware anchor features to query
            if _avp_anchor_features is not None:
                query_features.append(_avp_anchor_features)  # (B, Nobj, Nq, avp_output_dim)

            # Concatenate all query features
            anchor_query = torch.cat(query_features, dim=-1)  # (B,Nobj,Nq, 9/12 + avp_dim)

            # Extract anchor positions for RoPE (current frame positions at anchors)
            anchor_pos_at_curr = torch.gather(curr_pos_bt, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)

            # Penetration sign was already applied to all vertices' dist_vec before the encoder;
            # the anchor query therefore carries the sign information as well.

            _ref_src = _object_centered(ref_pose, vmask_all.squeeze(-1).bool()) if self._ref_pose_object_centered else ref_pose
            if self._ref_pose_object_centered and not getattr(self, '_refpose_logged_fwd', False):
                self._refpose_logged_fwd = True
                print(f"[RefPose] object-centred (x/y only, z absolute) reference pose ACTIVE (forward): |shift| max={((ref_pose - _ref_src).abs().max()).item():.3f} m", flush=True)
            ref_pose_at_anchors = torch.gather(_ref_src, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)

            # Check for NaN/Inf before predictor
            if DEBUG_NAN and (torch.isnan(anchor_query).any() or torch.isinf(anchor_query).any()):
                logger.error(f"NaN/Inf detected in anchor_query!")
                logger.error(f"  distvec_at_anchors: nan={torch.isnan(distvec_at_anchors).sum().item()}, inf={torch.isinf(distvec_at_anchors).sum().item()}")
                logger.error(f"  vel_at_anchors: nan={torch.isnan(vel_at_anchors).sum().item()}, inf={torch.isinf(vel_at_anchors).sum().item()}")
                logger.error(f"  relpos_at_anchors: nan={torch.isnan(relpos_at_anchors).sum().item()}, inf={torch.isinf(relpos_at_anchors).sum().item()}")

            if DEBUG_NAN and (torch.isnan(ref_pose_at_anchors).any() or torch.isinf(ref_pose_at_anchors).any()):
                logger.error(f"NaN/Inf detected in ref_pose_at_anchors!")
                logger.error(f"  ref_pose_at_anchors range: [{ref_pose_at_anchors.min().item():.4f}, {ref_pose_at_anchors.max().item():.4f}]")

            for layer_idx, obj_layer in enumerate(all_obj_layers if DEBUG_NAN else ()):
                if torch.isnan(obj_layer).any() or torch.isinf(obj_layer).any():
                    logger.error(f"NaN/Inf detected in all_obj_layers[{layer_idx}]!")
                    logger.error(f"  Shape: {obj_layer.shape}, nan={torch.isnan(obj_layer).sum().item()}, inf={torch.isinf(obj_layer).sum().item()}")

            # Predict accelerations per anchor using grouped format (B, Nobj, Nq, *)
            accel_grouped = self.predictor(
                all_obj_layers,
                anchor_query,           # (B, Nobj, Nq, 9/12)
                ref_pose=ref_pose_at_anchors,  # (B, Nobj, Nq, 3)
                query_pos=anchor_pos_at_curr,  # (B, Nobj, Nq, 3)
                key_pos=obj_anchor_flat,        # (B, Nobj, Nq*3) - multi-point RoPE
                multi_scale_state_features=None if self._no_global_state else all_state_layers,
                state_pos=None if self._no_global_state else state_pos,
                obj_mask=obj_mask  # Object validity mask for padding
            )  # Returns (B, Nobj, Nq, 3)

            # Check for NaN/Inf after predictor
            if DEBUG_NAN and (torch.isnan(accel_grouped).any() or torch.isinf(accel_grouped).any()):
                logger.error(f"NaN/Inf detected in predictor output!")
                logger.error(f"  accel_grouped: nan={torch.isnan(accel_grouped).sum().item()}, inf={torch.isinf(accel_grouped).sum().item()}")
                logger.error(f"  accel_grouped range: min={accel_grouped[~torch.isnan(accel_grouped)].min().item() if (~torch.isnan(accel_grouped)).any() else 'all nan'}, max={accel_grouped[~torch.isnan(accel_grouped)].max().item() if (~torch.isnan(accel_grouped)).any() else 'all nan'}")

            return vmask_t, accel_grouped, all_obj_layers, all_state_layers, state_tokens

        # ========== TIME-BATCHED NETWORK PASS (training.time_batched_forward) ==========
        # With pure teacher forcing every step's network input is built from ground-truth frames, so the network
        # part of all T-2 steps runs once on the steps stacked along the batch (step-major: rows [g*B, (g+1)*B)
        # belong to step g+1). The encoder keeps per-step BatchNorm statistics (grouped masked BN: same
        # statistics, gradients and running-statistics updates as step-by-step calls), and the per-step physics
        # (Verlet, Kabsch, ground handling) below is unchanged. Equal to the step-by-step path up to
        # floating-point rounding and the dropout random stream.
        _tb_accel = None
        if self.training and self._time_batched_forward and T > 2:
            if teacher_forcing_ratio < 1.0:
                raise ValueError("training.time_batched_forward requires teacher_forcing_ratio == 1 (inputs must be ground truth)")
            Tw = T - 2
            rope_memo_reset()

            def _rep(x):
                return x.unsqueeze(0).expand(Tw, *x.shape).reshape(Tw * x.shape[0], *x.shape[1:])
            _ref_b = _rep(ref_pose) if (torch.is_tensor(ref_pose) and ref_pose.dim() == 4 and ref_pose.shape[0] == B) else ref_pose
            _vr = batch.get('vtx_rank')
            _, _acc_s, _, _, _ = _net_step(
                None,
                obj_pos[:, 1:T - 1].transpose(0, 1).reshape(Tw * B, Nobj, Nv, 3),
                obj_pos[:, 0:T - 2].transpose(0, 1).reshape(Tw * B, Nobj, Nv, 3),
                Tw * B, Nobj, Nv, Nq, _rep(obj_mask),
                vtx_mask[:, 1:T - 1].transpose(0, 1).reshape(Tw * B, Nobj, Nv),
                _rep(anchor_idx), _ref_b,
                _rep(phys_params) if phys_params is not None else None,
                _rep(skip_step_size_batch) if skip_step_size_batch is not None else None,
                _rep(_vr) if _vr is not None else None,
                None, Tw, Tw, _knn_active_slots,
                enc_pack=(EncoderPack(vtx_mask[:, 1:T - 1].transpose(0, 1).reshape(Tw * B * Nobj, Nv),
                                      _rep(_vr).reshape(Tw * B * Nobj, Nv) if _vr is not None else None,
                                      len(self.shared_obj_encoder.hierarchical_layers),
                                      self.shared_obj_encoder.hierarchical_subsets, Tw,
                                      small=self.shared_obj_encoder.subset_small_objects)
                          if self._encoder_packed else None),
                knn_pk=_knn_pack_fwd)
            _tb_accel = _acc_s.view(Tw, B, *_acc_s.shape[1:])

        # Time loop predicts for steps t in [1, T-2] -> next positions x_{t+1}
        for t in range(1, T - 1):
            rope_memo_reset()  # rotary angles are memoized within one step only
            # Teacher vs student forcing (no special-case first step)
            use_teacher_forcing = (torch.rand(1).item() < teacher_forcing_ratio)

            if use_teacher_forcing or len(all_pred_pos_for_rollout) == 0:
                prev_pos_bt = obj_pos[:, t-1]   # (B,Nobj,Nv,3)
                curr_pos_bt = obj_pos[:, t]
            else:
                # Use last predicted (STE blended) as current for autoregressive stability
                curr_pos_bt = all_pred_pos_for_rollout[-1]  # (B,Nobj,Nv,3)
                if len(all_pred_pos_for_rollout) >= 2:
                    prev_pos_bt = all_pred_pos_for_rollout[-2]
                else:
                    prev_pos_bt = obj_pos[:, t-1]

            if _tb_accel is not None:
                # time-batched: this step's accelerations were computed with all teacher-forced steps
                vmask_t = vtx_mask[:, t]
                accel_grouped = _tb_accel[t - 1]
                all_obj_layers = all_state_layers = None
            else:
                vmask_t, accel_grouped, all_obj_layers, all_state_layers, state_tokens = _net_step(
                    t, curr_pos_bt, prev_pos_bt, B, Nobj, Nv, Nq, obj_mask, vtx_mask[:, t], anchor_idx, ref_pose,
                    phys_params, skip_step_size_batch, batch.get('vtx_rank'), state_tokens, 1, 1, _knn_active_slots,
                    enc_pack=_enc_pack_fwd, knn_pk=_knn_pack_fwd)
            # Already in grouped format, no need to reshape
            accel_flat = accel_grouped  # (B, Nobj, Nq, 3)
            accel_base_norm = accel_flat  # Network output: base acceleration (step_size-independent)

            # Verlet integration (dt_verlet=1)
            dt_squared = self.dt_verlet * self.dt_verlet
            with torch.amp.autocast('cuda', enabled=False):
                # No normalizer — network output is raw acceleration
                accel_base_denorm = accel_base_norm

                # ============================================================
                # Explicit step_size² scaling for multi-timescale training
                # ============================================================
                # Physical reasoning:
                #   - Network outputs "base acceleration" a_base (step_size-independent)
                #   - Actual acceleration = a_base × step_size²
                #   - This explicit scaling removes the burden from network to learn s² relationship
                #   - Benefit: Better generalization to unseen step_sizes, easier optimization
                # ============================================================
                if self._if_time_step and skip_step_size_batch is not None:
                    step_size_sq = (skip_step_size_batch ** 2).view(B, 1, 1, 1)  # (B,1,1,1)
                    accel_scaled_denorm = accel_base_denorm * step_size_sq  # (B,Nobj,Nq,3)
                    # Also scale normalized version for loss computation
                    accel_no_icp_norm_anchor = accel_base_norm * step_size_sq
                else:
                    accel_scaled_denorm = accel_base_denorm
                    accel_no_icp_norm_anchor = accel_base_norm

                # Build base next positions for all vertices as zero-accel extrapolation
                # No detach: allow gradients to flow through position history for long-horizon training
                next_pos_bt = 2 * curr_pos_bt - prev_pos_bt  # (B,Nobj,Nv,3)

                # Vectorized scatter using scatter_reduce (faster than per-object loop)
                # Step 1: Gather all anchor positions at once
                anchor_idx_exp = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
                curr_anchors = torch.gather(curr_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                prev_anchors = torch.gather(prev_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)

                # Step 2: Compute next positions at anchors using Verlet (with scaled acceleration)
                next_anchors = accel_scaled_denorm * dt_squared + 2 * curr_anchors - prev_anchors  # (B,Nobj,Nq,3)

                # Step 3: Compute updates (delta from base prediction)
                base_anchors = torch.gather(next_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                anchor_updates = next_anchors - base_anchors  # (B,Nobj,Nq,3)

                # Mask padding objects' updates to zero
                # Padding objects have anchor_idx=0 which would scatter to wrong positions
                obj_mask_scatter = obj_mask.unsqueeze(-1).unsqueeze(-1).float()  # (B,Nobj,1,1)
                anchor_updates = anchor_updates * obj_mask_scatter  # padding objects -> zero updates

                # Step 4: Vectorized scatter using scatter_reduce
                # Flatten batch and object dimensions for efficient scatter
                # Use reshape() to handle non-contiguous tensors
                B_Nobj_Nv = B * Nobj * Nv
                next_pos_flat = next_pos_bt.reshape(B_Nobj_Nv, 3)  # (B*Nobj*Nv, 3)

                # Create flat indices for scatter_reduce
                # anchor_idx is (B, Nobj, Nq), we need indices in flat (B*Nobj*Nv) space
                obj_offsets = torch.arange(Nobj, device=anchor_idx.device).view(1, Nobj, 1) * Nv  # (1,Nobj,1)
                batch_offsets = torch.arange(B, device=anchor_idx.device).view(B, 1, 1) * (Nobj * Nv)  # (B,1,1)
                flat_anchor_idx = anchor_idx + obj_offsets + batch_offsets  # (B,Nobj,Nq)
                flat_anchor_idx = flat_anchor_idx.reshape(B * Nobj * Nq)  # (B*Nobj*Nq,)

                # Flatten updates
                anchor_updates_flat = anchor_updates.reshape(B * Nobj * Nq, 3)  # (B*Nobj*Nq, 3)

                # Apply updates using scatter_reduce (more efficient than loop)
                next_pos_flat = torch.scatter_reduce(
                    next_pos_flat,
                    dim=0,
                    index=flat_anchor_idx.unsqueeze(-1).expand(-1, 3),
                    src=anchor_updates_flat,
                    reduce='sum',
                    include_self=True
                )

                # Reshape back
                next_pos_bt = next_pos_flat.reshape(B, Nobj, Nv, 3)

                # Per-object rigid projection using anchors (Kabsch on anchors), then apply to full Nv
                if registration_refs is None:
                    registration_refs = curr_pos_bt.detach().clone()  # (B,Nobj,Nv,3)

                # Vectorized Kabsch: Process ONLY valid objects (filter padding BEFORE SVD)
                # Step 1: Gather anchor positions for all objects at once
                anchor_idx_exp = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
                last_anchor_all = torch.gather(registration_refs, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                curr_anchor_all = torch.gather(next_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)

                # Reshape to (B*Nobj, Nq, 3) for batch processing
                B_Nobj = B * Nobj
                last_anchor_flat = last_anchor_all.view(B_Nobj, Nq, 3)
                curr_anchor_flat = curr_anchor_all.view(B_Nobj, Nq, 3)

                # Filter padding objects before SVD; they would produce garbage R/t
                obj_mask_flat = obj_mask.view(B_Nobj)  # (B*Nobj,)
                valid_idx = _kabsch_valid_idx_fwd  # indices of valid objects (fixed within the forward)
                n_valid = valid_idx.shape[0]

                # Initialize R/t for ALL objects (identity R, zero t)
                R_flat = torch.eye(3, device=device, dtype=last_anchor_flat.dtype).unsqueeze(0).expand(B_Nobj, -1, -1).clone()
                t_vec_flat = torch.zeros(B_Nobj, 1, 3, device=device, dtype=last_anchor_flat.dtype)

                if n_valid > 0:
                    # Gather only valid objects' anchors
                    last_anchor_valid = last_anchor_flat[valid_idx]  # (Nvalid, Nq, 3)
                    curr_anchor_valid = curr_anchor_flat[valid_idx]  # (Nvalid, Nq, 3)


                    if self.if_roma:
                        # RoMa mode: batch process only valid objects
                        # RoMa requires FP32 for numerical stability
                        with torch.amp.autocast('cuda', enabled=False):
                            R_valid, t_vec_valid = roma.rigid_points_registration(
                                last_anchor_valid.float(), 
                                curr_anchor_valid.float()
                            )
                        R_valid = R_valid.float()  # Ensure FP32 for downstream
                        t_vec_valid = t_vec_valid.float().unsqueeze(1)  # (Nvalid, 1, 3)

                    else:
                        # Standard SVD Kabsch: batch process only valid objects
                        centroid_last = last_anchor_valid.mean(dim=1, keepdim=True)  # (Nvalid,1,3)
                        centroid_curr = curr_anchor_valid.mean(dim=1, keepdim=True)  # (Nvalid,1,3)
                        last_c = last_anchor_valid - centroid_last
                        curr_c = curr_anchor_valid - centroid_curr

                        # Batch SVD for valid objects only
                        H = torch.bmm(last_c.transpose(1, 2), curr_c)  # (Nvalid,3,3)
                        with torch.amp.autocast('cuda', enabled=False):
                            H_float = H.float()
                            U, S, Vt = torch.linalg.svd(H_float)  # Batch SVD!
                        U = U.to(H.dtype); Vt = Vt.to(H.dtype)
                        R_valid = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2))  # (Nvalid,3,3)

                        # Fix reflections
                        with torch.amp.autocast('cuda', enabled=False):
                            det_R = torch.det(R_valid.float()).to(R_valid.dtype)
                        need_flip = det_R < 0
                        if need_flip.any():
                            Fm = torch.eye(3, device=R_valid.device, dtype=R_valid.dtype).unsqueeze(0).expand(n_valid, -1, -1).clone()
                            Fm[need_flip, 2, 2] = -1
                            R_valid = torch.bmm(torch.bmm(Vt.transpose(1, 2), Fm), U.transpose(1, 2))

                        # Translation
                        rotated_centroid = torch.bmm(R_valid, centroid_last.transpose(1, 2)).transpose(1, 2)
                        t_vec_valid = centroid_curr - rotated_centroid  # (Nvalid,1,3)

                    # Scatter valid R/t back to full tensors
                    R_flat[valid_idx] = R_valid
                    t_vec_flat[valid_idx] = t_vec_valid

                # Reshape back to (B,Nobj,3,3) and (B,Nobj,1,3)
                R_all = R_flat.view(B, Nobj, 3, 3)
                t_vec_all = t_vec_flat.view(B, Nobj, 1, 3)

                # ========== VECTORIZED TRANSFORMATION APPLICATION ==========
                # Apply R, t to all objects in a single bmm call (no per-object loop)
                # registration_refs: (B, Nobj, Nv, 3)
                # R_all: (B, Nobj, 3, 3)
                # t_vec_all: (B, Nobj, 1, 3)

                # Reshape for batch bmm: (B, Nobj, Nv, 3) -> (B*Nobj, Nv, 3)
                last_flat = registration_refs.view(B * Nobj, Nv, 3)
                R_flat = R_all.view(B * Nobj, 3, 3)
                t_flat = t_vec_all.view(B * Nobj, 1, 3)

                # Single bmm for ALL objects: y = R @ x + t (RoMa convention)
                # last_flat.transpose(1, 2): (B*Nobj, 3, Nv)
                # bmm result: (B*Nobj, 3, Nv)
                # transpose back + t: (B*Nobj, Nv, 3)
                y_flat = torch.bmm(R_flat, last_flat.transpose(1, 2)).transpose(1, 2) + t_flat
                # Use reshape() because transpose() makes tensor non-contiguous
                yproj_bt = y_flat.reshape(B, Nobj, Nv, 3).float()  # FP32 for numerical stability

                # Raw position BEFORE Kabsch/STE (for position loss supervision)
                next_pos_raw_norm = next_pos_bt

                # Apply gradient strategy based on transform type
                if self.if_roma:
                    # RoMa mode: fully differentiable, gradients flow through Kabsch
                    blended_next = yproj_bt
                else:
                    # Standard mode: straight-through estimator (STE)
                    # Use Kabsch value but only pass gradient to original prediction
                    blended_next = next_pos_bt + (yproj_bt - next_pos_bt).detach()

                # Recompute accelerations from blended positions (for acceleration_after_icp loss)
                accel_icp_denorm = (blended_next - 2 * curr_pos_bt + prev_pos_bt) / dt_squared
                accel_icp_norm = accel_icp_denorm

            # ========== INTERSECTION HANDLING ==========
            # Resolve object-ground and object-object penetrations (STE-differentiable)
            intersection_cfg = self.config.get('model', {}).get('intersection_handling', {})
            rigid_next = blended_next  # Kabsch output before the ground/object correction
            blended_next = handle_intersections(
                blended_next, obj_mask, vmask_t, intersection_cfg
            )

            # Update registration refs to the post-correction position (consistent with rollout); with
            # model.kabsch_reference 'fixed' every step fits the same reference mesh (the first current frame)
            if not self._kabsch_reference_fixed:
                registration_refs = blended_next.detach()

            # Store positions:
            # - STE blended positions for autoregressive rollout (more stable)
            # - Raw positions for loss supervision (direct gradient)
            # - ICP-transformed positions for position_after_icp loss
            all_pred_pos_for_rollout.append(blended_next)  # For next step input
            all_pred_pos_raw.append(next_pos_bt)  # For position_before_icp loss
            all_pred_pos_raw_norm.append(next_pos_raw_norm)  # For position_before_icp loss (normalized)
            # ICP/Kabsch-transformed positions for position_after_icp loss. after_kabsch_loss 'rigid': the Kabsch
            # output itself, the state the acceleration_after_icp loss uses (the stated after-Kabsch losses);
            # 'handled': the state after the intersection correction.
            blended_next_norm = rigid_next if self._after_kabsch_loss_rigid else blended_next
            all_pred_pos_icp_norm.append(blended_next_norm)  # For position_after_icp loss
            all_accel_no_icp_norm_anchor_seq.append(accel_no_icp_norm_anchor)

            # Save current features for next step's penetration prediction
            prev_obj_layers = all_obj_layers
            prev_state_layers = all_state_layers
            all_accel_icp_norm_seq.append(accel_icp_norm)

        # Stack time dimension
        # Design principle:
        # - Position outputs: RAW predictions (before Kabsch) AND ICP-transformed (after Kabsch)
        # - Acceleration outputs: before and after Kabsch transform
        # - Autoregressive rollout: uses ICP-transformed positions internally for stability
        output = {
            'obj_positions_denorm': torch.stack(all_pred_pos_raw, dim=1),      # (B,T-2,Nobj,Nv,3) - raw (before Kabsch)
            'obj_positions_norm': torch.stack(all_pred_pos_raw_norm, dim=1),    # (B,T-2,Nobj,Nv,3) - raw normalized (before Kabsch)
            'obj_positions_icp_norm': torch.stack(all_pred_pos_icp_norm, dim=1),  # (B,T-2,Nobj,Nv,3) - ICP-transformed normalized (after Kabsch)
            'obj_positions_ste': torch.stack(all_pred_pos_for_rollout, dim=1),  # (B,T-2,Nobj,Nv,3) - ICP blended (matches rollout)
            'obj_accelerations_no_icp_norm_anchor': torch.stack(all_accel_no_icp_norm_anchor_seq, dim=1),  # before Kabsch
            'obj_accelerations_norm': torch.stack(all_accel_icp_norm_seq, dim=1),  # after Kabsch
        }

        # Include penetration predictions for ALL vertices if collected
        penetration_enabled = self.config['model'].get('if_penetration', self.config['model'].get('if_directional_distance', False))
        if penetration_enabled and len(all_pen_pred_vertices_seq) > 0:
            try:
                output['obj_pen_pred_vertices'] = torch.stack(all_pen_pred_vertices_seq, dim=1)  # (B,T-2,Nobj,Nv,1)
            except Exception as _e:
                logger.warning(f"Failed to stack pen predictions: {_e}")

        if return_all:
            # Not storing states for multi-object yet
            pass

        return output
    
    def rollout(self, initial_batch, num_steps, dt=0.01, return_extras=True, warmup_frames=2):
        """
        Rollout predictions from initial conditions (multi-object).

        Args:
            initial_batch: dict from dataloader containing at least
                - obj_positions (B, T0, Nobj, Nv, 3)
                - valid_obj_mask (B, Nobj)
                - valid_vtx_mask (B, T0, Nobj, Nv)
                - ref_obj_positions (B or 1, Nobj, Nv, 3) [optional]
            num_steps: total frames to output (including provided GT warmup frames)
            dt: timestep (unused; integration handled in model forward)
            warmup_frames: number of GT warmup frames (default=2 for our model, HopNet uses 4)
        Returns:
            dict with 'obj_positions' of shape (B, num_steps, Nobj, Nv, 3)
                and optionally 'obj_body_dist_vec2' and 'obj_accelerations_anchor' if return_extras=True
        """
        self.eval()
        with torch.no_grad():
            if 'obj_positions' not in initial_batch:
                raise ValueError("rollout expects multi-object 'obj_positions' in initial_batch")

            obj_pos_all = initial_batch['obj_positions']  # (B,T0,Nobj,Nv,3)
            valid_obj_mask = initial_batch['valid_obj_mask']  # (B,Nobj)
            valid_vtx_mask_all = initial_batch['valid_vtx_mask']  # (B,T0,Nobj,Nv)

            # Read physics parameters and time step if enabled
            phys_params = None  # (B, Nobj, 3): [mass, friction, restitution]
            skip_step_size_batch = None  # (B,)
            if self._if_phys_param:
                if 'phys_params' not in initial_batch:
                    raise KeyError("rollout requires 'phys_params' in initial_batch when if_phys_param=true")
                phys_params = initial_batch['phys_params'].to(obj_pos_all.device)  # (B, Nobj, 3)
            if self._if_time_step:
                if 'skip_step_size' not in initial_batch:
                    raise KeyError("rollout requires 'skip_step_size' in initial_batch when if_time_step=true")
                skip_step_size_batch = initial_batch['skip_step_size'].to(obj_pos_all.device)  # (B,)

                # Optional debug logging of the rollout step size
                if_debug = self.config.get('data', {}).get('random_step_size', {}).get('if_debug', False)
                if if_debug:
                    logger.info(f"[DEBUG Rollout] step_size for rollout: {skip_step_size_batch.cpu().numpy().tolist()}, num_steps={num_steps}")

                # Validate step_size is within training range
                # Check if multi-timescale training is enabled and validate step_size
                random_step_config = self.config.get('data', {}).get('random_step_size', {})
                if random_step_config.get('enable', False):
                    trained_step_sizes = random_step_config.get('step_sizes', [1])  # e.g., [1, 2, 5, 10]
                    for b_idx, step_size_val in enumerate(skip_step_size_batch.cpu().numpy()):
                        if step_size_val not in trained_step_sizes:
                            logger.warning(
                                f"[Rollout Warning] Sample {b_idx}: step_size={step_size_val} is NOT in trained range {trained_step_sizes}!\n"
                                f"Model was only trained on step_sizes {trained_step_sizes}. Prediction may be inaccurate due to extrapolation."
                            )
                        elif if_debug:
                            logger.info(f"[Rollout Validation] Sample {b_idx}: step_size={step_size_val} is within trained range {trained_step_sizes}")

            B, T0, Nobj, Nv, _ = obj_pos_all.shape
            device = obj_pos_all.device

            # Require at least warmup_frames GT frames (minimum 2 for Verlet initialization)
            min_required = max(2, warmup_frames)
            if T0 < min_required:
                raise ValueError(f"rollout requires at least {min_required} GT frames (warmup_frames={warmup_frames}), got {T0}")

            # Reference pose: MUST be global frame 0 (same as training)
            if 'ref_obj_positions' not in initial_batch:
                raise ValueError("rollout requires 'ref_obj_positions' (global frame 0) in initial_batch")
            ref_pose = initial_batch['ref_obj_positions']  # (B?,Nobj,Nv,3)
            if ref_pose.dim() == 3:
                ref_pose = ref_pose.unsqueeze(0).expand(B, -1, -1, -1).to(device)
            else:
                ref_pose = ref_pose.to(device)

            # ========== Center scene by first frame (eval/rollout alignment) ==========
            # If enabled, subtract the global center of the first frame from all coordinates
            # so the rollout coordinate system matches the training setup.
            # center_t0 is stored per sample so visualization can apply the same shift to disk-loaded GT.
            _center_t0_per_sample = []  # Store per-sample center for visualization GT alignment
            if self.config.get('data', {}).get('center_scene_by_t0', False):
                for b in range(B):
                    vmask_b0 = valid_vtx_mask_all[b, 0]  # (Nobj, Nv)
                    v_t0 = obj_pos_all[b, 0][vmask_b0]  # Valid vertices at t=0
                    if v_t0.shape[0] > 0:
                        center_t0 = v_t0.mean(dim=0)  # (3,)
                        _center_t0_per_sample.append(center_t0.cpu().numpy())
                        
                        logger.info(f"[Center-T0 ROLLOUT] b={b}: center_t0={center_t0.cpu().numpy().tolist()}, "
                                   f"obj_pos range before: [{obj_pos_all[b].min().item():.4f}, {obj_pos_all[b].max().item():.4f}]")
                        
                        # Shift positions
                        obj_pos_all[b] = obj_pos_all[b] - center_t0
                        ref_pose[b] = ref_pose[b] - center_t0
                        
                        logger.info(f"[Center-T0 ROLLOUT] b={b} after: "
                                   f"obj_pos range: [{obj_pos_all[b].min().item():.4f}, {obj_pos_all[b].max().item():.4f}]")
                    else:
                        _center_t0_per_sample.append(None)
                logger.info(f"[Center-T0 ROLLOUT] Applied centering to {len(_center_t0_per_sample)} samples")

            # Require anchors for query and supervision (no fallback)
            if 'anchor_idx' not in initial_batch:
                raise KeyError("rollout requires 'anchor_idx' in initial_batch for anchor-based querying")
            anchor_idx = initial_batch['anchor_idx'].to(device).long()  # (B,Nobj,Nq)
            if anchor_idx.dim() != 3 or anchor_idx.shape[0] != B or anchor_idx.shape[1] != Nobj:
                raise ValueError(f"anchor_idx must be (B,Nobj,Nq); got {tuple(anchor_idx.shape)} vs (B={B},Nobj={Nobj})")
            Nq = int(anchor_idx.shape[2])

            # Initialize previous/current positions (use last 2 frames of warmup for Verlet)
            # For warmup_frames=2: use frame 0, 1
            # For warmup_frames=4: use frame 2, 3 (GT frames 0,1,2,3 then predict from frame 4)
            prev_idx = warmup_frames - 2
            curr_idx = warmup_frames - 1
            prev_pos_bt = obj_pos_all[:, prev_idx]  # (B,Nobj,Nv,3)
            curr_pos_bt = obj_pos_all[:, curr_idx]
            registration_refs = curr_pos_bt.detach().clone()

            # Initialize persistent state tokens (carry across steps like training)
            # When no_global_state=True, pass None to decoder
            if self._no_global_state:
                state_tokens = None
                state_pos = None
            else:
                state_tokens, state_pos = self.state_repr.get_state_and_positions(B, device)

            # Collect frames: include ALL GT warmup frames first
            out_seq = [obj_pos_all[:, i] for i in range(warmup_frames)]

            # Steps to predict beyond the provided GT frames in the output horizon
            remaining = max(0, num_steps - warmup_frames)

            # Collect per-step input body_dist_vec2 for visualization (curr frame at each step)
            # Only collect if return_extras=True (video gen) to save memory during RMSE eval
            all_dist_vec2_seq = [] if return_extras else None  # list of (B, Nobj, Nv, 3)
            # Collect predicted anchor accelerations (denormalized) for visualization
            all_accel_anchor_denorm_seq = [] if return_extras else None  # list of (B, Nobj, Nq, 3)

            # Previous step's features for penetration prediction (step>0)
            prev_obj_layers_rollout = None
            prev_state_layers_rollout = None

            # Optional CUDA-graph fast path (RIGIDFORMER_CUDA_GRAPH=1): after the
            # early steps whose vertex-mask index is still advancing, every step
            # executes an identical op stream on fixed shapes, so one step is
            # captured as a CUDA graph and replayed for the remainder. Cross-step
            # tensors are kept in fixed-address buffers that the graph itself
            # rotates, which makes each replay self-contained.
            use_cuda_graph = (os.environ.get('RIGIDFORMER_CUDA_GRAPH', '0') == '1'
                              and not return_extras and torch.cuda.is_available())
            # Constant scalars used inside the step; created once so the step body
            # performs no host-to-device transfers (required for graph capture).
            _attn_zero = torch.tensor(0.0, dtype=torch.float32, device=device)
            _attn_ninf = torch.tensor(float('-inf'), dtype=torch.float32, device=device)
            # Kabsch padding filter: valid_obj_mask is fixed for the whole rollout,
            # so the valid-object index set is computed once (nonzero would force a
            # host sync every step and cannot run under graph capture).
            _kabsch_mask_flat = valid_obj_mask.reshape(B * Nobj)
            _kabsch_valid_idx = _kabsch_mask_flat.nonzero(as_tuple=True)[0]
            _kabsch_n_valid = int(_kabsch_valid_idx.shape[0])
            use_cuda_graph = use_cuda_graph and bool(getattr(self, 'if_roma', False))
            graph_obj = None
            graph_obj2 = None
            graph_capture_step = T0 - 1  # first step index with the saturated mask
            if use_cuda_graph and remaining <= graph_capture_step + 1:
                use_cuda_graph = False

            # Progress bar for rollout (only show on rank 0)
            import torch.distributed as dist
            from tqdm import tqdm
            show_pbar = True
            if dist.is_available() and dist.is_initialized():
                show_pbar = (dist.get_rank() == 0)

            # Object slots holding a valid object in at least one sample; fixed for the rollout, so
            # computed once here (a host sync outside the step loop and any graph capture).
            _knn_active_slots_r = (torch.nonzero(valid_obj_mask.any(dim=0)).flatten().tolist()
                                   if self._knn_num_groups <= 1 else None)
            # Valid-vertex packs, one per mask index the loop uses (min(warmup_frames - 1 + step, T0 - 1)); built here, outside
            # the step loop and any graph capture.
            _knn_packs_r, _enc_packs_r = {}, {}
            if self._knn_packed or self._encoder_packed:
                _vr_r = initial_batch.get('vtx_rank')
                for _mi in range(1, valid_vtx_mask_all.shape[1]):
                    _vm_r = valid_vtx_mask_all[:, _mi]
                    if self._knn_packed:
                        _knn_packs_r[_mi] = knn_pack(_vm_r, valid_obj_mask)
                    if self._encoder_packed:
                        _enc_packs_r[_mi] = EncoderPack(
                            _vm_r.reshape(B * Nobj, Nv), _vr_r.reshape(B * Nobj, Nv) if _vr_r is not None else None,
                            len(self.shared_obj_encoder.hierarchical_layers), self.shared_obj_encoder.hierarchical_subsets, 1,
                            small=self.shared_obj_encoder.subset_small_objects)

            # Anchor validation. Anchors and vertex masks are fixed for the whole
            # rollout, so every per-step check is equivalent to checking each mask
            # index once up front; this keeps the step loop free of host syncs.
            for _mask_idx in range(1, T0):
                for i in range(Nobj):
                    vmask_curr = valid_vtx_mask_all[:, _mask_idx, i]      # (B,Nv)
                    idx_i_chk = anchor_idx[:, i]                          # (B,Nq)
                    mask_at_anchors = torch.gather(vmask_curr, dim=1, index=idx_i_chk)
                    obj_valid_b = valid_obj_mask[:, i]                    # (B,)
                    num_valid = vmask_curr.sum(dim=1)                     # (B,)
                    if int(((num_valid == 0) & obj_valid_b).sum().item()) > 0:
                        raise RuntimeError(f"Zero valid vertices for rollout at mask_idx={_mask_idx}, obj={i} for valid samples")
                    need = (num_valid < Nq) & obj_valid_b
                    if bool(need.any()):
                        min_valid = int(num_valid[obj_valid_b].min().item()) if bool(obj_valid_b.any()) else 0
                        raise RuntimeError(f"Not enough valid vertices for rollout at mask_idx={_mask_idx}, obj={i}: min_valid={min_valid} < anchor_k={Nq}")
                    bad_mask = (~mask_at_anchors) & obj_valid_b.unsqueeze(1)
                    if bool(bad_mask.any()):
                        bad = bad_mask.sum().item()
                        raise RuntimeError(f"Invalid anchors for rollout: {bad} indices invalid at mask_idx={_mask_idx}, obj={i}.")

            pbar = tqdm(range(remaining), desc='Rollout', disable=not show_pbar, leave=False)
            for step in pbar:
                rope_memo_reset()  # rotary angles are memoized within one step only (keeps graph capture self-contained)
                _capturing = use_cuda_graph and step == graph_capture_step
                if _capturing:
                    # Move cross-step carries into fixed-address buffers, then
                    # record this step. Kernels are only recorded during capture,
                    # so the step is materialized by an immediate replay below.
                    prev_pos_bt = prev_pos_bt.clone()
                    curr_pos_bt = curr_pos_bt.clone()
                    registration_refs = registration_refs.clone()
                    if prev_obj_layers_rollout is not None:
                        prev_obj_layers_rollout = [x.clone() for x in prev_obj_layers_rollout]
                    if prev_state_layers_rollout is not None:
                        prev_state_layers_rollout = [x.clone() for x in prev_state_layers_rollout]
                    torch.cuda.synchronize()
                    graph_obj = torch.cuda.CUDAGraph()
                    _graph_stream = torch.cuda.Stream()
                    _prev_stream = torch.cuda.current_stream()
                    torch.cuda.set_stream(_graph_stream)
                    # thread_local: NCCL watchdog / DataLoader threads must not
                    # poison the capture under DDP evaluation.
                    graph_obj.capture_begin(capture_error_mode='thread_local')
                # ========== STATE RESET (optional, to avoid train-test mismatch) ==========
                # When rollout_state_reset=True, reset state tokens at each step to match
                # training distribution (training only sees ~2 state updates per sequence).
                # This prevents state drift during long rollouts (e.g., 400 frames).
                if self._rollout_state_reset:
                    if self._no_global_state:
                        state_tokens = None
                        state_pos = None
                    else:
                        state_tokens, state_pos = self.state_repr.get_state_and_positions(B, device)

                # Compute velocities from last two frames
                vel_bt = curr_pos_bt - prev_pos_bt  # (B,Nobj,Nv,3)

                # Prepare far-away filler for invalid vertices in reference lists
                FAR = 1e6

                mask_idx_curr = min(warmup_frames - 1 + step, T0 - 1)  # current frame's mask

                # ========== VECTORIZED KNN FOR ALL OBJECTS AT ONCE ==========
                # curr_pos_bt: (B, Nobj, Nv, 3)
                # valid_vtx_mask_all[:, mask_idx_curr]: (B, Nobj, Nv)
                vmask_step = valid_vtx_mask_all[:, mask_idx_curr]  # (B, Nobj, Nv)

                # Compute KNN distance vectors for ALL objects simultaneously
                # Only computes for valid objects and valid vertices!
                if self._knn_packed:
                    dist_vec_all = compute_knn_packed(curr_pos_bt, _knn_packs_r[mask_idx_curr])
                else:
                    dist_vec_all = compute_knn_all_objects_vectorized(
                        curr_pos_bt, vmask_step, FAR, valid_obj_mask=valid_obj_mask,
                        num_groups=self._knn_num_groups, active_slots=_knn_active_slots_r
                    )

                # Apply plane constraint adjustment (vectorized)
                # Use reshape() instead of view() to handle non-contiguous tensors
                curr_pos_flat = curr_pos_bt.reshape(B, Nobj * Nv, 3)
                dist_vec_flat = dist_vec_all.reshape(B, Nobj * Nv, 3)
                dist_vec2_flat = distance_2_plane(curr_pos_flat, [0, 0, 0], "z", dist_vec_flat)
                dist_vec2_bt = dist_vec2_flat.reshape(B, Nobj, Nv, 3)

                # Store raw dist_vec2 for visualization (only if return_extras=True)
                if return_extras:
                    all_dist_vec2_seq.append(dist_vec2_bt)

                # Relative position (vectorized): curr_pos - ref_pose
                relpos_cat = curr_pos_bt - ref_pose  # (B, Nobj, Nv, 3)

                # Velocity (already in correct shape)
                vel_cat = vel_bt  # (B, Nobj, Nv, 3)

                # Dist vectors
                distvec_cat = dist_vec2_bt  # (B, Nobj, Nv, 3)

                # Get current step's vertex mask for all objects (vectorized)
                mask_idx_curr_all = mask_idx_curr
                vmask_curr_all = valid_vtx_mask_all[:, mask_idx_curr_all]  # (B,Nobj,Nv)
                vmask_curr_all_3d = vmask_curr_all.unsqueeze(-1)  # (B,Nobj,Nv,1)

                # Apply masks BEFORE normalization (vectorized across all objects)
                distvec_cat = torch.where(vmask_curr_all_3d, distvec_cat, torch.zeros_like(distvec_cat))
                vel_cat = torch.where(vmask_curr_all_3d, vel_cat, torch.zeros_like(vel_cat))
                relpos_cat = torch.where(vmask_curr_all_3d, relpos_cat, torch.zeros_like(relpos_cat))

                # Apply network-predicted directional sign to ALL vertices' dist_vec (before normalization)
                # Only for step>0 (student forcing steps with predicted positions, may have penetration)
                # For step=0, use GT positions from warmup (no penetration, sign=+1)
                penetration_enabled = self.config['model'].get('if_penetration', self.config['model'].get('if_directional_distance', False))
                if penetration_enabled and step > 0:
                    # Predict penetration sign using PREVIOUS step's features and CURRENT predicted positions
                    # curr_pos_bt = blended_next (AFTER Kabsch + intersection handling)
                    pen_pred_all_vertices = self.predictor.predict_penetration_vertices(
                        prev_obj_layers_rollout,  # Use previous step's decoder outputs
                        prev_state_layers_rollout,  # Use previous step's state features
                        pos_full=curr_pos_bt  # AFTER Kabsch transform (matches training forward)
                    )  # (B, Nobj, Nv, 1) in [-1,1]

                    # Apply predicted sign to dist_vec for ALL vertices
                    pen_sign_all = torch.sign(pen_pred_all_vertices + 1e-6)  # (B, Nobj, Nv, 1) in {-1,+1}
                    distvec_cat = distvec_cat * pen_sign_all  # (B, Nobj, Nv, 3)

                # No normalization — raw values, matching the forward path
                # (features are masked once, then the penetration sign is applied)

                # Broadcast physics parameters and time_step to vertex dimension if enabled (rollout)
                phys_broadcast = None  # (B, Nobj, Nv, 3)
                time_broadcast = None  # (B, Nobj, Nv, 2)

                if self._if_phys_param:
                    # Broadcast per-object physics params (B,Nobj,3) to all vertices (B,Nobj,Nv,3)
                    phys_broadcast = phys_params.unsqueeze(2).expand(B, Nobj, Nv, 3)  # (B,Nobj,Nv,3)
                    # Mask with vertex mask to ensure padding vertices get zero physics params
                    phys_broadcast = torch.where(vmask_curr_all_3d, phys_broadcast, torch.zeros_like(phys_broadcast))

                if self._if_time_step:
                    # Broadcast both s and s² to all objects and vertices (B,Nobj,Nv,2)
                    # Must match the forward path
                    step_size_s = skip_step_size_batch.float()  # (B,)
                    step_size_sq = (skip_step_size_batch ** 2).float()  # (B,)
                    time_features = torch.stack([step_size_s, step_size_sq], dim=-1)  # (B, 2)
                    time_broadcast = time_features.view(B, 1, 1, 2).expand(B, Nobj, Nv, 2)  # (B,Nobj,Nv,2)
                    # Mask for consistency
                    # vmask_curr_all_3d is (B,Nobj,Nv,1), expand to match (B,Nobj,Nv,2)
                    vmask_time = vmask_curr_all_3d.expand(-1, -1, -1, 2)  # (B,Nobj,Nv,2)
                    time_broadcast = torch.where(vmask_time, time_broadcast, torch.zeros_like(time_broadcast))

                # PointNet Encoder: process each object's Nv vertices -> per-object global feature
                # Input dimension depends on flags:
                #   Base: 9D (dist_vec + velocity + rel_pos)
                #   +if_phys_param=true: +3D
                #   +if_time_step=true (no FiLM): +2D (s and s²)
                # Split into components and reshape to (B*Nobj, Nv, D) for encoder
                dist_reshaped = distvec_cat.view(B * Nobj, Nv, 3)
                vel_reshaped = vel_cat.view(B * Nobj, Nv, 3)
                relpos_reshaped = relpos_cat.view(B * Nobj, Nv, 3)

                # Optional features

                phys_reshaped = None
                if self._if_phys_param:
                    phys_reshaped = phys_broadcast.view(B * Nobj, Nv, 3)

                time_reshaped = None
                if self._if_time_step and not self._film_conditioning:
                    time_reshaped = time_broadcast.view(B * Nobj, Nv, 2)  # (s, s²)

                # Prepare vertex mask for encoder pooling
                # vmask_curr_all_3d is (B,Nobj,Nv,1), need (B*Nobj, Nv)
                vtx_mask_for_encoder = vmask_curr_all_3d.squeeze(-1).reshape(B * Nobj, Nv)  # (B*Nobj, Nv)
                height_reshaped = None
                if self.shared_obj_encoder.height_input:
                    _h = torch.where(vmask_curr_all_3d, curr_pos_bt[..., 2:3], torch.zeros_like(curr_pos_bt[..., 2:3]))
                    height_reshaped = _h.reshape(B * Nobj, Nv, 1)
                rank_reshaped = None
                if self.shared_obj_encoder.hierarchical_subsets:
                    if 'vtx_rank' not in initial_batch:
                        raise KeyError("model.encoder.hierarchical_subsets requires data.fps_order: true (batch lacks 'vtx_rank')")
                    rank_reshaped = initial_batch['vtx_rank'].reshape(B * Nobj, Nv)

                # Call encoder with all features (encoder will concatenate them internally)
                # AVP: capture per-point features if AVP is enabled
                obj_features_encoded, _per_vtx_feat, _ = self.shared_obj_encoder(
                    dist_reshaped,
                    vel_reshaped,
                    relpos_reshaped,
                        phys_params=phys_reshaped,
                    time_step=time_reshaped,
                    valid_mask=vtx_mask_for_encoder,
                    height=height_reshaped,
                    vtx_rank=rank_reshaped,
                    pack=_enc_packs_r.get(mask_idx_curr_all) if self._encoder_packed else None
                )  # obj_features_encoded: (B*Nobj, enc_out_dim), _per_vtx_feat: (B*Nobj, backbone_dim, Nv)

                # AVP: Anchor-Vertex Pooling — enrich anchor queries with collision-aware features
                _avp_anchor_features = None  # Will be (B, Nobj, Nq, avp_output_dim) if AVP enabled
                if self.avp is not None and _per_vtx_feat is not None:
                    # _per_vtx_feat is (B*Nobj, backbone_dim, Nv) — Conv1d channels-first format
                    # Transpose to (B*Nobj, Nv, backbone_dim) for AVP
                    _pvf = _per_vtx_feat.transpose(1, 2)  # (B*Nobj, Nv, backbone_dim)
                    # Anchor positions in flat format
                    _curr_pos_flat = curr_pos_bt.reshape(B * Nobj, Nv, 3)  # (B*Nobj, Nv, 3)
                    _anchor_pos_flat = torch.gather(
                        _curr_pos_flat, dim=1,
                        index=anchor_idx.reshape(B * Nobj, -1).unsqueeze(-1).expand(-1, -1, 3)
                    )  # (B*Nobj, Nq, 3)
                    _avp_anchor_features = self.avp(
                        _pvf, _curr_pos_flat, _anchor_pos_flat, vtx_mask_for_encoder
                    )  # (B*Nobj, Nq, avp_output_dim)
                    _avp_anchor_features = _avp_anchor_features.view(B, Nobj, anchor_idx.shape[2], -1)  # (B, Nobj, Nq, avp_output_dim)
                    # Zero out padding objects (proj bias makes padding non-zero; consistent with other features)
                    _avp_anchor_features = _avp_anchor_features * valid_obj_mask.unsqueeze(-1).unsqueeze(-1).float()

                obj_features_encoded = obj_features_encoded.view(B, Nobj, -1)  # (B, Nobj, enc_out_dim)

                # Project to token dimension for Decoder
                obj_features = self.obj_feature_proj(obj_features_encoded)  # (B, Nobj, D)

                # FiLM: modulate encoder features with step-size [s, s²]
                if self.encoder_film is not None and skip_step_size_batch is not None:
                    time_feat_film = torch.stack([
                        skip_step_size_batch.float(),
                        (skip_step_size_batch ** 2).float()
                    ], dim=-1)  # (B, 2)
                    obj_features = self.encoder_film(obj_features, time_feat_film)  # (B, Nobj, D)

                # Object anchor positions for RoPE
                # Same approach as forward(): gather anchors, flatten to multi-point for RoPE.
                # Replace padding objects' positions with zeros before the gather:
                # padding objects have anchor_idx=0, which would gather arbitrary values into RoPE inputs.
                obj_mask_pos = valid_obj_mask.unsqueeze(-1).unsqueeze(-1)  # (B, Nobj, 1, 1)
                safe_curr_pos_bt = torch.where(obj_mask_pos, curr_pos_bt, torch.zeros_like(curr_pos_bt))

                idx_exp_3_for_decoder = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B, Nobj, Nq, 3)
                obj_anchor_positions = torch.gather(safe_curr_pos_bt, dim=2, index=idx_exp_3_for_decoder)  # (B, Nobj, Nq, 3)
                # Flatten to (B, Nobj, Nq*3) for multi-point RoPE
                Nq = anchor_idx.shape[2]  # Number of anchors per object

                # ============== PE Type Routing (rollout) ==============
                # valid_vtx_mask_t is needed for geometry encoder
                valid_vtx_mask_t = valid_vtx_mask_all[:, mask_idx_curr]  # (B, Nobj, Nv)
                if self._pe_type in ('arope', 'se3'):
                    # SE(3) uses arope anchors for cross-attention RoPE; SE(3) bias computed separately
                    obj_anchor_flat = obj_anchor_positions.view(B, Nobj, Nq * 3)  # (B, Nobj, 12)
                elif self._pe_type in ('obb', 'pca'):
                    obj_anchor_flat = self._geometry_encoder(
                        safe_curr_pos_bt, valid_vtx_mask_t, obj_mask=valid_obj_mask
                    )  # (B, Nobj, 12)
                else:
                    raise ValueError(f"Unknown pe_type='{self._pe_type}' in rollout().")

                # ============== SE(3) Pairwise Bias Computation (rollout) ==============
                _se3_bias = None
                if self._pe_type == 'se3' and self._se3_bias_module is not None:
                    from src.models.se3_attention import extract_object_pose_features
                    _rotmats, _centers, _scales = extract_object_pose_features(
                        safe_curr_pos_bt, ref_pose, valid_vtx_mask_t, valid_obj_mask
                    )
                    _se3_bias = self._se3_bias_module(_rotmats, _centers, _scales, valid_obj_mask)

                # Decoder: SDPA additive mask + feature zeroing (same as forward)
                obj_attn_mask = torch.where(
                    valid_obj_mask.unsqueeze(1).unsqueeze(1),
                    _attn_zero,
                    _attn_ninf
                )
                pad_mask_float = valid_obj_mask.unsqueeze(-1).to(dtype=obj_features.dtype)
                obj_features = obj_features * pad_mask_float
                obj_anchor_flat = obj_anchor_flat * valid_obj_mask.unsqueeze(-1).to(dtype=obj_anchor_flat.dtype)

                # Per-layer FiLM: pass time_features for step-size conditioning in each layer
                _decoder_time_feat = time_feat_film if (self._film_conditioning and self._if_time_step) else None
                all_state_layers, all_obj_layers = self.decoder(
                    state_tokens,
                    obj_features,
                    state_pos,
                    obj_anchor_flat,
                    mask=obj_attn_mask,
                    se3_bias=_se3_bias,
                    time_features=_decoder_time_feat  # Per-layer FiLM
                )
                # No need to remove any tokens - all_obj_layers directly contains object features

                # Update state (skip when no_global_state)
                if not self._no_global_state:
                    state_tokens = all_state_layers[-1]

                # Build anchor query for Predictor (sample at anchor points Nq)
                # Query dimension depends on flags:
                #   Base: 9D (dist_vec + vel + rel_pos)
                #   +if_phys_param=true: +3D = 12D
                #   +if_time_step=true (no FiLM): +2D (s and s²)
                idx_exp_3 = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
                distvec_at_anchors = torch.gather(distvec_cat, dim=2, index=idx_exp_3)
                vel_at_anchors = torch.gather(vel_cat, dim=2, index=idx_exp_3)
                relpos_at_anchors = torch.gather(relpos_cat, dim=2, index=idx_exp_3)

                # Build query features list
                query_features = [distvec_at_anchors, vel_at_anchors, relpos_at_anchors]

                # Add physics parameters at anchors if enabled
                if self._if_phys_param:
                    phys_at_anchors = torch.gather(phys_broadcast, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)
                    query_features.append(phys_at_anchors)

                # Add time step at anchors if enabled AND not using FiLM
                if self._if_time_step and not self._film_conditioning:
                    idx_exp_2 = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 2)  # (B,Nobj,Nq,2)
                    time_at_anchors = torch.gather(time_broadcast, dim=2, index=idx_exp_2)  # (B,Nobj,Nq,2)
                    query_features.append(time_at_anchors)

                # AVP: Append collision-aware anchor features to query
                if _avp_anchor_features is not None:
                    query_features.append(_avp_anchor_features)  # (B, Nobj, Nq, avp_output_dim)

                # Concatenate all query features
                anchor_query = torch.cat(query_features, dim=-1)  # (B,Nobj,Nq, 9/12 + avp_dim)

                # Extract anchor positions for RoPE (current frame positions at anchors)
                anchor_pos_at_curr = torch.gather(curr_pos_bt, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)

                # Penetration sign was already applied to all vertices' dist_vec;
                # the anchor query therefore carries the sign information as well.

                _ref_src = _object_centered(ref_pose, vmask_step.bool()) if self._ref_pose_object_centered else ref_pose
                if self._ref_pose_object_centered and not getattr(self, '_refpose_logged_roll', False):
                    self._refpose_logged_roll = True
                    print(f"[RefPose] object-centred (x/y only, z absolute) reference pose ACTIVE (rollout): |shift| max={((ref_pose - _ref_src).abs().max()).item():.3f} m", flush=True)
                ref_pose_at_anchors = torch.gather(_ref_src, dim=2, index=idx_exp_3)  # (B,Nobj,Nq,3)

                # Predict accelerations at anchors using grouped format (B, Nobj, Nq, *)
                accel_norm_anchor = self.predictor(
                    all_obj_layers,
                    anchor_query,           # (B, Nobj, Nq, 9/12)
                    ref_pose=ref_pose_at_anchors,  # (B, Nobj, Nq, 3)
                    query_pos=anchor_pos_at_curr,  # (B, Nobj, Nq, 3)
                    key_pos=obj_anchor_flat,        # (B, Nobj, Nq*3) - multi-point RoPE
                    multi_scale_state_features=None if self._no_global_state else all_state_layers,
                    state_pos=None if self._no_global_state else state_pos,
                    obj_mask=valid_obj_mask  # Object validity mask for padding
                )  # Returns (B, Nobj, Nq, 3)

                # Force FP32 for Verlet integration and position computation
                # (must match the training forward path for numerical stability)
                with torch.amp.autocast('cuda', enabled=False):
                    accel_base_denorm = accel_norm_anchor.float()

                    # Explicit step_size² scaling for multi-timescale rollout
                    # (must match the forward path)
                    if self._if_time_step and skip_step_size_batch is not None:
                        step_size_sq = (skip_step_size_batch.float() ** 2).view(B, 1, 1, 1)  # (B,1,1,1) FP32
                        accel_scaled_denorm = accel_base_denorm * step_size_sq  # FP32 * FP32 = FP32
                    else:
                        accel_scaled_denorm = accel_base_denorm

                    # Store scaled acceleration for visualization (matches forward behavior)
                    if return_extras:
                        all_accel_anchor_denorm_seq.append(accel_scaled_denorm)

                    # Verlet integration on anchors; initialize full next positions as zero-accel extrapolation
                    dt_squared = self.dt_verlet * self.dt_verlet
                    # Positions must be FP32 for numerical stability
                    curr_pos_bt_f = curr_pos_bt.float()
                    prev_pos_bt_f = prev_pos_bt.float()
                    next_pos_bt = 2 * curr_pos_bt_f - prev_pos_bt_f  # (B,Nobj,Nv,3) - FP32

                    # Vectorized scatter using scatter_reduce (same optimization as forward)
                    # Step 1: Gather all anchor positions at once (use FP32 positions)
                    anchor_idx_exp = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
                    curr_anchors = torch.gather(curr_pos_bt_f, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                    prev_anchors = torch.gather(prev_pos_bt_f, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)

                    # Step 2: Compute next positions at anchors using Verlet (with scaled acceleration)
                    next_anchors = accel_scaled_denorm * dt_squared + 2 * curr_anchors - prev_anchors  # (B,Nobj,Nq,3)

                    # Step 3: Compute updates (delta from base prediction)
                    base_anchors = torch.gather(next_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                    anchor_updates = next_anchors - base_anchors  # (B,Nobj,Nq,3)

                    # Mask padding objects' updates to zero
                    # Padding objects have anchor_idx=0 which would scatter to wrong positions
                    obj_mask_scatter = valid_obj_mask.unsqueeze(-1).unsqueeze(-1).float()  # (B,Nobj,1,1)
                    anchor_updates = anchor_updates * obj_mask_scatter  # padding objects -> zero updates

                    # Step 4: Vectorized scatter using scatter_reduce
                    # Flatten batch and object dimensions for efficient scatter
                    # Use reshape() to handle non-contiguous tensors
                    B_Nobj_Nv = B * Nobj * Nv
                    next_pos_flat = next_pos_bt.reshape(B_Nobj_Nv, 3)  # (B*Nobj*Nv, 3)

                    # Create flat indices for scatter_reduce
                    obj_offsets = torch.arange(Nobj, device=anchor_idx.device).view(1, Nobj, 1) * Nv  # (1,Nobj,1)
                    batch_offsets = torch.arange(B, device=anchor_idx.device).view(B, 1, 1) * (Nobj * Nv)  # (B,1,1)
                    flat_anchor_idx = anchor_idx + obj_offsets + batch_offsets  # (B,Nobj,Nq)
                    flat_anchor_idx = flat_anchor_idx.reshape(B * Nobj * Nq)  # (B*Nobj*Nq,)

                    # Flatten updates
                    anchor_updates_flat = anchor_updates.reshape(B * Nobj * Nq, 3)  # (B*Nobj*Nq, 3)

                    # Apply updates using scatter_reduce (more efficient than loop)
                    next_pos_flat = torch.scatter_reduce(
                        next_pos_flat,
                        dim=0,
                        index=flat_anchor_idx.unsqueeze(-1).expand(-1, 3),
                        src=anchor_updates_flat,
                        reduce='sum',
                        include_self=True
                    )

                    # Reshape back
                    next_pos_bt = next_pos_flat.reshape(B, Nobj, Nv, 3)

                # Per-object rigid projection using anchors; apply to full Nv
                # Vectorized Kabsch: Process ONLY valid objects (filter padding BEFORE SVD)
                # Step 1: Gather anchor positions for all objects at once
                anchor_idx_exp = anchor_idx.unsqueeze(-1).expand(-1, -1, -1, 3)  # (B,Nobj,Nq,3)
                last_anchor_all = torch.gather(registration_refs, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)
                curr_anchor_all = torch.gather(next_pos_bt, dim=2, index=anchor_idx_exp)  # (B,Nobj,Nq,3)

                # Reshape to (B*Nobj, Nq, 3) for batch processing
                B_Nobj = B * Nobj
                last_anchor_flat = last_anchor_all.view(B_Nobj, Nq, 3)
                curr_anchor_flat = curr_anchor_all.view(B_Nobj, Nq, 3)

                # Filter padding objects before SVD; they would produce garbage R/t
                # (indices precomputed before the loop; the mask never changes)
                valid_idx = _kabsch_valid_idx
                n_valid = _kabsch_n_valid

                # Initialize R/t for ALL objects (identity R, zero t)
                # FP32 so the RoMa FP32 output is not downcast during scatter
                R_flat = torch.eye(3, device=device, dtype=torch.float32).unsqueeze(0).expand(B_Nobj, -1, -1).clone()
                t_vec_flat = torch.zeros(B_Nobj, 1, 3, device=device, dtype=torch.float32)

                if n_valid > 0:
                    # Gather only valid objects' anchors
                    last_anchor_valid = last_anchor_flat[valid_idx]  # (Nvalid, Nq, 3)
                    curr_anchor_valid = curr_anchor_flat[valid_idx]  # (Nvalid, Nq, 3)

                    if _capturing:
                        # The SVD inside the rigid solve is not graph-capturable, so
                        # the step is recorded as two graphs around an eager solve.
                        # Close graph 1 here and execute it to materialize its outputs.
                        graph_obj.capture_end()
                        graph_obj.replay()
                        torch.cuda.current_stream().synchronize()

                    if self.if_roma:
                        # RoMa mode: batch process only valid objects
                        # RoMa requires FP32 for numerical stability
                        with torch.amp.autocast('cuda', enabled=False):
                            R_valid, t_vec_valid = roma.rigid_points_registration(
                                last_anchor_valid.float(), 
                                curr_anchor_valid.float()
                            )
                        R_valid = R_valid.float()  # Ensure FP32 for downstream
                        t_vec_valid = t_vec_valid.float().unsqueeze(1)  # (Nvalid, 1, 3)

                    else:
                        # Standard SVD Kabsch: batch process only valid objects
                        centroid_last = last_anchor_valid.mean(dim=1, keepdim=True)  # (Nvalid,1,3)
                        centroid_curr = curr_anchor_valid.mean(dim=1, keepdim=True)  # (Nvalid,1,3)
                        last_c = last_anchor_valid - centroid_last
                        curr_c = curr_anchor_valid - centroid_curr

                        # Batch SVD for valid objects only
                        H = torch.bmm(last_c.transpose(1, 2), curr_c)  # (Nvalid,3,3)
                        with torch.amp.autocast('cuda', enabled=False):
                            H_float = H.float()
                            U, S, Vt = torch.linalg.svd(H_float)  # Batch SVD!
                        U = U.to(H.dtype); Vt = Vt.to(H.dtype)
                        R_valid = torch.bmm(Vt.transpose(1, 2), U.transpose(1, 2))  # (Nvalid,3,3)

                        # Fix reflections
                        with torch.amp.autocast('cuda', enabled=False):
                            det_R = torch.det(R_valid.float()).to(R_valid.dtype)
                        need_flip = det_R < 0
                        if need_flip.any():
                            Fm = torch.eye(3, device=R_valid.device, dtype=R_valid.dtype).unsqueeze(0).expand(n_valid, -1, -1).clone()
                            Fm[need_flip, 2, 2] = -1
                            R_valid = torch.bmm(torch.bmm(Vt.transpose(1, 2), Fm), U.transpose(1, 2))

                        # Translation
                        rotated_centroid = torch.bmm(R_valid, centroid_last.transpose(1, 2)).transpose(1, 2)
                        t_vec_valid = centroid_curr - rotated_centroid  # (Nvalid,1,3)

                    if _capturing:
                        # Fixed-address buffers carry the solve result into graph 2.
                        _rt_R_buf = R_valid.clone()
                        _rt_t_buf = t_vec_valid.clone()
                        R_valid = _rt_R_buf
                        t_vec_valid = _rt_t_buf
                        graph_obj2 = torch.cuda.CUDAGraph()
                        graph_obj2.capture_begin(pool=graph_obj.pool(), capture_error_mode='thread_local')
                    elif graph_obj2 is not None:
                        _rt_R_buf.copy_(R_valid)
                        _rt_t_buf.copy_(t_vec_valid)
                        R_valid = _rt_R_buf
                        t_vec_valid = _rt_t_buf

                    # Scatter valid R/t back to full tensors
                    R_flat[valid_idx] = R_valid
                    t_vec_flat[valid_idx] = t_vec_valid

                # Reshape back to (B,Nobj,3,3) and (B,Nobj,1,3)
                R_all = R_flat.view(B, Nobj, 3, 3)
                t_vec_all = t_vec_flat.view(B, Nobj, 1, 3)

                # ========== VECTORIZED TRANSFORMATION APPLICATION ==========
                # Apply R, t to all objects in a single bmm call (no per-object loop)
                # registration_refs: (B, Nobj, Nv, 3)
                # R_all: (B, Nobj, 3, 3)
                # t_vec_all: (B, Nobj, 1, 3)

                # Reshape for batch bmm: (B, Nobj, Nv, 3) -> (B*Nobj, Nv, 3)
                last_flat = registration_refs.view(B * Nobj, Nv, 3)
                R_flat = R_all.view(B * Nobj, 3, 3)
                t_flat = t_vec_all.view(B * Nobj, 1, 3)

                # Single bmm for ALL objects: y = R @ x + t (RoMa convention)
                # Convert to FP32 to prevent BF16 precision loss during autoregressive rollout
                y_flat = torch.bmm(R_flat, last_flat.transpose(1, 2)).transpose(1, 2) + t_flat
                # Use reshape() because transpose() makes tensor non-contiguous
                yproj_bt = y_flat.reshape(B, Nobj, Nv, 3).float()  # FP32 for numerical stability

                # Apply gradient strategy based on transform type
                if self.if_roma:
                    # RoMa mode: fully differentiable, gradients flow through Kabsch
                    blended_next = yproj_bt
                else:
                    # Standard mode: straight-through estimator (STE)
                    blended_next = next_pos_bt + (yproj_bt - next_pos_bt).detach()

                # ========== INTERSECTION HANDLING (EVAL) ==========
                intersection_cfg = self.config.get('model', {}).get('intersection_handling', {})
                blended_next = handle_intersections(
                    blended_next, valid_obj_mask, vmask_step, intersection_cfg
                )

                if _capturing:
                    # Rotate the carry buffers in place inside the captured region
                    # so every replay advances the state by one step.
                    prev_pos_bt.copy_(curr_pos_bt)
                    curr_pos_bt.copy_(blended_next)
                    if not self._kabsch_reference_fixed:
                        registration_refs.copy_(blended_next)
                    if prev_obj_layers_rollout is not None:
                        for _buf, _new in zip(prev_obj_layers_rollout, all_obj_layers):
                            _buf.copy_(_new)
                    if prev_state_layers_rollout is not None:
                        for _buf, _new in zip(prev_state_layers_rollout, all_state_layers):
                            _buf.copy_(_new)
                    graph_obj2.capture_end()
                    torch.cuda.set_stream(_prev_stream)
                    torch.cuda.synchronize()
                    graph_obj2.replay()  # capture only records; execute part 2 now
                    torch.cuda.synchronize()
                    out_seq.append(curr_pos_bt.clone())
                    pbar.close()
                    break
                # Registration refs use the post-collision position, matching the forward path ('advancing');
                # model.kabsch_reference 'fixed' keeps the last warm-up frame as the reference for every step
                if not self._kabsch_reference_fixed:
                    registration_refs = blended_next.detach()

                out_seq.append(blended_next)

                # Save current features for next step's penetration prediction
                prev_obj_layers_rollout = all_obj_layers
                prev_state_layers_rollout = all_state_layers

                prev_pos_bt = curr_pos_bt
                curr_pos_bt = blended_next

            if graph_obj is not None:
                target_len = warmup_frames + remaining
                while len(out_seq) < target_len:
                    graph_obj.replay()
                    with torch.amp.autocast('cuda', enabled=False):
                        _R_new, _t_new = roma.rigid_points_registration(
                            last_anchor_valid.float(),
                            curr_anchor_valid.float()
                        )
                    _rt_R_buf.copy_(_R_new.float())
                    _rt_t_buf.copy_(_t_new.float().unsqueeze(1))
                    graph_obj2.replay()
                    out_seq.append(curr_pos_bt.clone())

            if graph_obj is not None:
                # Tear the graphs down and return their private pools to the
                # allocator: the evaluator calls rollout once per batch and per
                # step size, and retired pools would otherwise accumulate.
                graph_obj.reset()
                if graph_obj2 is not None:
                    graph_obj2.reset()
                graph_obj = None
                graph_obj2 = None
                if os.environ.get('RIGIDFORMER_GRAPH_KEEPCACHE', '0') != '1':
                    torch.cuda.empty_cache()

            full = torch.stack(out_seq, dim=1)  # (B, 2+remaining, Nobj, Nv, 3)

            # Single numerical-health check per rollout (one host sync total, not
            # one per step): any non-finite prediction aborts loudly here.
            if not torch.isfinite(full).all():
                raise RuntimeError("Non-finite values in rollout output; "
                                   "inspect the checkpoint and input data.")

            # ========== Map back to original absolute space ==========
            # If center_scene_by_t0 was applied, add center_t0 back to predictions
            # so that the output is in the original world coordinate frame;
            # centering remains a purely internal operation of the model.
            if self.config.get('data', {}).get('center_scene_by_t0', False) and len(_center_t0_per_sample) > 0:
                for b in range(B):
                    if _center_t0_per_sample[b] is not None:
                        center_t0_tensor = torch.from_numpy(_center_t0_per_sample[b]).to(full.device)
                        full[b] = full[b] + center_t0_tensor
                logger.info(f"[Center-T0 ROLLOUT] Mapped predictions back to original absolute space")

            output = {
                'obj_positions': full
            }

            # Return body_dist_vec2 and accelerations only if requested (for visualization)
            # Skip during RMSE eval to save ~4GB memory on long rollouts
            if return_extras and remaining > 0:
                output['obj_body_dist_vec2'] = torch.stack(all_dist_vec2_seq, dim=1)
                output['obj_accelerations_anchor'] = torch.stack(all_accel_anchor_denorm_seq, dim=1)

            return output

            


class Trainer:
    """Training manager for the model."""
    
    def __init__(self, config, model, train_loader, val_loader, device='cuda', viz_params=None):
        self.config = config
        self.viz_params = viz_params or {}

        # Fast fail: start the video renderer if video saving is enabled.
        # Must run before any CUDA context is created (e.g. model.to(device))
        # to avoid EGL/CUDA memory conflicts.
        if self.config['training']['checkpoint'].get('save_video', False):
            # Only main process needs visualization
            if not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0:
                if self.config['training']['checkpoint'].get('video_renderer', 'filament') == 'polyscope':
                    self._polyscope_self_test()
                else:
                    self._filament_renderer()
        # Avoid moving DDP-wrapped model again
        if isinstance(model, DDP):
            self.model = model
        else:
            # Try to move model to device with OOM handling (only in inference mode)
            try:
                self.model = model.to(device)
                logger.info(f"Model moved to {device}")
            except torch.cuda.OutOfMemoryError as e:
                # Check if we're in inference-only mode (will be set later, so check train_loader)
                is_inference = len(train_loader) == 0
                if is_inference:
                    logger.error(f"CUDA OOM when moving model to GPU: {e}")
                    logger.warning("Inference-only mode: falling back to CPU execution")
                    device = 'cpu'
                    torch.cuda.empty_cache()
                    self.model = model.to(device)
                    logger.info("Model moved to CPU (execution will be slower)")
                else:
                    logger.error("CUDA OOM during training initialization")
                    logger.error("Reduce batch size or model size, or use a GPU with more memory")
                    raise
            except RuntimeError as e:
                if 'out of memory' in str(e).lower():
                    is_inference = len(train_loader) == 0
                    if is_inference:
                        logger.error(f"GPU memory error: {e}")
                        logger.warning("Inference-only mode: falling back to CPU execution")
                        device = 'cpu'
                        torch.cuda.empty_cache()
                        self.model = model.to(device)
                        logger.info("Model moved to CPU (execution will be slower)")
                    else:
                        logger.error("GPU OOM during training initialization")
                        logger.error("Reduce batch size or model size, or use a GPU with more memory")
                        raise
                else:
                    raise
        self.device = device
        self.train_loader = train_loader
        self.val_loader = val_loader
        # Cache for dataset mesh faces to avoid repeated disk IO
        self._faces_cache = {}
        # Cache for video datasets (per step_size) to avoid repeated indexing
        # Key: step_size, Value: MultiObjectDynamics dataset
        self._test_video_datasets = {}   # Test split: {step_size: dataset}
        self._train_video_datasets = {}  # Train split: {step_size: dataset}
        self._val_video_datasets = {}    # Val split: {step_size: dataset}

        # Read validation enable flag from config
        self.enable_validation = self.config['training'].get('enable_validation', True)
        if not self.enable_validation:
            logger.info("Validation disabled by config (training.enable_validation=false)")

        # Detect inference-only mode (empty train_loader)
        self.inference_only = len(train_loader) == 0
        if self.inference_only:
            logger.info("Inference-only mode detected (empty train_loader), skipping optimizer/scheduler setup")

        # Setup optimizer based on config type
        if not self.inference_only:
            optimizer_type = config['training']['optimizer'].get('type', None)
            if optimizer_type is None:
                raise ValueError(
                    "Optimizer type must be specified in training.optimizer.type\n"
                    "Supported: 'AdamW'"
                )
        else:
            optimizer_type = None

        if optimizer_type is None:
            # Inference-only mode, optimizer setup will be skipped
            pass
        elif optimizer_type.lower() == 'adamw':
            # Standard AdamW optimizer (matching baselines)
            opt_cfg = config['training']['optimizer']
            lr = float(opt_cfg.get('lr', 1e-4))
            betas = tuple(opt_cfg.get('betas', [0.9, 0.999]))
            weight_decay = float(opt_cfg.get('weight_decay', 0.01))
            # Adam's eps (torch default 1e-8). The anchor-space losses are in metres per frame^2, so gradients can
            # be small enough for eps to damp the step; training.optimizer.eps makes it a recorded per-run choice.
            adam_eps = float(opt_cfg.get('eps', 1e-8))

            # Split params: fusion layers may use a higher learning rate
            fusion_params = []
            other_params = []
            fusion_param_names = []
            # encoder_lr_mult > 1 gives the point encoder its own learning rate; used when
            # fine-tuning a checkpoint after changing the encoder's pooling semantics.
            enc_params = []
            enc_lr_mult = float(opt_cfg.get('encoder_lr_mult', 1.0))
            # decoder_lr_mult > 1 does the same for the Transformer decoder and the anchor predictor
            # (the modules that consume the rotary embedding).
            dec_params = []
            dec_lr_mult = float(opt_cfg.get('decoder_lr_mult', 1.0))
            for name, param in model.named_parameters():
                if not param.requires_grad:
                    continue
                if 'predictor.fusion_norm1' in name or 'predictor.fusion_w1' in name or 'predictor.fusion_g1' in name:
                    fusion_params.append(param)
                    fusion_param_names.append(name)
                elif enc_lr_mult != 1.0 and 'shared_obj_encoder' in name:
                    enc_params.append(param)
                elif dec_lr_mult != 1.0 and (name.startswith('decoder.') or name.startswith('predictor.')
                                             or name.startswith('module.decoder.') or name.startswith('module.predictor.')):
                    dec_params.append(param)
                else:
                    other_params.append(param)

            fusion_lr_mult = float(opt_cfg.get('fusion_lr_mult', 10.0))
            fusion_lr = lr * fusion_lr_mult

            # decay_matrices_only splits the groups, so it takes this path also when no group differs in lr (e.g. the
            # fusion layers frozen with training.optimizer.frozen_params)
            if fusion_params or enc_params or dec_params or bool(opt_cfg.get('decay_matrices_only', False)):
                param_groups = [
                    {'params': other_params, 'lr': lr},
                ]
                if fusion_params:
                    param_groups.append({'params': fusion_params, 'lr': fusion_lr})
                if enc_params:
                    param_groups.append({'params': enc_params, 'lr': lr * enc_lr_mult})
                    logger.info(f"Encoder parameter group: {len(enc_params)} tensors at lr={lr * enc_lr_mult:.2e} (x{enc_lr_mult})")
                if dec_params:
                    param_groups.append({'params': dec_params, 'lr': lr * dec_lr_mult})
                    logger.info(f"Decoder/predictor parameter group: {len(dec_params)} tensors at lr={lr * dec_lr_mult:.2e} (x{dec_lr_mult})")
                # training.optimizer.decay_matrices_only: weight decay on weight matrices / convolution kernels only; norm
                # gains, biases, scalars (e.g. the anchor-vertex pooling log-radius) and the learned register-token
                # embeddings are not decayed. Every group is split in two with the same learning rate.
                if bool(opt_cfg.get('decay_matrices_only', False)):
                    _split = []
                    _embed = {id(q) for n_, q in model.named_parameters() if 'register_tokens' in n_}
                    for _g in param_groups:
                        _mat = [q for q in _g['params'] if q.ndim >= 2 and id(q) not in _embed]
                        _vec = [q for q in _g['params'] if q.ndim < 2 or id(q) in _embed]
                        _split.append({'params': _mat, 'lr': _g['lr']})
                        if _vec:
                            _split.append({'params': _vec, 'lr': _g['lr'], 'weight_decay': 0.0})
                    param_groups = _split
                    logger.info(f"[Optimizer] weight decay {weight_decay} on {sum(len(g['params']) for g in param_groups if 'weight_decay' not in g)} "
                                f"weight matrices only; {sum(len(g['params']) for g in param_groups if 'weight_decay' in g)} "
                                f"norm / bias / scalar tensors undecayed")
                self.optimizer = optim.AdamW(
                    param_groups,
                    betas=betas,
                    eps=adam_eps,
                    weight_decay=weight_decay
                )
                logger.info(f"Optimizer: AdamW with 2 groups (eps={adam_eps:g}):")
                logger.info(f"  - Base params ({len(other_params)}): lr={lr}")
                logger.info(f"  - Fusion params ({len(fusion_params)}): lr={fusion_lr} ({fusion_lr_mult}x)")
                for fn in fusion_param_names:
                    logger.info(f"    - {fn}")
            else:
                if bool(opt_cfg.get('decay_matrices_only', False)):
                    raise ValueError('training.optimizer.decay_matrices_only needs the grouped optimizer path')
                self.optimizer = optim.AdamW(
                    model.parameters(),
                    lr=lr,
                    betas=betas,
                    eps=adam_eps,
                    weight_decay=weight_decay
                )
                logger.info(f"Optimizer: AdamW (lr={lr}, betas={betas}, wd={weight_decay}, eps={adam_eps:g})")
        else:
            raise ValueError(
                f"Unsupported optimizer type: {optimizer_type}\n"
                "Supported: 'AdamW'"
            )
        
        # Setup scheduler based on config (skip in inference-only mode)
        if not self.inference_only:
            scheduler_type = config['training']['scheduler'].get('type', 'cosine')
            if scheduler_type == 'constant':
                # No learning rate scheduling - keep constant
                self.scheduler = optim.lr_scheduler.LambdaLR(
                    self.optimizer,
                    lr_lambda=lambda epoch: 1.0  # Keep learning rate constant
                )
            elif scheduler_type == 'cosine':
                T_max = config['training']['max_epochs']
                # Support both keys 'min_lr' and 'eta_min'
                eta_min = float(config['training']['scheduler'].get('eta_min', config['training']['scheduler'].get('min_lr', 1e-6)))
                warmup_epochs = int(config['training']['scheduler'].get('warmup_epochs', 0))
                
                if warmup_epochs > 0:
                    # Use start_factor directly (recommended: 0.1 = start from 10% of base_lr)
                    start_factor = float(config['training']['scheduler'].get('start_factor', 0.1))

                    logger.info(f"Initializing Scheduler with Warmup: {warmup_epochs} epochs, start_factor={start_factor}")
                    
                    try:
                        from torch.optim.lr_scheduler import LinearLR, SequentialLR
                        
                        # Warmup phase
                        warmup_scheduler = LinearLR(
                            self.optimizer, 
                            start_factor=start_factor, 
                            end_factor=1.0, 
                            total_iters=warmup_epochs
                        )
                        
                        # Cosine phase (starts after warmup)
                        # T_max for cosine is the remaining epochs
                        cosine_scheduler = optim.lr_scheduler.CosineAnnealingLR(
                            self.optimizer,
                            T_max=max(1, T_max - warmup_epochs),
                            eta_min=eta_min
                        )
                        
                        self.scheduler = SequentialLR(
                            self.optimizer, 
                            schedulers=[warmup_scheduler, cosine_scheduler], 
                            milestones=[warmup_epochs]
                        )
                    except ImportError:
                        logger.warning("SequentialLR/LinearLR not found (PyTorch < 1.10?). Fallback to simple CosineAnnealingLR without warmup.")
                        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
                            self.optimizer,
                            T_max=T_max,
                            eta_min=eta_min
                        )
                else:
                    # Standard Cosine without warmup
                    self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
                        self.optimizer,
                        T_max=T_max,
                        eta_min=eta_min
                    )
            else:
                # Default to cosine
                self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
                    self.optimizer,
                    T_max=config['training']['max_epochs'],
                    eta_min=float(config['training']['scheduler'].get('eta_min', 1e-6))
                )
        else:
            # Inference-only mode: no optimizer/scheduler needed
            self.optimizer = None
            self.scheduler = None
        
        # Setup loss functions
        # Use SmoothL1Loss (Huber loss) for more stable gradients
        self.position_loss = nn.SmoothL1Loss()  # Default beta=1.0
        self.velocity_loss = nn.SmoothL1Loss()
        self.acceleration_loss = nn.SmoothL1Loss()
        
        # Loss weights
        self.loss_weights = config['training']['loss_weights']

        # Learnable transform checks
        self.is_learnable_transform = (
            self.config.get('model', {}).get('transform', {}).get('method', 'kabsch') == 'learnable'
        )
        self.require_gt_motion = bool(self.config.get('data', {}).get('load_gt_motion', False))
        self._did_learnable_assert = False
        
        # Setup logging and checkpointing with same timestamp (allow override via ENV)
        timestamp = os.environ.get('EXP_TIMESTAMP', datetime.now().strftime('%Y%m%d_%H%M%S'))
        self.timestamp = timestamp  # Store for later use
        
        # Logging directory with optional experiment tag
        exp_tag = self.config.get('experiment_tag', None)
        log_base = Path(config['paths']['log_dir']) / exp_tag if exp_tag else Path(config['paths']['log_dir'])
        self.log_dir = log_base / f'run_{timestamp}'
        self.log_dir.mkdir(parents=True, exist_ok=True)

        # Only rank0 writes TensorBoard to avoid 8x event files in DDP
        def _is_main_proc():
            return (not (dist.is_available() and dist.is_initialized())) or dist.get_rank() == 0

        self.is_main = _is_main_proc()

        class _NoopWriter:
            def add_scalar(self, *args, **kwargs):
                return
            def add_text(self, *args, **kwargs):
                return
            def close(self):
                return

        if self.is_main and self.config.get('logging', {}).get('use_tensorboard', True):
            self.writer = SummaryWriter(self.log_dir)
        else:
            self.writer = _NoopWriter()
        
        # Optional file logging (disabled by default); only rank0 writes
        if self.is_main and self.config.get('logging', {}).get('file_logging', False):
            # Put file log under the run directory next to TensorBoard
            log_file_path = self.log_dir / 'training.log'
            file_handler = logging.FileHandler(log_file_path)
            file_handler.setLevel(logging.INFO)
            file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
            # Add to root logger so ALL modules (evaluator, dataloader, etc.) write to file
            logging.getLogger().addHandler(file_handler)
            logger.info(f"  - Log file: {log_file_path}")
        
        # Checkpoint directory with same timestamp and optional tag
        base_checkpoint_dir = Path(config['paths']['checkpoint_dir']) / exp_tag if exp_tag else Path(config['paths']['checkpoint_dir'])
        self.checkpoint_dir = base_checkpoint_dir / f'exp_{timestamp}'
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        
        # Log the experiment directories
        logger.info(f"Experiment {timestamp} directories:")
        logger.info(f"  - TensorBoard: {self.log_dir}")
        # (log file path printed only when file_logging enabled)
        logger.info(f"  - Checkpoints: {self.checkpoint_dir}")

        # Save config to log directory (run_XXXXXX) for reproducibility
        config_save_path = self.log_dir / 'config.yaml'
        with open(config_save_path, 'w') as f:
            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
        logger.info(f"  - Config saved to: {config_save_path}")
        
        self.best_val_loss = float('inf')
        self.epoch = 0
        # First epoch of the training loop. Stays 0 for a fresh run or a weights-only warm start;
        # set to (completed epoch + 1) when a full training checkpoint is resumed.
        self.start_epoch = 0

        # Optional exponential moving average of the weights (training.ema: {enable, decay}).
        # Rank 0 only: after DDP's all-reduce every rank holds identical weights. Off by default,
        # in which case training is unchanged.
        self.ema = None
        _ema_cfg = config['training'].get('ema', {}) or {}
        if bool(_ema_cfg.get('enable', False)) and not self.inference_only and self.is_main:
            _ema_base = self.model.module if isinstance(self.model, DDP) else self.model
            self.ema = WeightEMA(_ema_base, float(_ema_cfg.get('decay', 0.9999)))
            logger.info(f"[EMA] weight averaging ACTIVE on rank 0: decay={self.ema.decay} "
                        f"({len(self.ema.shadow)} state entries, fp32 shadow)")

        # GPU Augmentation (rotation) - only for training
        if not self.inference_only:
            from src.data.gpu_augmentation import GPUAugmentation
            self.gpu_augmentation = GPUAugmentation(config, self.device)
            logger.info(f"GPU augmentation initialized (rotation={self.gpu_augmentation.rotation_enabled}, "
                        f"reflection={self.gpu_augmentation.reflection_enabled}, rigid_noise={self.gpu_augmentation.noise_enabled})")
        else:
            self.gpu_augmentation = None

        # Scheduled Sampling configuration
        self.scheduled_sampling_config = config['training'].get('scheduled_sampling', {})
        self.use_scheduled_sampling = self.scheduled_sampling_config.get('enabled', False)
        
        if self.use_scheduled_sampling:
            logger.info(f"Scheduled Sampling enabled with mode: {self.scheduled_sampling_config.get('mode', 'linear')}")

        # ========== COMPREHENSIVE CONFIG LOGGING ==========
        loss_w = config['training'].get('loss_weights', {})
        ss_cfg = self.scheduled_sampling_config
        logger.info("=" * 70)
        logger.info("[CONFIG SUMMARY] Training Configuration at Startup:")
        logger.info(f"  Loss Weights:")
        logger.info(f"    position_before_icp:     {loss_w.get('position_before_icp', 'N/A')}")
        logger.info(f"    position_after_icp:      {loss_w.get('position_after_icp', 'N/A')}")
        logger.info(f"    acceleration_before_icp: {loss_w.get('acceleration_before_icp', 'N/A')}")
        logger.info(f"    acceleration_after_icp:  {loss_w.get('acceleration_after_icp', 'N/A')}")
        logger.info(f"    penetration_sign:        {loss_w.get('penetration_sign', 'N/A')}")
        logger.info(f"    rotation_anchor:         {loss_w.get('rotation_anchor', 0.0)}")
        logger.info(f"    if_time_step_scale:      {loss_w.get('if_time_step_scale', False)}")
        logger.info(f"  Scheduled Sampling:")
        logger.info(f"    enabled:          {ss_cfg.get('enabled', False)}")
        logger.info(f"    mode:             {ss_cfg.get('mode', 'N/A')}")
        logger.info(f"    teacher_epochs:   {ss_cfg.get('teacher_epochs', 'N/A')}")
        logger.info(f"    ramping_epochs:   {ss_cfg.get('ramping_epochs', 'N/A')}")
        logger.info(f"  Model Features:")
        # Access model attrs (model_ref = unwrapped model, self.model may be DDP-wrapped)
        model_ref = model if not hasattr(model, 'module') else model.module
        logger.info(f"    if_phys_param:    {getattr(model_ref, '_if_phys_param', 'N/A')}")
        logger.info(f"    if_time_step:     {getattr(model_ref, '_if_time_step', 'N/A')}")
        logger.info(f"    film_conditioning:{getattr(model_ref, '_film_conditioning', 'N/A')}")
        logger.info(f"    encoder_film:     {getattr(model_ref, 'encoder_film', None) is not None}")
        logger.info(f"    encoder_input_dim:{getattr(model_ref, '_encoder_input_dim', 'N/A')}")
        logger.info(f"    pe_type:          {getattr(model_ref, '_pe_type', 'N/A')}")
        logger.info(f"    if_roma:          {getattr(model_ref, 'if_roma', 'N/A')}")
        logger.info(f"  Training:")
        logger.info(f"    sequence_length:  {config['data'].get('sequence_length', 'N/A')}")
        logger.info(f"    eval_warmup:      {config.get('evaluation', {}).get('eval_warmup_frames', 'N/A')}")
        logger.info(f"    optimizer:        {config['training']['optimizer'].get('type', 'N/A')}")
        logger.info(f"    lr:               {config['training']['optimizer'].get('lr', 'N/A')}")
        logger.info("=" * 70)



        # Initialize evaluator for RMSE metrics
        from src.training.evaluator import Evaluator
        self.evaluator = Evaluator(self)
        logger.info(f"Evaluator initialized: metrics_enabled={self.evaluator.metrics_enabled}, "
                   f"print_every_n_epochs={self.evaluator.print_every_n_epochs}")

        # Ensure DDP sampler epoch set before first epoch
        try:
            if hasattr(self.train_loader, 'sampler') and hasattr(self.train_loader.sampler, 'set_epoch'):
                self.train_loader.sampler.set_epoch(self.epoch)
            if hasattr(self.val_loader, 'sampler') and hasattr(self.val_loader.sampler, 'set_epoch'):
                self.val_loader.sampler.set_epoch(self.epoch)
        except Exception:
            pass

    def _filament_renderer(self):
        """The process's offscreen renderer for videos (src/utils/rollout_render.py), created on first use."""
        from src.utils import rollout_render
        ck = self.config['training']['checkpoint']
        width, height = (int(v) for v in ck.get('video_size', [1280, 720]))
        return rollout_render.get_renderer(width, height, int(ck.get('video_supersample', 2)),
                                           floor=ck.get('video_floor', 'light'))

    def _polyscope_self_test(self):
        """Attempt to initialize Polyscope headless EGL once at startup; raise on failure."""
        os.environ.setdefault('PYOPENGL_PLATFORM', 'egl')
        os.environ.setdefault('POLYSCOPE_BACKEND', 'openGL3_egl')
        os.environ.setdefault('POLYSCOPE_ALLOW_HEADLESS', '1')
        import polyscope as ps
        # Allow headless backends and select egl before init
        if hasattr(ps, 'set_allow_headless'):
            ps.set_allow_headless(True)
        # Set EGL device index to match CUDA_VISIBLE_DEVICES (for multi-GPU systems)
        egl_device_id = os.environ.get('EGL_DEVICE_ID', None)
        if egl_device_id is not None and hasattr(ps, 'set_egl_device_index'):
            ps.set_egl_device_index(int(egl_device_id))
            logger.info(f"Polyscope EGL device index set to: {egl_device_id}")
        # Initialize the EGL backend directly
        ps.init('openGL3_egl')
        # Shutdown immediately; we only test availability
        if hasattr(ps, 'shutdown'):
            ps.shutdown()
        logger.info("Polyscope headless (EGL) initialization test passed")
        
    def get_teacher_forcing_ratio(self, epoch):
        """Calculate teacher forcing ratio based on scheduled sampling configuration."""
        if not self.use_scheduled_sampling:
            return 1.0  # Always use teacher forcing if disabled
        
        config = self.scheduled_sampling_config
        mode = config.get('mode', 'linear')
        
        if mode == 'three_stage':
            # Three-stage schedule like character-motion-vaes
            teacher_epochs = config.get('teacher_epochs', 100)
            ramping_epochs = config.get('ramping_epochs', 200)
            
            if epoch < teacher_epochs:
                # Pure teacher forcing
                return 1.0
            elif epoch < teacher_epochs + ramping_epochs:
                # Linear ramping
                progress = (epoch - teacher_epochs) / ramping_epochs
                return 1.0 - progress  # From 1.0 to 0.0
            else:
                # Pure student forcing
                return 0.0
                
        elif mode == 'linear':
            # Linear decay
            start_epoch = config.get('start_epoch', 100)
            end_epoch = config.get('end_epoch', 1000)
            min_ratio = config.get('min_teacher_ratio', 0.1)
            max_ratio = config.get('max_teacher_ratio', 1.0)
            
            if epoch < start_epoch:
                return max_ratio
            elif epoch >= end_epoch:
                return min_ratio
            else:
                progress = (epoch - start_epoch) / (end_epoch - start_epoch)
                return max_ratio - progress * (max_ratio - min_ratio)
                
        elif mode == 'exponential':
            # Exponential decay
            start_epoch = config.get('start_epoch', 100)
            decay_rate = config.get('decay_rate', 0.99)
            min_ratio = config.get('min_teacher_ratio', 0.1)
            
            if epoch < start_epoch:
                return 1.0
            else:
                epochs_since_start = epoch - start_epoch
                ratio = decay_rate ** epochs_since_start
                return max(ratio, min_ratio)
                
        elif mode == 'sigmoid':
            # Inverse sigmoid decay
            start_epoch = config.get('start_epoch', 100)
            k = config.get('k', 5.0)
            min_ratio = config.get('min_teacher_ratio', 0.1)
            
            if epoch < start_epoch:
                return 1.0
            else:
                epochs_since_start = epoch - start_epoch
                ratio = k / (k + np.exp(epochs_since_start / k))
                return max(ratio, min_ratio)
        else:
            logger.warning(f"Unknown scheduled sampling mode: {mode}, using teacher forcing")
            return 1.0
    
    def train_epoch(self):
        """Train for one epoch."""
        self.model.train()
        # AMP is required for Flash Attention (fp16 or bf16)
        use_amp = True
        # Check for BF16 support (preferred on Ampere and newer)
        amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        scaler = torch.amp.GradScaler('cuda', enabled=(use_amp and amp_dtype == torch.float16)) # Scaler needed only for fp16
        total_loss = 0
        num_batches = 0
        # Per-iteration linear warmup from zero (training.scheduler.warmup_iter_epochs, in epochs).
        # A fresh Adam state takes uniform +-lr steps on every parameter for its first ~1000
        # iterations; ramping the rate from zero keeps a converged warm start intact.
        _wu_frac = float(self.config['training'].get('scheduler', {}).get('warmup_iter_epochs', 0.0))
        _iters_per_epoch = len(self.train_loader)
        _wu_iters = int(round(_wu_frac * _iters_per_epoch))
        _epoch_base_lrs = [g['lr'] for g in self.optimizer.param_groups]
        # Component accumulators (averaged per epoch)
        comp_sum_pos_before = 0.0
        comp_sum_pos_after = 0.0
        comp_sum_acc_after = 0.0
        comp_sum_acc_before = 0.0
        comp_sum_edge = 0.0
        comp_sum_pen = 0.0
        comp_sum_pen_acc = 0.0


        # Compute teacher forcing ratio from scheduled sampling config
        if self.use_scheduled_sampling:
            teacher_forcing_ratio = self.get_teacher_forcing_ratio(self.epoch)
        else:
            teacher_forcing_ratio = 1.0  # Scheduled sampling off -> pure teacher forcing

        if self.is_main:
            logger.info(f"[Train Epoch {self.epoch}] teacher_forcing_ratio={teacher_forcing_ratio:.3f} "
                       f"(scheduled_sampling={'ON' if self.use_scheduled_sampling else 'OFF'})")

        pbar = tqdm(self.train_loader, desc=f'Training Epoch {self.epoch}', disable=not self.is_main)
        for batch in pbar:
            # Move batch to device
            batch = {k: v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else v
                    for k, v in batch.items()}

            # Apply GPU augmentation (rotation) - only for training
            if self.gpu_augmentation is not None:
                batch = self.gpu_augmentation(batch, split='train')

            # Forward pass (Mixed Precision)
            with torch.amp.autocast('cuda', enabled=use_amp, dtype=amp_dtype):
                predictions = self.model(batch, teacher_forcing_ratio=teacher_forcing_ratio)

            # Access underlying module when using DDP
            _model_ref = self.model.module if isinstance(self.model, DDP) else self.model
            # For multi-object batches, we will compute GT inside the multi-object loss branch below
            
            # Loss aggregation
            if 'obj_positions_denorm' in predictions:
                # Multi-object path: aggregate across all valid vertices using masks
                import torch.nn.functional as F
                T0 = batch['obj_positions'].shape[1]
                # obj_positions: (B, T, Nobj, Nv, 3)
                obj_mask = batch['valid_obj_mask'].unsqueeze(1).unsqueeze(-1)  # (B,1,Nobj,1)
                # Position predictions align to frames 2..T0-1
                vmask_pos = batch['valid_vtx_mask'][:, 2:T0]  # (B,T-2,Nobj,Nv)
                full_mask_pos = (vmask_pos & obj_mask).to(predictions['obj_positions_norm'].dtype)
                # Acceleration predictions align to centers 1..T0-2 (curr index in forward)
                vmask_acc = batch['valid_vtx_mask'][:, 1:T0-1]  # (B,T-2,Nobj,Nv)
                full_mask_acc = (vmask_acc & obj_mask).to(predictions['obj_positions_norm'].dtype)

                # ============================================================
                # Anchor-only losses: all losses are computed at the Nq anchor points.
                # Rationale: the network predicts accelerations at anchors, Kabsch R,t
                # come from anchors, and all Nv vertices are a deterministic function
                # of the anchors (rigid body). A full-vertex loss would be redundant
                # and dilute the gradient signal.
                # ============================================================

                # Prepare anchor index for gathering (shared across all losses)
                if 'anchor_idx' not in batch:
                    raise KeyError("Missing 'anchor_idx' in batch for anchor-based supervision")
                anchor_idx_b = batch['anchor_idx']  # (B,Nobj,Nq)
                Bsz = batch['obj_positions'].shape[0]
                Tm2_pos = T0 - 2  # Number of prediction steps
                Nq_loss = anchor_idx_b.shape[2]

                # Anchor gather index: (B,Nobj,Nq) -> (B,T-2,Nobj,Nq,3) for 5D tensors
                anchor_idx_pos = anchor_idx_b.unsqueeze(1).unsqueeze(-1).expand(
                    -1, Tm2_pos, -1, -1, 3
                )  # (B,T-2,Nobj,Nq,3)

                # Anchor mask: obj_mask only (all anchors valid by construction for valid objects)
                # obj_mask is (B,1,Nobj,1) -> (B,T-2,Nobj,Nq)
                anchor_mask_pos = obj_mask.expand(-1, Tm2_pos, -1, Nq_loss).to(predictions['obj_positions_norm'].dtype)
                # Divisor of the anchor losses: this rank's valid-anchor count, or with training.loss_weights.
                # global_anchor_mean the global count / world size, so that DDP's gradient average is the gradient of
                # the mean over every valid anchor of the global batch (ranks hold different numbers of objects)
                _anchor_count = anchor_mask_pos.float().sum()
                if (self.config['training']['loss_weights'].get('global_anchor_mean', False)
                        and dist.is_available() and dist.is_initialized()):
                    _anchor_count = _anchor_count.detach().clone()
                    dist.all_reduce(_anchor_count, op=dist.ReduceOp.SUM)
                    _anchor_count = _anchor_count / dist.get_world_size()

                # 1) Position loss BEFORE Kabsch (anchor-only)
                gt_pos = batch['obj_positions'][:, 2:T0]  # (B,T-2,Nobj,Nv,3)
                gt_pos_norm = gt_pos

                # Gather at anchors
                gt_pos_anchor = torch.gather(gt_pos_norm, dim=3, index=anchor_idx_pos)  # (B,T-2,Nobj,Nq,3)
                pred_pos_anchor = torch.gather(predictions['obj_positions_norm'], dim=3, index=anchor_idx_pos)  # (B,T-2,Nobj,Nq,3)

                # Compute position error at anchors
                pos_error = pred_pos_anchor - gt_pos_anchor  # (B,T-2,Nobj,Nq,3)

                # Apply step_size² normalization if enabled
                if_time_step_scale_pos = self.config['training']['loss_weights'].get('if_time_step_scale', True)
                if if_time_step_scale_pos:
                    if 'skip_step_size' not in batch:
                        raise KeyError(
                            "Missing 'skip_step_size' in batch for position loss normalization!\n"
                            "Multi-timescale training with if_time_step_scale=true requires step_size."
                        )
                    step_sizes_pos = batch['skip_step_size'].float()  # (B,)
                    step_size_scale_sq_pos = step_sizes_pos ** 2  # (B,)
                    scale_pos = step_size_scale_sq_pos.view(Bsz, 1, 1, 1, 1)  # (B,1,1,1,1)
                    pos_error_normalized = pos_error / scale_pos  # (B,T-2,Nobj,Nq,3)
                else:
                    pos_error_normalized = pos_error

                pos_l1 = F.smooth_l1_loss(
                    pos_error_normalized, torch.zeros_like(pos_error_normalized), reduction='none'
                ).sum(dim=-1)  # (B,T-2,Nobj,Nq)
                position_loss_before_icp = (pos_l1.float() * anchor_mask_pos.float()).sum() / (_anchor_count + 1e-8)

                # 1b) Position loss AFTER Kabsch (anchor-only)
                if 'obj_positions_icp_norm' in predictions:
                    pred_pos_icp_anchor = torch.gather(predictions['obj_positions_icp_norm'], dim=3, index=anchor_idx_pos)  # (B,T-2,Nobj,Nq,3)
                    pos_error_icp = pred_pos_icp_anchor - gt_pos_anchor  # (B,T-2,Nobj,Nq,3)
                    if if_time_step_scale_pos:
                        pos_error_icp_normalized = pos_error_icp / scale_pos
                    else:
                        pos_error_icp_normalized = pos_error_icp
                    pos_icp_l1 = F.smooth_l1_loss(
                        pos_error_icp_normalized, torch.zeros_like(pos_error_icp_normalized), reduction='none'
                    ).sum(dim=-1)  # (B,T-2,Nobj,Nq)
                    position_loss_after_icp = (pos_icp_l1.float() * anchor_mask_pos.float()).sum() / (_anchor_count + 1e-8)
                else:
                    position_loss_after_icp = torch.tensor(0.0, device=self.device)

                # 1c) Orientation loss (training.loss_weights.rotation_anchor, anchor-only): the after-Kabsch anchor error
                # without its mean over the object's anchors (the translation part), divided by the RMS distance of the
                # ground-truth anchors from their centroid. This is about the rotation error in radians whatever the
                # object's size; the position losses weigh a rotation error by the object's size, while the orientation
                # metric (degrees per object) does not.
                rot_anchor_w = float(self.loss_weights.get('rotation_anchor', 0.0))
                if rot_anchor_w > 0:
                    if 'obj_positions_icp_norm' not in predictions:
                        raise KeyError("training.loss_weights.rotation_anchor needs 'obj_positions_icp_norm' in the predictions")
                    rot_error = centered_anchor_error(pred_pos_icp_anchor, gt_pos_anchor)  # (B,T-2,Nobj,Nq,3)
                    if if_time_step_scale_pos:
                        rot_error = rot_error / scale_pos
                    rot_l1 = F.smooth_l1_loss(rot_error, torch.zeros_like(rot_error), reduction='none').sum(dim=-1)
                    rotation_anchor_loss = (rot_l1.float() * anchor_mask_pos.float()).sum() / (_anchor_count + 1e-8)
                else:
                    rotation_anchor_loss = torch.tensor(0.0, device=self.device)

                # 2) Acceleration losses in normalized space
                # Targets: the acceleration that, integrated from the frames the network was given (the rigid-noise
                # augmentation's perturbed frames when present), reaches the ground-truth next frame; with clean
                # inputs this is the ground-truth acceleration
                _inputs = batch.get('obj_positions_input', batch['obj_positions'])
                gt_prev = _inputs[:, 0:T0-2]
                gt_curr = _inputs[:, 1:T0-1]
                gt_next = batch['obj_positions'][:, 2:T0]
                gt_accel_den = gt_next - 2 * gt_curr + gt_prev  # (B,T-2,Nobj,Nv,3)
                gt_accel_norm = gt_accel_den

                # 2) Acceleration losses (anchor-only for both before and after Kabsch)
                a_icp_norm = predictions['obj_accelerations_norm']  # (B,T-2,Nobj,Nv,3) — after Kabsch
                if 'obj_accelerations_no_icp_norm_anchor' not in predictions:
                    raise KeyError("Missing 'obj_accelerations_no_icp_norm_anchor' in predictions")
                a_no_icp_norm_anchor = predictions['obj_accelerations_no_icp_norm_anchor']  # (B,T-2,Nobj,Nq,3)
                Bsz_acc, Tm2, Nobj_acc, Nv_acc, _ = a_icp_norm.shape

                # Gather GT accel and predicted accel_after_icp at anchor points
                # Reuse anchor_idx_b from position loss section above
                idx_exp_acc = anchor_idx_b.unsqueeze(1).unsqueeze(-1).expand(-1, Tm2, -1, -1, 3)  # (B,T-2,Nobj,Nq,3)
                gt_accel_anchor = torch.gather(gt_accel_norm, dim=3, index=idx_exp_acc)  # (B,T-2,Nobj,Nq,3)
                a_icp_anchor = torch.gather(a_icp_norm, dim=3, index=idx_exp_acc)  # (B,T-2,Nobj,Nq,3)

                # Step_size² normalization
                if_time_step_scale = self.config['training']['loss_weights'].get('if_time_step_scale', True)
                if_debug = self.config['training']['loss_weights'].get('if_debug', False)

                # Compute acceleration errors at anchors
                error_icp = a_icp_anchor - gt_accel_anchor  # (B,T-2,Nobj,Nq,3)
                error_no_icp = a_no_icp_norm_anchor - gt_accel_anchor  # (B,T-2,Nobj,Nq,3)

                if if_time_step_scale:
                    if 'skip_step_size' not in batch:
                        raise KeyError(
                            "Missing 'skip_step_size' in batch!\n"
                            "Multi-timescale training with if_time_step_scale=true requires step_size.\n"
                            "Please ensure dataloader provides 'skip_step_size' field."
                        )
                    step_sizes = batch['skip_step_size'].float()  # (B,)
                    step_size_scale_sq = step_sizes ** 2  # (B,)
                    scale_acc = step_size_scale_sq.view(Bsz_acc, 1, 1, 1, 1)  # (B,1,1,1,1)
                    error_icp_normalized = error_icp / scale_acc  # (B,T-2,Nobj,Nq,3)
                    error_no_icp_normalized = error_no_icp / scale_acc  # (B,T-2,Nobj,Nq,3)
                else:
                    error_icp_normalized = error_icp
                    error_no_icp_normalized = error_no_icp

                # Anchor-only L1 losses
                a_ste_l1 = F.smooth_l1_loss(error_icp_normalized, torch.zeros_like(error_icp_normalized), reduction='none').sum(dim=-1)  # (B,T-2,Nobj,Nq)
                a_no_icp_l1 = F.smooth_l1_loss(error_no_icp_normalized, torch.zeros_like(error_no_icp_normalized), reduction='none').sum(dim=-1)  # (B,T-2,Nobj,Nq)

                # Anchor mask: obj_mask only (all anchors valid for valid objects); the same anchors as the position
                # losses (same objects, T-2 steps), so the same divisor
                anchor_full_mask = obj_mask.expand(-1, Tm2, -1, Nq_loss).to(a_icp_norm.dtype)  # (B,T-2,Nobj,Nq)
                if anchor_full_mask.shape != anchor_mask_pos.shape:
                    raise ValueError(f"acceleration and position anchor masks differ: {tuple(anchor_full_mask.shape)} vs "
                                     f"{tuple(anchor_mask_pos.shape)}")
                accel_loss_after_icp_ste = (a_ste_l1.float() * anchor_full_mask.float()).sum() / (_anchor_count + 1e-8)
                accel_loss_before_icp = (a_no_icp_l1.float() * anchor_full_mask.float()).sum() / (_anchor_count + 1e-8)

                # 3) Penetration sign supervision (optional) - ALL VERTICES
                pen_loss = torch.tensor(0.0, device=self.device)
                pen_acc = torch.tensor(0.0, device=self.device)
                penetration_enabled = self.config['model'].get('if_penetration', self.config['model'].get('if_directional_distance', False))
                do_pen = penetration_enabled and ('obj_pen_pred_vertices' in predictions)
                if do_pen:
                    try:
                        # Prepare GT signs for ALL vertices for steps 1..T0-2 (aligned with predictions time)
                        # For each step t (forward uses prev=t-1, curr=t), use current frame vertices to test inside others
                        with torch.no_grad():
                            import numpy as _np
                            import trimesh as _trimesh  # type: ignore
                            Bsz = batch['obj_positions'].shape[0]
                            Tm2 = predictions['obj_pen_pred_vertices'].shape[1]
                            Nobj_ = batch['obj_positions'].shape[2]
                            Nv = batch['obj_positions'].shape[3]
                            gt_sign_vertices = torch.ones((Bsz, Tm2, Nobj_, Nv, 1), dtype=torch.float32, device=self.device)
                            valid_obj = batch['valid_obj_mask']  # (B,Nobj)
                            for b in range(Bsz):
                                # dataset name per-sample is a string (from dataloader list)
                                if 'dataset_name' not in batch or not isinstance(batch['dataset_name'], (list, tuple)):
                                    raise RuntimeError("Batch missing per-sample 'dataset_name' list for penetration GT")
                                ds_name = batch['dataset_name'][b]
                                if not isinstance(ds_name, str) or not ds_name:
                                    raise RuntimeError(f"Invalid dataset_name for sample {b}: {ds_name}")
                                for tstep in range(Tm2):
                                    # Skip tstep=0 (t=1, placeholder, GT positions, no penetration)
                                    if tstep == 0:
                                        continue

                                    # For tstep>=1: use predicted mesh positions to construct GT penetration
                                    # tstep=1 -> t=2 -> uses predictions[0]; tstep=2 -> t=3 -> uses predictions[1]
                                    # Use STE blended positions (after Kabsch projection), which match
                                    # forward's curr_pos_bt in student forcing mode:
                                    #   obj_positions_denorm = next_pos_bt (before Kabsch)
                                    #   obj_positions_ste = blended_next (after Kabsch, STE)
                                    V_all = predictions['obj_positions_ste'][b, tstep-1]  # (Nobj,Nv,3), after transform
                                    # Vertex mask for the corresponding frame: tstep=1 -> frame2, tstep=2 -> frame3
                                    vmask_all = batch['valid_vtx_mask'][b, tstep+1]  # (Nobj,Nv) - mask for frame tstep+1
                                    for i in range(Nobj_):
                                        if not valid_obj[b, i]:
                                            continue
                                        Vi = V_all[i][vmask_all[i]]  # (Nvi,3)
                                        if Vi.shape[0] == 0:
                                            continue
                                        inside_any = _np.zeros((Vi.shape[0],), dtype=_np.bool_)
                                        # Test against all other objects' meshes (triangle faces from dataset directory)
                                        for j in range(Nobj_):
                                            if j == i or not valid_obj[b, j]:
                                                continue
                                            Vj = V_all[j][vmask_all[j]].detach().cpu().numpy()
                                            # Load faces from batch (must exist; no fallback)
                                            try:
                                                # Use permuted faces from batch - this ensures faces match the permuted object order
                                                if 'faces_list' not in batch or batch['faces_list'] is None:
                                                    raise RuntimeError(
                                                        f"Missing 'faces_list' in batch! DataLoader must provide faces. "
                                                        f"Check that faces_list is included in sample dict."
                                                    )
                                                faces_np = batch['faces_list'][b][j].astype(_np.int64, copy=False)
                                                # Filter faces by current Nv mask (drop faces referencing padded verts)
                                                idx_valid = _np.where(vmask_all[j].detach().cpu().numpy())[0]
                                                if idx_valid.size == 0:
                                                    continue
                                                # Build mapping original->compact indices
                                                rev = -_np.ones((int(vmask_all[j].shape[0]),), dtype=_np.int64)
                                                rev[idx_valid] = _np.arange(idx_valid.size, dtype=_np.int64)
                                                faces_compact = rev[faces_np]
                                                faces_compact = faces_compact[(faces_compact >= 0).all(axis=1)]
                                                if faces_compact.size == 0 or faces_compact.shape[0] == 0:
                                                    continue
                                                mesh_j = _trimesh.Trimesh(vertices=Vj, faces=faces_compact, process=False)
                                                inside_j = mesh_j.contains(Vi.detach().cpu().numpy())
                                                inside_any |= inside_j
                                            except Exception as __e:
                                                logger.warning(f"Pen GT faces/contains failed (ds={ds_name}, obj={j}): {__e}")
                                        # Map vertex-level GT to full Nv space (ALL vertices)
                                        # inside_any is compact (Vi.shape[0]); need to expand to full Nv space
                                        idx_valid_i_val = _np.where(vmask_all[i].detach().cpu().numpy())[0]
                                        Nv_full = int(vmask_all[i].shape[0])
                                        inside_full = _np.zeros((Nv_full,), dtype=_np.bool_)
                                        if idx_valid_i_val.size > 0 and inside_any.shape[0] > 0:
                                            inside_full[idx_valid_i_val] = inside_any
                                        # Convert to sign for ALL vertices
                                        gt_sign = _np.where(inside_full, -1.0, 1.0).astype(_np.float32)
                                        gt_sign_vertices[b, tstep, i, :, 0] = torch.from_numpy(gt_sign).to(self.device)

                        # Pred scores and loss/acc (ALL VERTICES)
                        pen_pred = predictions['obj_pen_pred_vertices']  # (B,T-2,Nobj,Nv,1) in [-1,1]
                        pen_l2 = (pen_pred - gt_sign_vertices).pow(2).squeeze(-1)  # (B,T-2,Nobj,Nv)

                        # Vertex-level mask for loss aggregation
                        vtx_mask_for_loss = batch['valid_vtx_mask'][:, 1:Tm2+1, :, :]  # (B,T-2,Nobj,Nv)

                        # Time-level mask: Skip t=1 (tstep=0) where no penetration prediction was made
                        # Only supervise t>=2 (tstep>=1) where actual prediction occurred
                        time_mask = torch.ones((Tm2,), dtype=torch.float32, device=self.device)
                        time_mask[0] = 0.0  # Mask out t=1 (first step, GT positions, placeholder prediction)
                        time_mask = time_mask.view(1, Tm2, 1, 1)  # (1, T-2, 1, 1)

                        # Combined mask: vertex validity AND time validity
                        vertex_full_mask = vtx_mask_for_loss.to(pen_l2.dtype) * time_mask  # (B,T-2,Nobj,Nv)
                        pen_loss = (pen_l2.float() * vertex_full_mask.float()).sum() / (vertex_full_mask.float().sum() + 1e-8)

                        # Accuracy: sign match
                        pred_sign = torch.sign(pen_pred)
                        acc_bool = (pred_sign == torch.sign(gt_sign_vertices)).squeeze(-1)
                        pen_acc = (acc_bool * vertex_full_mask.bool()).sum().float() / (vertex_full_mask.bool().sum().float() + 1e-8)
                    except Exception as _e:
                        logger.warning(f"Penetration GT computation failed: {_e}")
            else:
                raise ValueError("Expected multi-object predictions with 'obj_positions_denorm'.")
                

            ##############################
            # Visualization of predictions vs ground truth
            # Same frequency as video saving
            if self.config['training']['checkpoint'].get('save_video', False):
                video_every = self.config['training']['checkpoint'].get('video_every_n_epochs', 5)
                if self.epoch % video_every == 0 and num_batches == 1:  # Only visualize first batch
                    self._visualize_predictions_vs_gt(
                        predictions, batch, batch_idx=num_batches-1
                    )
            
            # Compute edge length loss to preserve rigid body structure
            edge_loss = 0.0
            if self.loss_weights.get('edge_length', 0.0) > 0:
                if 'obj_positions_denorm' in predictions:
                    # Proxy edge-length loss, computed per object relative to t=0
                    init_obj = batch['obj_positions'][:, 0]  # (B,Nobj,Nv,3)
                    # Sum edge loss over objects using mask
                    vmask0 = batch['valid_vtx_mask'][:, 0]  # (B,Nobj,Nv)
                    obj_mask_b = batch['valid_obj_mask'].unsqueeze(-1)  # (B,Nobj,1)
                    mask0 = (vmask0 & obj_mask_b)
                    # Flatten objects and apply edge length loss per object then average over valid
                    pred_den = predictions['obj_positions_denorm']  # (B,T-2,Nobj,Nv,3)
                    # Use t-0 anchor per object
                    # Compute mean over time first, then per-object edge len towards init
                    pred_mean = pred_den.mean(dim=1)  # (B,Nobj,Nv,3)
                    # Edge loss proxy: mean squared displacement from initial pose
                    ed = (pred_mean - init_obj).pow(2).sum(dim=-1)
                    edge_loss = (ed.float() * mask0.float()).sum() / (mask0.float().sum() + 1e-8)
            
            # Combine losses
            pos_before_icp_w = self.loss_weights.get('position_before_icp', 0.0)  # Before Kabsch
            pos_after_icp_w = self.loss_weights.get('position_after_icp', 0.0)  # After Kabsch
            accel_after_icp_w = self.loss_weights.get('acceleration_after_icp', 0.0)  # After Kabsch
            accel_before_icp_w = self.loss_weights.get('acceleration_before_icp', 0.0)  # Before Kabsch

            loss = (
                pos_before_icp_w * position_loss_before_icp
                + pos_after_icp_w * position_loss_after_icp
                + accel_after_icp_w * accel_loss_after_icp_ste
                + accel_before_icp_w * accel_loss_before_icp
                + rot_anchor_w * rotation_anchor_loss
            )

            # Add penetration loss with its own weight if enabled
            pen_w = self.loss_weights.get('penetration_sign', 0.0)
            if pen_w > 0 and 'obj_pen_pred_vertices' in predictions:
                loss = loss + pen_w * pen_loss

            if self.loss_weights.get('edge_length', 0.0) > 0:
                loss = loss + self.loss_weights['edge_length'] * edge_loss

            # Per-step-size loss weighting: weight = 1/step_size
            # s=1: x1.0, s=5: x0.2, s=10: x0.1 -> prioritizes fine-grained (s=1) accuracy
            loss_weight_per_ss = self.config['training']['loss_weights'].get('loss_weight_per_step_size', False)
            if loss_weight_per_ss and 'skip_step_size' in batch:
                ss_val = batch['skip_step_size'].float().mean().item()  # All samples in batch share same step_size
                ss_weight = 1.0 / max(ss_val, 1.0)
                loss = loss * ss_weight

            # Diagnostic: log a detailed breakdown before backward when the loss is suspiciously high
            loss_val_pre = loss.item()
            if loss_val_pre > 100 or torch.isnan(loss) or torch.isinf(loss):
                ss_info = batch['skip_step_size'].cpu().tolist() if 'skip_step_size' in batch else 'N/A'
                logger.error(f"[Loss Explosion] Loss={loss_val_pre:.4e} at batch {num_batches}, step_size={ss_info}")
                logger.error(f"  pos_before_icp={position_loss_before_icp.item():.6e} (w={pos_before_icp_w})")
                logger.error(f"  pos_after_icp ={position_loss_after_icp.item():.6e} (w={pos_after_icp_w})")
                logger.error(f"  acc_before_icp={accel_loss_before_icp.item():.6e} (w={accel_before_icp_w})")
                logger.error(f"  acc_after_icp ={accel_loss_after_icp_ste.item():.6e} (w={accel_after_icp_w})")
                logger.error(f"  rot_anchor    ={rotation_anchor_loss.item():.6e} (w={rot_anchor_w})")
                # Log which component caused the explosion
                weighted = {
                    'pos_before': pos_before_icp_w * position_loss_before_icp.item(),
                    'pos_after': pos_after_icp_w * position_loss_after_icp.item(),
                    'acc_before': accel_before_icp_w * accel_loss_before_icp.item(),
                    'acc_after': accel_after_icp_w * accel_loss_after_icp_ste.item(),
                }
                culprit = max(weighted, key=weighted.get)
                logger.error(f"  Weighted: {weighted} -> culprit={culprit}")
                logger.error(f"  time_step_id range: {batch['time_step_id'].min().item()}-{batch['time_step_id'].max().item()}")

            # Non-finite loss: abort on every rank together before anything reaches the
            # parameters or the optimizer state (a rank raising alone would leave the others
            # blocked in the gradient all-reduce). No fallback: the [Loss Explosion] breakdown
            # above is the diagnostic.
            _bad = (~torch.isfinite(loss.detach())).to(torch.float32).reshape(1)
            if dist.is_available() and dist.is_initialized():
                dist.all_reduce(_bad, op=dist.ReduceOp.MAX)
            if _bad.item() > 0:
                raise RuntimeError(
                    f"Non-finite loss ({loss_val_pre}) at epoch {self.epoch}, batch {num_batches}; "
                    f"aborting before the optimizer step")

            # Backward pass
            self.optimizer.zero_grad()
            loss = loss.float()
            scaler.scale(loss).backward()

            # Gradient clipping with AMP requires unscale first
            clip_val = self.config['training'].get('gradient_clip', 0.0)
            if clip_val and clip_val > 0:
                scaler.unscale_(self.optimizer)
                # Gradients are identical on all ranks after DDP's all-reduce, so a non-finite
                # norm raises on every rank together instead of being written into the weights.
                # (With an fp16 GradScaler, inf gradients are expected occasionally and the
                # scaler skips that step itself, so the check is left to the scaler there.)
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), 
                    clip_val,
                    norm_type=self.config['training'].get('gradient_clip_norm_type', 2),
                    error_if_nonfinite=not scaler.is_enabled(),
                )
                # Log when grad_norm was clipped significantly
                if grad_norm > clip_val * 10:
                    logger.warning(f"[GRAD CLIP] grad_norm={grad_norm:.2f} >> clip_val={clip_val}, loss={loss_val_pre:.4e}")
            if _wu_iters > 0:
                _gstep = self.epoch * _iters_per_epoch + num_batches
                if _gstep < _wu_iters:
                    _f = (_gstep + 1) / _wu_iters
                    for _g, _b in zip(self.optimizer.param_groups, _epoch_base_lrs):
                        _g['lr'] = _b * _f
                    if _gstep in (0, _wu_iters // 2, _wu_iters - 1) and self.is_main:
                        logger.info("[Warmup-iter] step %d/%d: lr = %s" % (_gstep + 1, _wu_iters, ['%.2e' % _g['lr'] for _g in self.optimizer.param_groups]))
            scaler.step(self.optimizer)
            scaler.update()
            if self.ema is not None:
                self.ema.update()
            
            # Update progress bar (only on rank 0 when DDP)
            loss_value = loss.item()

            # Check for inf/nan and log warning
            if torch.isnan(loss).any() or torch.isinf(loss).any():
                logger.warning(f"[Batch {num_batches}] Loss is inf/nan: {loss_value}")
                logger.warning(f"  Accel after ICP: {accel_loss_after_icp_ste.item()}, Accel before ICP: {accel_loss_before_icp.item()}")
                logger.warning(f"  Position before ICP: {position_loss_before_icp.item()}, Position after ICP: {position_loss_after_icp.item()}")
                if self.loss_weights.get('edge_length', 0.0) > 0:
                    logger.warning(f"  Edge loss: {edge_loss.item() if isinstance(edge_loss, torch.Tensor) else edge_loss}")
                # Skip this batch for average calculation
            else:
                total_loss += loss_value
                num_batches += 1

            if self.is_main:
                pbar.set_postfix({'loss': loss_value})

            # Synchronize batch losses across all ranks for accurate logging
            if dist.is_available() and dist.is_initialized():
                # Prepare all loss components for sync
                edge_loss_val = edge_loss.item() if isinstance(edge_loss, torch.Tensor) else edge_loss
                # Pack pen metrics too (fallback to zeros)
                pen_loss_val = pen_loss.item() if isinstance(pen_loss, torch.Tensor) else float(pen_loss)
                pen_acc_val = pen_acc.item() if isinstance(pen_acc, torch.Tensor) else float(pen_acc)
                pos_after_icp_val = position_loss_after_icp.item() if isinstance(position_loss_after_icp, torch.Tensor) else float(position_loss_after_icp)
                loss_components = torch.tensor([
                    loss.item(),
                    accel_loss_after_icp_ste.item(),
                    accel_loss_before_icp.item(),
                    position_loss_before_icp.item(),
                    pos_after_icp_val,
                    edge_loss_val,
                    pen_loss_val,
                    pen_acc_val,
                ], device=self.device)
                # Sum across all ranks (Gloo backend does not support AVG), then divide by world size
                dist.all_reduce(loss_components, op=dist.ReduceOp.SUM)
                if dist.is_available() and dist.is_initialized():
                    world_size = dist.get_world_size()
                    loss_components = loss_components / max(world_size, 1)
                loss_sync, accel_after_icp_sync, accel_before_icp_sync, pos_before_icp_sync, pos_after_icp_sync, edge_sync, pen_loss_sync, pen_acc_sync = loss_components.tolist()
            else:
                # Single GPU: use local values
                loss_sync = loss.item()
                accel_after_icp_sync = accel_loss_after_icp_ste.item()
                accel_before_icp_sync = accel_loss_before_icp.item()
                pos_before_icp_sync = position_loss_before_icp.item()
                pos_after_icp_sync = position_loss_after_icp.item() if isinstance(position_loss_after_icp, torch.Tensor) else float(position_loss_after_icp)
                edge_sync = edge_loss.item() if isinstance(edge_loss, torch.Tensor) else edge_loss
                pen_loss_sync = pen_loss.item() if isinstance(pen_loss, torch.Tensor) else float(pen_loss)
                pen_acc_sync = pen_acc.item() if isinstance(pen_acc, torch.Tensor) else float(pen_acc)


            # Log to tensorboard (only rank 0 in DDP)
            if self.is_main:
                global_step = self.epoch * len(self.train_loader) + num_batches
                # Rigid alpha is fixed at 1.0
                self.writer.add_scalar('train/rigid_alpha', 1.0, global_step)
                self.writer.add_scalar('train/loss', loss_sync, global_step)
                self.writer.add_scalar('train/accel_loss_after_icp', accel_after_icp_sync, global_step)
                self.writer.add_scalar('train/accel_loss_before_icp', accel_before_icp_sync, global_step)
                self.writer.add_scalar('train/position_loss_before_icp', pos_before_icp_sync, global_step)
                if rot_anchor_w > 0:
                    self.writer.add_scalar('train/rotation_anchor_loss', rotation_anchor_loss.item(), global_step)
                self.writer.add_scalar('train/position_loss_after_icp', pos_after_icp_sync, global_step)
                if self.loss_weights.get('edge_length', 0.0) > 0:
                    self.writer.add_scalar('train/edge_loss', edge_sync, global_step)
                if pen_w > 0 and 'obj_pen_pred_vertices' in predictions:
                    self.writer.add_scalar('train/pen_sign_loss', pen_loss_sync, global_step)
                    self.writer.add_scalar('train/pen_sign_acc', pen_acc_sync, global_step)


            # Accumulate epoch sums (use synced values)
            comp_sum_pos_before += float(pos_before_icp_sync)
            comp_sum_pos_after += float(pos_after_icp_sync)
            comp_sum_acc_after += float(accel_after_icp_sync)
            comp_sum_acc_before += float(accel_before_icp_sync)
            comp_sum_edge += float(edge_sync)
            comp_sum_pen += float(pen_loss_sync)
            comp_sum_pen_acc += float(pen_acc_sync)

        
        if num_batches > 0:
            avg_loss = total_loss / num_batches
            avg_pos_before = comp_sum_pos_before / num_batches
            avg_pos_after = comp_sum_pos_after / num_batches
            avg_acc_after = comp_sum_acc_after / num_batches
            avg_acc_before = comp_sum_acc_before / num_batches
            avg_edge = comp_sum_edge / num_batches
            avg_pen = comp_sum_pen / max(num_batches, 1)
            avg_pen_acc = comp_sum_pen_acc / max(num_batches, 1)

        else:
            logger.warning("No valid batches in training epoch (all losses were inf/nan)")
            avg_loss = float('inf')
            avg_pos_before = float('nan')
            avg_pos_after = float('nan')
            avg_acc_after = float('nan')
            avg_acc_before = float('nan')
            avg_edge = float('nan')
            avg_pen = float('nan')
            avg_pen_acc = float('nan')


        # Synchronize loss across all ranks in DDP mode
        if dist.is_available() and dist.is_initialized():
            loss_tensor = torch.tensor(avg_loss, device=self.device)
            dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM)
            world_size = dist.get_world_size()
            avg_loss = (loss_tensor / max(world_size, 1)).item()

        self.last_train_loss = avg_loss  # Store for overfitting mode
        if _wu_iters > 0:
            # undo the per-iteration warm-up scaling before the epoch scheduler advances from these rates
            for _g, _b in zip(self.optimizer.param_groups, _epoch_base_lrs):
                _g['lr'] = _b
        return {
            'loss': float(avg_loss),
            'pos_before_icp': float(avg_pos_before),
            'pos_after_icp': float(avg_pos_after),
            'accel_after_icp': float(avg_acc_after),
            'accel_before_icp': float(avg_acc_before),
            'edge': float(avg_edge),
            'pen_sign_loss': float(avg_pen),
            'pen_sign_acc': float(avg_pen_acc),

        }
    


    def validate(self):
        """Validate the model."""
        # Skip validation if no validation data
        if self.val_loader is None or len(self.val_loader) == 0:
            logger.info("No validation data, using training loss as validation metric")
            return self.last_train_loss if hasattr(self, 'last_train_loss') else float('inf')

        try:
            return self._validate_impl()
        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
            if 'out of memory' in str(e).lower():
                if self.inference_only:
                    logger.error(f"CUDA OOM during validation: {e}")
                    logger.warning("Inference-only mode: moving model to CPU and retrying validation")

                    # Move model to CPU
                    self.device = 'cpu'
                    self.model = self.model.to('cpu')
                    torch.cuda.empty_cache()
                    logger.info("Model moved to CPU (execution will be slower)")

                    # Retry validation on CPU
                    return self._validate_impl()
                else:
                    logger.error("CUDA OOM during training validation - stopping training")
                    logger.error("Reduce batch size or model size, or use a GPU with more memory")
                    raise
            else:
                raise

    def _validate_impl(self):
        """Internal validation implementation (can be called with GPU or CPU)."""
        self.model.eval()
        _model_ref = self.model.module if isinstance(self.model, DDP) else self.model
        if not hasattr(_model_ref, '_pe_type'):
            raise RuntimeError("Model missing _pe_type attribute; cannot validate without PE type.")
        total_loss = 0
        total_pos_loss = 0
        total_accel_loss = 0
        total_edge_loss = 0
        num_batches = 0

        with torch.no_grad():
            pbar = tqdm(self.val_loader, desc='Validation', disable=not self.is_main)
            for batch in pbar:
                # Move batch to device
                batch = {k: v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else v
                        for k, v in batch.items()}

                # Forward pass (Flash Attention handles dtype internally in decoder.py)
                predictions = self.model(batch, teacher_forcing_ratio=0.0)  # No teacher forcing during validation
                
                # Compute losses (multi-object path)
                if 'obj_positions_denorm' in predictions and 'obj_positions' in batch:
                    # Sequence length from multi-object tensor
                    seq_len = batch['obj_positions'].shape[1]
                    if seq_len <= 2:
                        logger.warning("Single/short sequence in validation - skipping")
                        continue
                    # Ground truth positions for steps 1..T-2
                    gt_obj_pos = batch['obj_positions'][:, 2:seq_len]  # (B,T-2,Nobj,Nv,3)
                    pred_obj_pos = predictions['obj_positions_denorm']    # (B,T-2,Nobj,Nv,3)
                    # Ensure vertex dimension matches between pred and gt (defensive alignment)
                    Nv_pred = pred_obj_pos.shape[3]
                    Nv_gt = gt_obj_pos.shape[3]
                    if Nv_pred != Nv_gt:
                        Nv_common = min(Nv_pred, Nv_gt)
                        pred_obj_pos = pred_obj_pos[:, :, :, :Nv_common, :]
                        gt_obj_pos = gt_obj_pos[:, :, :, :Nv_common, :]
                    # Masks at those steps (align vertex dim accordingly)
                    vmask = batch['valid_vtx_mask'][:, 1:seq_len-1]
                    if vmask.shape[3] != gt_obj_pos.shape[3]:
                        vmask = vmask[:, :, :, :gt_obj_pos.shape[3]]
                    obj_mask = batch['valid_obj_mask'].unsqueeze(1).unsqueeze(-1)  # (B,1,Nobj,1)
                    full_mask = vmask & obj_mask  # (B,T-2,Nobj,Nv)
                    # Position loss (denorm space)
                    pos_err = (pred_obj_pos - gt_obj_pos).pow(2).sum(dim=-1)  # (B,T-2,Nobj,Nv)
                    pos_loss_before_icp = (pos_err * full_mask).sum() / (full_mask.sum() + 1e-8)

                    # Acceleration losses (must match training computation!)
                    _model_ref = self.model.module if isinstance(self.model, DDP) else self.model
                    accel_loss_after_icp_ste = 0.0
                    accel_loss_before_icp = 0.0

                    if 'obj_accelerations_norm' in predictions:
                        a_ste = predictions['obj_accelerations_norm']
                        gt_prev = batch['obj_positions'][:, 0:seq_len-2]
                        gt_curr = batch['obj_positions'][:, 1:seq_len-1]
                        gt_next = batch['obj_positions'][:, 2:seq_len]
                        gt_accel = gt_next - 2 * gt_curr + gt_prev

                        # Multi-timescale training: optionally normalize acceleration error by step_size²
                        # (must match the training computation for consistent validation)
                        accel_err_ste = a_ste - gt_accel  # (B,T-2,Nobj,Nv,3)

                        # Apply step_size² normalization if enabled (match training)
                        if_time_step_scale = self.config['training']['loss_weights'].get('if_time_step_scale', True)
                        if if_time_step_scale:
                            if 'skip_step_size' not in batch:
                                raise KeyError(
                                    "Missing 'skip_step_size' in validation batch!\n"
                                    "Multi-timescale training with if_time_step_scale=true requires step_size."
                                )
                            step_sizes = batch['skip_step_size'].float()  # (B,)
                            Bsz_val = step_sizes.shape[0]
                            scale_full = (step_sizes ** 2).view(Bsz_val, 1, 1, 1, 1)  # (B,1,1,1,1)
                            accel_err_ste = accel_err_ste / scale_full  # Normalize by step_size²
                        # else: no normalization, use raw error

                        accel_err_ste_sq = accel_err_ste.pow(2).sum(dim=-1)  # (B,T-2,Nobj,Nv)
                        accel_loss_after_icp_ste = (accel_err_ste_sq * full_mask).sum() / (full_mask.sum() + 1e-8)

                    if 'obj_accelerations_no_icp_norm_anchor' in predictions and 'anchor_idx' in batch:
                        # Predicted anchor accelerations (raw): (B,T-2,Nobj,Nq,3)
                        a_no_icp_anchor = predictions['obj_accelerations_no_icp_norm_anchor']
                        Tm2 = a_no_icp_anchor.shape[1]
                        # Ground truth accelerations (full Nv)
                        gt_prev = batch['obj_positions'][:, 0:seq_len-2]
                        gt_curr = batch['obj_positions'][:, 1:seq_len-1]
                        gt_next = batch['obj_positions'][:, 2:seq_len]
                        gt_accel = gt_next - 2 * gt_curr + gt_prev  # (B,T-2,Nobj,Nv,3)
                        # Downsample GT accel to anchors for fair comparison
                        idx_exp = batch['anchor_idx'].unsqueeze(1).unsqueeze(-1).expand(-1, Tm2, -1, -1, 3)  # (B,T-2,Nobj,Nq,3)
                        gt_accel_anchor = torch.gather(gt_accel, dim=3, index=idx_exp)  # (B,T-2,Nobj,Nq,3)

                        # Multi-timescale training: optionally normalize acceleration error by step_size²
                        # (must match the training computation for consistent validation)
                        accel_err_no_icp = a_no_icp_anchor - gt_accel_anchor  # (B,T-2,Nobj,Nq,3)

                        # Apply step_size² normalization if enabled (match training)
                        if_time_step_scale = self.config['training']['loss_weights'].get('if_time_step_scale', True)
                        if if_time_step_scale:
                            if 'skip_step_size' not in batch:
                                raise KeyError(
                                    "Missing 'skip_step_size' in validation batch!\n"
                                    "Multi-timescale training with if_time_step_scale=true requires step_size."
                                )
                            step_sizes = batch['skip_step_size'].float()  # (B,)
                            Bsz_val = step_sizes.shape[0]
                            scale_anchor = (step_sizes ** 2).view(Bsz_val, 1, 1, 1, 1)  # (B,1,1,1,1)
                            accel_err_no_icp = accel_err_no_icp / scale_anchor  # Normalize by step_size²
                        # else: no normalization, use raw error

                        accel_err_no_icp_sq = accel_err_no_icp.pow(2).sum(dim=-1)  # (B,T-2,Nobj,Nq)
                        # Anchor mask derived from object validity only (anchors are prevalidated)
                        anchor_full_mask = batch['valid_obj_mask'].unsqueeze(1).unsqueeze(-1).expand(-1, Tm2, -1, a_no_icp_anchor.shape[3])
                        accel_loss_before_icp = (accel_err_no_icp_sq * anchor_full_mask).sum() / (anchor_full_mask.sum() + 1e-8)
                    # Edge length proxy (optional)
                    edge_loss = 0.0
                    if self.loss_weights.get('edge_length', 0.0) > 0:
                        init_obj = batch['obj_positions'][:, 0]  # (B,Nobj,Nv,3)
                        pred_mean = pred_obj_pos.mean(dim=1)      # (B,Nobj,Nv,3)
                        vmask0 = batch['valid_vtx_mask'][:, 0]    # (B,Nobj,Nv)
                        obj_mask_b = batch['valid_obj_mask'].unsqueeze(-1)  # (B,Nobj,1)
                        mask0 = vmask0 & obj_mask_b
                        ed = (pred_mean - init_obj).pow(2).sum(dim=-1)
                        edge_loss = (ed * mask0).sum() / (mask0.sum() + 1e-8)
                else:
                    logger.warning("Validation batch is missing multi-object tensors; skipping batch.")
                    continue
                    
                # Total loss (must match training weights!)
                loss = self.loss_weights['position_before_icp'] * pos_loss_before_icp
                if self.loss_weights.get('acceleration_after_icp', 0.0) > 0 and accel_loss_after_icp_ste > 0:
                    loss = loss + self.loss_weights['acceleration_after_icp'] * accel_loss_after_icp_ste
                if self.loss_weights.get('acceleration_before_icp', 0.0) > 0 and accel_loss_before_icp > 0:
                    loss = loss + self.loss_weights['acceleration_before_icp'] * accel_loss_before_icp
                if self.loss_weights.get('edge_length', 0.0) > 0:
                    loss = loss + self.loss_weights['edge_length'] * edge_loss

                total_loss += loss.item()
                total_pos_loss += pos_loss_before_icp.item() if isinstance(pos_loss_before_icp, torch.Tensor) else pos_loss_before_icp
                total_accel_loss += accel_loss_after_icp_ste.item() if isinstance(accel_loss_after_icp_ste, torch.Tensor) else accel_loss_after_icp_ste
                total_edge_loss += edge_loss.item() if isinstance(edge_loss, torch.Tensor) else edge_loss
                num_batches += 1
                
                pbar.set_postfix({'val_loss': loss.item()})
        
        avg_loss = total_loss / num_batches
        avg_pos_loss = total_pos_loss / num_batches
        avg_accel_loss = total_accel_loss / num_batches
        avg_edge_loss = total_edge_loss / num_batches

        # Synchronize validation metrics across all ranks in DDP mode
        if dist.is_available() and dist.is_initialized():
            metrics = torch.tensor([avg_loss, avg_pos_loss, avg_accel_loss, avg_edge_loss], device=self.device)
            dist.all_reduce(metrics, op=dist.ReduceOp.SUM)
            world_size = dist.get_world_size()
            metrics = metrics / max(world_size, 1)
            avg_loss, avg_pos_loss, avg_accel_loss, avg_edge_loss = metrics.tolist()

        # Log to tensorboard (only rank 0 in DDP)
        if self.is_main:
            # Log STE rigid alpha (always enabled)
            self.writer.add_scalar('val/rigid_alpha', 1.0, self.epoch)
            self.writer.add_scalar('val/loss', avg_loss, self.epoch)
            self.writer.add_scalar('val/position_loss_before_icp', avg_pos_loss, self.epoch)
            self.writer.add_scalar('val/accel_loss_after_icp_ste', avg_accel_loss, self.epoch)
            if self.loss_weights.get('edge_length', 0.0) > 0:
                self.writer.add_scalar('val/edge_loss', avg_edge_loss, self.epoch)

        return avg_loss

    def check_forward_rollout_consistency(self):
        """
        One-time check to verify forward and rollout produce consistent results.
        Called once after epoch 1 to ensure training and inference paths are aligned.
        """
        logger.info("=" * 80)
        logger.info("Running forward-rollout consistency check...")
        logger.info("=" * 80)

        # Get underlying model (unwrap DDP if needed)
        model_ref = self.model.module if isinstance(self.model, DDP) else self.model
        model_ref.eval()

        # Check if validation loader exists
        if self.val_loader is None:
            logger.info("Validation disabled - skipping consistency check")
            return

        # Get a batch from validation set
        val_iter = iter(self.val_loader)
        try:
            batch = next(val_iter)
        except StopIteration:
            logger.warning("No validation data available for consistency check")
            return

        # Move to device
        batch = {k: v.to(self.device, non_blocking=True) if isinstance(v, torch.Tensor) else v
                for k, v in batch.items()}

        # Flash Attention handles dtype internally in decoder.py - rest of model runs in FP32
        with torch.no_grad():
            # Set fixed seed for reproducibility
            torch.manual_seed(42)
            torch.cuda.manual_seed_all(42)

            # Run forward with student forcing to align with rollout (teacher_forcing_ratio=0.0)
            forward_output = model_ref(batch, teacher_forcing_ratio=0.0)

            # Prepare batch for rollout (needs only first 2 frames)
            B, T, Nobj, Nv, _ = batch['obj_positions'].shape
            logger.info(f"Consistency check: T={T}, predicting {T-2} step(s)")

            # Reset seed before rollout for exact match
            torch.manual_seed(42)
            torch.cuda.manual_seed_all(42)

            rollout_batch = {
                'obj_positions': batch['obj_positions'][:, :2],  # First 2 frames for warmup
                'valid_obj_mask': batch['valid_obj_mask'],
                'valid_vtx_mask': batch['valid_vtx_mask'][:, :2],  # First 2 frames
                'ref_obj_positions': batch['ref_obj_positions'],
                'anchor_idx': batch['anchor_idx'],
            }
            # Per-sample inputs the rollout reads besides positions and masks (physical parameters, step size, FPS
            # ranks), and dataset_name (penetration sign prediction)
            for _k in ('phys_params', 'skip_step_size', 'vtx_rank', 'dataset_name'):
                if _k in batch:
                    rollout_batch[_k] = batch[_k]

            # Run rollout (inference path) - predict same number of steps as forward
            num_steps_to_predict = T  # Total frames (includes 2 warmup)
            rollout_output = model_ref.rollout(rollout_batch, num_steps=num_steps_to_predict)

            # Extract predictions for comparison
            # Compare STE-blended positions (used by rollout autoregression) to ensure path alignment
            # Forward produces predictions for frames [2, T-1] -> shape (B, T-2, Nobj, Nv, 3)
            forward_pred = forward_output.get('obj_positions_ste', forward_output['obj_positions_denorm'])

            # Rollout produces full sequence including warmup -> shape (B, T, Nobj, Nv, 3)
            # Extract predicted frames [2, T-1] to match forward
            rollout_pred = rollout_output['obj_positions'][:, 2:T]  # (B, T-2, Nobj, Nv, 3)

            logger.info("=" * 80)
            logger.info("Step-by-step comparison:")
            logger.info("=" * 80)

            # Compare input frames (should be identical)
            input_frame0_fw = batch['obj_positions'][:, 0]
            input_frame1_fw = batch['obj_positions'][:, 1]
            input_frame0_ro = rollout_batch['obj_positions'][:, 0]
            input_frame1_ro = rollout_batch['obj_positions'][:, 1]

            diff_frame0 = torch.abs(input_frame0_fw - input_frame0_ro).max().item()
            diff_frame1 = torch.abs(input_frame1_fw - input_frame1_ro).max().item()
            logger.info(f"Input frame 0 difference: {diff_frame0:.6e}")
            logger.info(f"Input frame 1 difference: {diff_frame1:.6e}")

            # Compare ref_pose
            ref_diff = torch.abs(batch['ref_obj_positions'] - rollout_batch['ref_obj_positions']).max().item()
            logger.info(f"Reference pose difference: {ref_diff:.6e}")

            # Compare outputs for each predicted frame
            num_pred_frames = forward_pred.shape[1]
            logger.info(f"Number of predicted frames: {num_pred_frames}")

            for frame_idx in range(num_pred_frames):
                fw_frame = forward_pred[:, frame_idx]  # (B, Nobj, Nv, 3)
                ro_frame = rollout_pred[:, frame_idx]  # (B, Nobj, Nv, 3)
                frame_diff = torch.abs(fw_frame - ro_frame)

                # Get mask for this frame (frame_idx corresponds to frame 2+frame_idx in original batch)
                actual_frame_idx = 2 + frame_idx
                if actual_frame_idx < batch['valid_vtx_mask'].shape[1]:
                    vmask = batch['valid_vtx_mask'][:, actual_frame_idx]  # (B, Nobj, Nv)
                    obj_mask = batch['valid_obj_mask'].unsqueeze(-1)  # (B, Nobj, 1)
                    full_mask = vmask & obj_mask
                    masked_diff = frame_diff * full_mask.unsqueeze(-1)
                    masked_max = masked_diff.max().item()
                    masked_mean = masked_diff.sum().item() / (full_mask.sum().item() * 3 + 1e-8)
                else:
                    masked_max = frame_diff.max().item()
                    masked_mean = frame_diff.mean().item()

                logger.info(f"  Frame {actual_frame_idx}: max_diff={frame_diff.max().item():.6e}, "
                           f"masked_max={masked_max:.6e}, masked_mean={masked_mean:.6e}")

                # Print per-object statistics for first predicted frame
                if frame_idx == 0:
                    logger.info(f"    Per-object breakdown for frame {actual_frame_idx}:")
                    for obj_idx in range(min(Nobj, 5)):  # Only first 5 objects
                        obj_diff = frame_diff[:, obj_idx].max().item()
                        logger.info(f"      Object {obj_idx}: max_diff={obj_diff:.6e}")

            logger.info("=" * 80)

            # Compute difference
            diff = torch.abs(forward_pred - rollout_pred)
            max_diff = diff.max().item()
            mean_diff = diff.mean().item()

            # Get masks for valid comparison
            obj_mask = batch['valid_obj_mask'].unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # (B,1,Nobj,1,1)
            vtx_mask = batch['valid_vtx_mask'][:, 2:T].unsqueeze(-1)  # (B,T-2,Nobj,Nv,1)
            full_mask = (obj_mask & vtx_mask).float()

            # Masked statistics
            masked_diff = diff * full_mask
            num_valid = full_mask.sum().item()
            if num_valid > 0:
                masked_mean_diff = masked_diff.sum().item() / num_valid
                masked_max_diff = masked_diff.max().item()
            else:
                masked_mean_diff = 0.0
                masked_max_diff = 0.0

            # Log results
            logger.info(f"Consistency Check Results:")
            logger.info(f"  Forward output shape: {forward_pred.shape}")
            logger.info(f"  Rollout output shape: {rollout_pred.shape}")
            logger.info(f"  Max absolute difference: {max_diff:.6e}")
            logger.info(f"  Mean absolute difference: {mean_diff:.6e}")
            logger.info(f"  Masked max difference: {masked_max_diff:.6e}")
            logger.info(f"  Masked mean difference: {masked_mean_diff:.6e}")
            logger.info(f"  Number of valid elements: {int(num_valid)}")

            # Tolerance check (allow small numerical errors from GPU operations and accumulation)
            # Relaxed from 1e-4 to 1e-3 to account for autoregressive error accumulation
            tolerance = 1e-3
            if masked_max_diff > tolerance:
                error_msg = (
                    f"\n{'='*80}\n"
                    f"Consistency check FAILED.\n"
                    f"Forward and rollout outputs differ by {masked_max_diff:.6e} (tolerance: {tolerance:.6e})\n"
                    f"This indicates train and inference paths are not aligned.\n"
                    f"{'='*80}\n"
                )
                logger.error(error_msg)
                raise RuntimeError(error_msg)
            else:
                logger.info("=" * 80)
                logger.info("Consistency check passed - forward and rollout are aligned")
                logger.info("=" * 80)

        # Restore training mode
        model_ref.train()

    def generate_visualization_video(self):
        """Generate visualization videos for multiple datasets (configurable per split)."""
        logger.info(f"generate_visualization_video called for epoch {self.epoch}")
        # Only rank 0 should generate videos in DDP
        if dist.is_available() and dist.is_initialized():
            try:
                if dist.get_rank() != 0:
                    return
            except Exception:
                pass

        if not self.config['training']['checkpoint'].get('save_video', False):
            logger.info("Video generation disabled in config")
            return

        video_every = self.config['training']['checkpoint'].get('video_every_n_epochs', 10)
        logger.info(f"video_every_n_epochs = {video_every}, epoch = {self.epoch}, should generate = {self.epoch % video_every == 0}")

        if self.epoch % video_every != 0:
            return

        logger.info(f"Generating visualization videos for epoch {self.epoch}")
        _video_failures = []

        # Get number of videos per split from config (default to 3 for backwards compatibility)
        num_videos_cfg = self.config['training']['checkpoint'].get('num_videos_per_split', {})
        num_train_videos = num_videos_cfg.get('train', 3)
        num_val_videos = num_videos_cfg.get('val', 0)
        num_test_videos = num_videos_cfg.get('test', 0)

        logger.info(f"Video config: {num_train_videos} train videos, {num_val_videos} val videos, {num_test_videos} test videos")

        # Load dataset split info to know which datasets are train/val/test
        split_file = self.config.get('data', {}).get('split_file', 'dataset_split.json')
        if os.path.exists(split_file):
            with open(split_file, 'r') as f:
                split_info = json.load(f)

            # Read all dataset lists
            train_datasets_all = split_info.get('train_datasets', [])
            val_datasets_all = split_info.get('val_datasets', [])
            test_datasets_all = split_info.get('test_datasets', [])

            # Select datasets for visualization based on config
            train_datasets = []
            val_datasets = []
            test_datasets = []

            # Train videos
            if num_train_videos > 0 and len(train_datasets_all) > 0:
                train_datasets = train_datasets_all[:num_train_videos]
                logger.info(f"Train videos: {len(train_datasets)} datasets selected")

            # Val videos (only if validation is enabled)
            if num_val_videos > 0 and len(val_datasets_all) > 0:
                if self.config['training'].get('enable_validation', False):
                    val_datasets = val_datasets_all[:num_val_videos]
                    logger.info(f"Val videos: {len(val_datasets)} datasets selected (validation enabled)")
                else:
                    logger.info(f"Val videos: skipped (validation disabled, enable_validation=false)")

            # Test videos (always available if datasets exist)
            if num_test_videos > 0 and len(test_datasets_all) > 0:
                test_datasets = test_datasets_all[:num_test_videos]
                logger.info(f"Test videos: {len(test_datasets)} datasets selected")

            # Create list of (dataset_name, split_type) tuples
            datasets_to_visualize = (
                [(d, 'train') for d in train_datasets] +
                [(d, 'val') for d in val_datasets] +
                [(d, 'test') for d in test_datasets]
            )
            logger.info(f"Visualizing datasets: {[f'{d[0]}({d[1]})' for d in datasets_to_visualize]}")
        else:
            raise RuntimeError(f"Split file not found: {split_file}")
        
        # Multi-timescale video generation: loop over eval_step_sizes
        # (falls back to data.skip_step_size if not set)
        eval_step_sizes = self.config['data'].get('eval_step_sizes', [self.config['data'].get('skip_step_size', 1)])

        logger.info(f"Video generation: evaluating at step_sizes {eval_step_sizes}")

        # Optional debug logging (data.random_step_size.if_debug)
        if_debug = self.config.get('data', {}).get('random_step_size', {}).get('if_debug', False)
        if if_debug:
            logger.info(f"[DEBUG Video Gen] eval_step_sizes={eval_step_sizes}, total datasets={len(datasets_to_visualize)}")

        # Put model in eval mode
        self.model.eval()

        # Outer loop: iterate over each step_size
        for step_size in eval_step_sizes:
            logger.info(f"\n{'='*80}")
            logger.info(f"[Video Generation] step_size = {step_size}")
            logger.info(f"{'='*80}")

            # Optional debug logging
            if if_debug:
                logger.info(f"[DEBUG Video Gen] Starting step_size={step_size} iteration ({eval_step_sizes.index(step_size)+1}/{len(eval_step_sizes)})")

            # Number of frames to generate for this step_size
            base_num_frames = self.config['training']['checkpoint'].get('video_num_frames', 50)
            num_frames = base_num_frames // step_size + 1

            # Inner loop: generate video for each dataset at this step_size
            for dataset_name, split_type in datasets_to_visualize:
                try:
                    logger.info(f"Generating video for {dataset_name} ({split_type}) with step_size={step_size}")

                    # Extract dataset number for video naming (e.g., "dataset_aug_100" -> "100")
                    dataset_num = dataset_name.split('_')[-1]

                    # One video per step size; visualization.if_overlay in the
                    # config adds a transparent ground-truth overlay when true.
                    video_path = self.log_dir / f'inference_epoch_{self.epoch:04d}_dataset_{dataset_num}_{split_type}_size_{step_size}.mp4'
                    inference_fps = self._create_video_multiobj_pred(dataset_name, video_path, split=split_type, step_size=step_size)
                    if video_path.exists():
                        file_size_kb = video_path.stat().st_size / 1024
                        logger.info(f"Video saved to {video_path} ({file_size_kb:.1f} KB, Inference: {inference_fps:.2f} FPS)")
                        self.writer.add_text('visualization', f'[Video](file://{video_path})', self.epoch)

                except Exception as e:
                    logger.error(f"Visualization error for {dataset_name} (step_size={step_size}): {e}")
                    import traceback
                    logger.error(traceback.format_exc())
                    _video_failures.append(f"{dataset_name} s={step_size}: {e}")
                    continue  # Continue with next dataset

        # Put model back in train mode after all videos
        self.model.train()
        logger.info(f"\nCompleted video generation for epoch {self.epoch} (all step_sizes)")
        if _video_failures and (self.config.get('inference', {}) or {}).get('mode', False):
            # a requested video that failed is an error of the run (scripts/render_videos.sh relies on the exit status);
            # during training the videos are diagnostics and the errors above are logged
            raise RuntimeError(f"{len(_video_failures)} video(s) failed: {_video_failures[:3]}")

    def _create_video_multiobj_gt(self, obj_positions_seq, faces_list, output_path, fps=100):
        """Render multi-object GT meshes per frame using provided faces (indices per object)."""
        try:
            os.environ.setdefault('PYOPENGL_PLATFORM', 'egl')
            os.environ.setdefault('POLYSCOPE_BACKEND', 'openGL3_egl')
            os.environ.setdefault('POLYSCOPE_ALLOW_HEADLESS', '1')
            import polyscope as ps
            import numpy as np
            import subprocess
        except ImportError:
            raise RuntimeError("Polyscope not available; install polyscope or disable save_video")

        frames_dir = self.log_dir / f'temp_frames_epoch_{self.epoch}_mobj'
        frames_dir.mkdir(exist_ok=True)
        try:
            if hasattr(ps, 'set_allow_headless'):
                ps.set_allow_headless(True)
            if hasattr(ps, 'set_backend'):
                ps.set_backend('openGL3_egl')
            ps.init('openGL3_egl')
        except Exception as e:
            raise RuntimeError(f"Failed to initialize Polyscope headless backend: {e}")
        # Env overrides for visualization
        try:
            _w = int(os.environ.get('WIDTH', '0'))
            _h = int(os.environ.get('HEIGHT', '0'))
        except Exception:
            _w, _h = 0, 0
        if _w > 0 and _h > 0 and hasattr(ps, 'set_window_size'):
            try:
                ps.set_window_size(_w, _h)
            except Exception:
                pass
        _fps_env = os.environ.get('FPS', '').strip()
        if _fps_env:
            try:
                fps = float(_fps_env)
            except Exception:
                pass
        try:
            _edge_w = float(os.environ.get('EDGE_WIDTH', 'nan'))
        except Exception:
            _edge_w = float('nan')
        ps.set_autocenter_structures(False)
        ps.set_ground_plane_mode("none")
        ps.set_ground_plane_height(0.0)
        ps.set_navigation_style('turntable')
        ps.set_up_dir('z_up')

        # Camera settings from config
        checkpoint_config = self.config['training']['checkpoint']
        camera_config = checkpoint_config.get('camera', {})
        camera_mode = camera_config.get('mode', 'auto')

        first = obj_positions_seq[0]
        all_pts = np.concatenate(first, axis=0)

        if camera_mode == 'fixed':
            # Use fixed camera from config
            center = np.array(camera_config.get('center', [0.0, 0.0, 0.5]), dtype=np.float64)
            eye = np.array(camera_config.get('eye', [3.0, -3.0, 3.0]), dtype=np.float64)
        else:
            # Auto camera: compute from scene bounds
            center = np.mean(all_pts, axis=0)
            bbox_min = np.min(all_pts, axis=0)
            bbox_max = np.max(all_pts, axis=0)
            radius = 0.5 * np.linalg.norm(bbox_max - bbox_min)
            auto_direction = np.array(camera_config.get('auto_direction', [2.0, -2.0, 2.0]), dtype=np.float64)
            auto_scale = camera_config.get('auto_scale', 3.0)
            eye = center + auto_direction * max(radius, 1e-3) * auto_scale

        # Ground plane settings from config
        ground_config = checkpoint_config.get('ground', {})
        ground_size = ground_config.get('size', 10.0)
        ground_color = ground_config.get('color', [0.8, 0.8, 0.8])
        ground_vertices = np.array([
            [-ground_size, -ground_size, 0.0],
            [ ground_size, -ground_size, 0.0],
            [ ground_size,  ground_size, 0.0],
            [-ground_size,  ground_size, 0.0]
        ])
        ground_faces = np.array([[0, 1, 2], [0, 2, 3]])

        for t, objs in enumerate(obj_positions_seq):
            # Register ground first
            ps.register_surface_mesh("ground", ground_vertices, ground_faces, enabled=True)
            ps.get_surface_mesh("ground").set_color(ground_color)
            try:
                if _edge_w == _edge_w:
                    ps.get_surface_mesh("ground").set_edge_width(_edge_w)
            except Exception:
                pass
            for i, V in enumerate(objs):
                F = faces_list[i]
                ps.register_surface_mesh(f"obj_{i}", V, F, enabled=True, smooth_shade=True)
                try:
                    if _edge_w == _edge_w:
                        ps.get_surface_mesh(f"obj_{i}").set_edge_width(_edge_w)
                        ps.get_surface_mesh(f"obj_{i}").set_edge_color([0.0, 0.0, 0.0])
                except Exception:
                    pass
            ps.look_at(eye, center)  # ps.look_at(camera_location, target)
            ps.screenshot(str(frames_dir / f'frame_{t:05d}.png'))
            ps.remove_all_structures()
            if t % 10 == 0:
                logger.info(f"  [multiobj] Generated frame {t}/{len(obj_positions_seq)}")

        ffmpeg_cmd = [
            'ffmpeg', '-hide_banner', '-loglevel', 'error',
            '-r', str(fps),
            '-i', str(frames_dir / 'frame_%05d.png'),
            '-vf', 'scale=trunc(iw/2)*2:trunc(ih/2)*2',
            '-c:v', 'mpeg4', '-q:v', '2', '-pix_fmt', 'yuv420p', '-y',
            str(output_path)
        ]
        subprocess.run(ffmpeg_cmd, check=True)

        # Clean up temporary frames
        import shutil
        shutil.rmtree(frames_dir, ignore_errors=True)
        logger.info(f"  [multiobj] Cleaned up temp frames: {frames_dir}")

    def _render_video_filament(self, ds_dir, batch, pred_frames, gt_frames, vmask_all, output_path, fps, times):
        """Render a rollout with src/utils/rollout_render.py. Every object is drawn with its visual mesh, moved by the
        rigid motion fitted to its collision-mesh vertices (predicted: opaque; ground truth, when given: translucent).
        `times` (s) of the key frames: the video plays them in real time at `fps`, interpolating in between."""
        import json
        from src.utils import rollout_render
        mesh_dir = self.config['training']['checkpoint'].get('video_meshes', 'data_generation/objects')
        info = json.loads((Path(ds_dir) / 'physics_params.json').read_text())
        objects = [info[f'obj_{i}'] for i in range(sum(k.startswith('obj_') for k in info))]
        if 'permutation' in batch and batch['permutation'] is not None:
            perm = batch['permutation'][0]
            perm = perm.cpu().numpy() if isinstance(perm, torch.Tensor) else np.asarray(perm)
            objects = [objects[int(k)] for k in perm]
        if pred_frames.shape[1] != len(objects):
            raise RuntimeError(f'{ds_dir}: {pred_frames.shape[1]} predicted objects, physics_params.json lists '
                               f'{len(objects)}')
        shapes = [o['shape'] for o in objects]
        vm = None if vmask_all is None else np.asarray(vmask_all[0], bool)            # (Nobj, Nv)

        def poses(frames):
            out = []
            for i, o in enumerate(objects):
                ref = float(o['size']) * rollout_render.collision_vertices(mesh_dir, o['shape'])
                nv = len(ref)
                if vm is not None and not (vm[i, :nv].all() and not vm[i, nv:].any()):
                    raise RuntimeError(f"{ds_dir}: object {i} has {int(vm[i].sum())} valid vertices, its "
                                       f"{o['shape']} collision mesh {nv}")
                P, resid = rollout_render.fit_poses(ref, np.asarray(frames[:, i, :nv], np.float64))
                if resid > 1e-3:
                    raise RuntimeError(f"{ds_dir}: object {i} is not a rigid motion of its {o['shape']} collision "
                                       f"mesh at size {o['size']:.4f} (largest residual {resid:.2e} m)")
                out.append(P)
            return out

        visual = rollout_render.load_render_meshes(mesh_dir, shapes)
        self._filament_renderer().render(
            output_path, [visual[s] for s in shapes], [o['size'] for o in objects],
            [o.get('material') == 'metal' for o in objects], poses(pred_frames), fps,
            ghost=None if gt_frames is None else poses(np.asarray(gt_frames)), times=np.asarray(times, np.float64))
        logger.info(f"  [video] {output_path} ({len(pred_frames)} predicted frames, real time at {fps:g} fps)")

    def _create_video_multiobj_pred(self, dataset_name: str, output_path, fps=100, split='val', step_size: int = None, override_if_overlay: int = None):
        """Render predicted per-object vertices if batch/masks are available for first sample.
        This function will run one batch from the dataloader for the given dataset name and render
        the model predictions with simple point rendering per object (different colors).

        Args:
            dataset_name: Name of dataset to visualize
            output_path: Path to save video
            fps: Frames per second for video
            split: 'train', 'val', or 'test'
            step_size: Frame skip step size for multi-timescale (defaults to config skip_step_size)
            override_if_overlay: If provided (0 or 1), override config's if_overlay setting

        Returns:
            float: Inference FPS (frames predicted per second)
        """
        # Multi-Timescale: Determine skip_step_size early (needed for dataset creation)
        if step_size is None:
            skip_step_size = int(self.config['data'].get('skip_step_size', 1))
        else:
            skip_step_size = int(step_size)

        # Optional debug logging (data.random_step_size.if_debug)
        if_debug = self.config.get('data', {}).get('random_step_size', {}).get('if_debug', False)
        if if_debug:
            logger.info(f"[DEBUG Video Gen] Starting video generation for dataset='{dataset_name}', split='{split}', step_size={skip_step_size}")

        # training.checkpoint.video_renderer: 'filament' (default; src/utils/rollout_render.py, physically based, the objects'
        # visual meshes) or 'polyscope' (the earlier diagnostic view of the collision meshes)
        video_renderer = self.config['training']['checkpoint'].get('video_renderer', 'filament')
        if video_renderer not in ('filament', 'polyscope'):
            raise ValueError(f"training.checkpoint.video_renderer must be 'filament' or 'polyscope', got {video_renderer!r}")
        import numpy as np
        import random
        from pathlib import Path
        if video_renderer == 'polyscope':
            try:
                os.environ.setdefault('PYOPENGL_PLATFORM', 'egl')
                os.environ.setdefault('POLYSCOPE_BACKEND', 'openGL3_egl')
                os.environ.setdefault('POLYSCOPE_ALLOW_HEADLESS', '1')
                import polyscope as ps
            except ImportError:
                raise RuntimeError("Polyscope not available; install polyscope or disable save_video")

        # ============================================================
        # Create a separate dataset per step_size.
        # Training datasets sample a random step_size per sample (multi-timescale),
        # while video generation requires a fixed step_size to match the network
        # conditioning, so an independent dataset is created and cached per
        # (split, step_size) pair. Cache key: {step_size: dataset} per split.
        # ============================================================

        from src.data.multi_dataloader import MultiObjectDynamics

        # Select cache dictionary based on split
        if split == 'train':
            cache_dict = self._train_video_datasets
        elif split == 'val':
            if self.val_loader is None:
                raise RuntimeError(f"Cannot generate val video: validation is disabled (enable_validation=false)")
            cache_dict = self._val_video_datasets
        elif split == 'test':
            cache_dict = self._test_video_datasets
        else:
            raise ValueError(f"Unknown split: {split}")

        # Check cache, create dataset if not exists
        if skip_step_size not in cache_dict:
            logger.info(f"Creating {split} dataset for step_size={skip_step_size} (first time, will be cached)")

            # Disable random_step_size to get fixed step_size sequences
            augmentation_config = self.config['data'].get('augmentation', {}).copy()
            # Ensure no rotation/permutation augmentation for video generation (deterministic)
            if 'rotation_augmentation' in augmentation_config:
                augmentation_config['rotation_augmentation'] = augmentation_config['rotation_augmentation'].copy()
                augmentation_config['rotation_augmentation']['enable'] = False
            if 'anchor_resampling' in augmentation_config:
                augmentation_config['anchor_resampling'] = dict(augmentation_config['anchor_resampling'], enable=False)
            if 'object_permutation' in augmentation_config:
                augmentation_config['object_permutation'] = augmentation_config['object_permutation'].copy()
                augmentation_config['object_permutation']['enable'] = False

            cache_dict[skip_step_size] = MultiObjectDynamics(
                data_path=self.config['data']['dataset_path'],
                split_file=self.config['data']['split_file'],
                split=split,
                sequence_length=self.config['data']['sequence_length'],
                skip_step_size=skip_step_size,  # Fixed step_size (NOT random!)
                random_step_size={'enable': False},  # Disable random sampling (use fixed skip_step_size)
                max_vertices_per_object=self.config['data']['max_vertices_per_object'],
                max_objects_per_frame=self.config['data']['max_objects_per_frame'],
                anchor_k=self.config['data'].get('anchor_k', 32),
                skip_first_frames=self.config['data'].get('skip_first_frames', 0),
                augmentation=augmentation_config,
                anchor_fps_start=self.config['data'].get('anchor_fps_start', 'index0'),
                fps_order=self.config['data'].get('fps_order', False),
                object_scale_file=self.config['data'].get('object_scale_file'),
                if_phys_param=self.config.get('model', {}).get('if_phys_param', False)
            )
            logger.info(f"{split.capitalize()} dataset (step_size={skip_step_size}) created and cached: "
                       f"{len(cache_dict[skip_step_size].all_sequences)} sequences")
        else:
            logger.info(f"Reusing cached {split} dataset for step_size={skip_step_size} "
                       f"({len(cache_dict[skip_step_size].all_sequences)} sequences)")

        ds = cache_dict[skip_step_size]

        # Pick first sequence that belongs to dataset_name
        sample_idx = None
        for i, s in enumerate(ds.all_sequences):
            if s['dataset'] == dataset_name:
                sample_idx = i
                break
        if sample_idx is None:
            raise RuntimeError(f"No sequence found for dataset {dataset_name}")

        sample = ds[sample_idx]
        # Physical frame of the sequence's first frame (data.skip_first_frames shifts it): the rollout and the
        # ground-truth overlay start there
        if 'frame_ids' not in sample or sample['frame_ids'].numel() == 0:
            raise KeyError("video generation needs the sample's frame_ids (physical frame of the first frame)")
        video_start_frame = int(sample['frame_ids'][0])
        # Move tensors to device and handle special keys (dataset_name, faces_list)
        to_dev = lambda x: x.unsqueeze(0).to(self.device) if isinstance(x, torch.Tensor) else x
        batch = {}
        for k, v in sample.items():
            if k == 'dataset_name':
                # dataset_name must be a list (as in training collate_fn)
                batch[k] = [v] if isinstance(v, str) else v
            elif k == 'faces_list':
                # faces_list is a list of arrays per object, wrap it to simulate collate_fn behavior
                # collate_fn does: [sample0['faces_list'], sample1['faces_list'], ...]
                # We need: [sample['faces_list']] (batch size 1)
                batch[k] = [v]
            elif k == 'permutation':
                # permutation is a numpy array, wrap it to simulate collate_fn behavior
                batch[k] = [v]
            else:
                batch[k] = to_dev(v)

        # Decide rollout length & output FPS to match source duration
        ds_dir = Path(self.config['data']['dataset_path']) / dataset_name
        simparams_path = ds_dir / 'simulation_parameters.json'
        # Prefer explicit config override for original fps
        fps_orig = 240.0
        try:
            fps_cfg = self.config['training']['checkpoint'].get('original_video_fps', None)
            if fps_cfg is not None:
                fps_orig = float(fps_cfg)
            elif simparams_path.exists():
                with open(simparams_path, 'r') as f:
                    simp = json.load(f)
                fps_orig = float(simp.get('fps', fps_orig))
        except Exception:
            pass

        # skip_step_size already set at the beginning of the function

        # Prefer explicit config override for original frame count
        cfg_base = self.config['training']['checkpoint'].get('original_video_num_frames', None)
        if cfg_base is None:
            try:
                T = len(list(ds_dir.glob('frame_*')))
                base_num_frames = max(0, T)
            except Exception:
                base_num_frames = 480
        else:
            base_num_frames = int(cfg_base)

        # Two-stage sampling for efficiency:
        # Stage 1: Rollout prediction with skip_step_size (e.g., 480//1=480 or 480//10=48)
        # Stage 2: Render subsampling to fixed 48 frames for video
        s = max(1, skip_step_size)
        # Full remaining sequence after skip_step_size, from the sequence's first physical frame
        num_frames_rollout = (max(1, base_num_frames - video_start_frame) + s - 1) // s

        if video_renderer == 'filament':
            # every rollout frame is a key frame: the renderer samples the motion in real time at
            # training.checkpoint.video_fps, interpolating between key frames (step sizes above 1)
            render_subsample = 1
        else:
            # Calculate render subsampling rate to get ~48 frames
            target_render_frames = 48
            render_subsample = max(1, num_frames_rollout // target_render_frames)  # Subsample rollout for rendering
        num_frames_render = (num_frames_rollout + render_subsample - 1) // render_subsample  # Actual render frames

        # Use real-time duration from original data
        # original_duration = base_num_frames / fps_orig (e.g., 600 / 100 = 6s)
        original_duration = base_num_frames / fps_orig
        if video_renderer == 'filament':
            fps_out = float(self.config['training']['checkpoint'].get('video_fps', 30))
        else:
            fps_out = num_frames_render / original_duration  # Maintain real-time playback

        logger.info(
            f"[Sampling Strategy] base={base_num_frames} frames, skip_step_size={skip_step_size} "
            f"-> rollout={num_frames_rollout} frames, render_subsample={render_subsample} "
            f"-> render={num_frames_render} frames @ {fps_out:.2f} fps ({original_duration:.1f}s duration)"
        )

        # Prepare rollout using initial GT frames from the sampled sequence
        _model_ref = self.model.module if isinstance(self.model, DDP) else self.model

        # Apply GPU augmentation (rotation) if enabled for visualization
        # This ensures the visualization shows what the training sees
        rotation_angle_deg = 0.0  # Default: no rotation
        if (self.gpu_augmentation is not None and self.gpu_augmentation.rotation_enabled
                and not self.gpu_augmentation.rotation_per_sample):
            # Apply rotation to batch (modifies obj_positions and ref_obj_positions); per-sample rotations, reflections
            # and input noise are not applied here (one angle is what the ground-truth overlay can follow)
            batch = self.gpu_augmentation.apply_rotation_augmentation(batch, split=split)
            rotation_angle_deg = batch.get('rotation_angle_deg', 0.0)
            if rotation_angle_deg != 0.0:
                logger.info(f"[Video Gen] Applied rotation augmentation: {rotation_angle_deg:.0f}° for verification")

        # Override batch['skip_step_size'] with the video-generation parameter.
        # Dataset samples carry the training step_size (randomly sampled in
        # multi-timescale mode), while video generation evaluates the model at a
        # specific temporal resolution, so the sampled value must be replaced.
        if 'skip_step_size' in batch:
            original_step_size = batch['skip_step_size'].item() if isinstance(batch['skip_step_size'], torch.Tensor) else batch['skip_step_size']
            if if_debug:
                logger.info(f"[DEBUG Video Gen] Overriding batch skip_step_size: {original_step_size} -> {skip_step_size}")
            else:
                logger.info(f"[Video Gen] Overriding batch skip_step_size: {original_step_size} -> {skip_step_size}")
        batch['skip_step_size'] = torch.tensor([skip_step_size], dtype=torch.float32, device=self.device)

        # Optional debug logging of the final step_size before rollout
        if if_debug:
            final_step_size = batch['skip_step_size'].item() if isinstance(batch['skip_step_size'], torch.Tensor) else batch['skip_step_size']
            logger.info(f"[DEBUG Video Gen] Final step_size before rollout: {final_step_size}, num_frames_rollout={num_frames_rollout}")

        # Measure inference speed (FPS)
        import time
        torch.cuda.synchronize() if torch.cuda.is_available() else None
        rollout_start_time = time.perf_counter()

        # Get eval_warmup_frames from config (consistent with evaluator.py)
        eval_warmup_frames = self.config['training'].get('metrics_test', {}).get('eval_warmup_frames', 2)

        # Flash Attention handles dtype internally in decoder.py - rest of model runs in FP32 for stability
        with torch.no_grad():
            rollout_out = _model_ref.rollout(
                initial_batch=batch,
                num_steps=num_frames_rollout,
                dt=1.0,
                warmup_frames=eval_warmup_frames
            )

        torch.cuda.synchronize() if torch.cuda.is_available() else None
        rollout_end_time = time.perf_counter()
        rollout_total_time = rollout_end_time - rollout_start_time

        # Calculate inference FPS (frames predicted per second)
        # num_frames_rollout includes eval_warmup_frames GT warmup frames
        num_predicted_frames = max(1, num_frames_rollout - eval_warmup_frames)
        inference_fps = num_predicted_frames / rollout_total_time if rollout_total_time > 0 else 0.0

        logger.info(
            f"[Inference Speed] dataset={dataset_name} split={split} | "
            f"warmup_frames={eval_warmup_frames} | "
            f"Total rollout frames: {num_frames_rollout} | Predicted frames: {num_predicted_frames} | "
            f"Rollout time: {rollout_total_time:.4f}s | "
            f"Inference FPS: {inference_fps:.2f} frames/sec | "
            f"Time per frame: {rollout_total_time/num_predicted_frames*1000:.2f}ms"
        )

        # Ensure FP32 for numpy/Polyscope (BF16 can cause visualization issues)
        pred_full = rollout_out['obj_positions'][0].float().cpu().numpy()  # (T, Nobj, Nv, 3)
        Ttot, Nobj, Nv, _ = pred_full.shape

        # Apply render subsampling to rollout predictions for video efficiency
        # Generate frame indices: uniformly sample from full rollout sequence
        render_frame_indices = np.arange(0, Ttot, render_subsample)[:num_frames_render]
        actual_render_frames = len(render_frame_indices)

        logger.info(
            f"[Render Sampling] Rollout has {Ttot} frames, rendering {actual_render_frames} frames "
            f"(subsample={render_subsample}, indices: {render_frame_indices[0]}..{render_frame_indices[-1]})"
        )

        # Subsample predictions for rendering
        pred_full_render = pred_full[render_frame_indices]  # (num_frames_render, Nobj, Nv, 3)

        # Get predicted and GT accelerations for visualization
        # Note: accelerations are computed at frame centers, so T frames -> T-2 accelerations
        pred_accel_anchor_seq = None    # (T_rollout-2, Nobj, Nq, 3) anchors only
        gt_accel_seq = None             # (T0-2, Nobj, Nv, 3)
        pred_accel_anchor_seq_render = None  # Subsampled for rendering
        gt_accel_seq_render = None           # Subsampled for rendering

        try:
            # Predicted anchor accelerations from rollout (denormalized)
            if 'obj_accelerations_anchor' not in rollout_out:
                raise RuntimeError("Rollout output missing 'obj_accelerations_anchor' for visualization")
            pred_accel_anchor_seq = rollout_out['obj_accelerations_anchor'][0].cpu().numpy()  # (T_rollout-2, Nobj, Nq, 3)

            # Subsample accelerations for rendering
            # Acceleration at index i corresponds to frames i, i+1, i+2 (centered at i+1)
            # For render_frame_indices, we need accelerations at indices render_frame_indices[t]-1 (if t >= 1)
            accel_indices_render = []
            for idx in render_frame_indices:
                if idx >= 1:  # Frame idx has acceleration at idx-1
                    accel_idx = idx - 1
                    if accel_idx < pred_accel_anchor_seq.shape[0]:
                        accel_indices_render.append(accel_idx)
            if len(accel_indices_render) > 0:
                pred_accel_anchor_seq_render = pred_accel_anchor_seq[accel_indices_render]
                logger.info(f"[Acceleration Sampling] Pred accel: {pred_accel_anchor_seq.shape} -> {pred_accel_anchor_seq_render.shape}")
            else:
                pred_accel_anchor_seq_render = None
        except Exception as e:
            logger.error(f"Failed to prepare predicted accelerations for visualization: {e}")
            pred_accel_anchor_seq_render = None

        try:
            # GT accelerations: compute from ground truth positions
            # batch['obj_positions'] shape: (B, T0, Nobj, Nv, 3)
            gt_pos = batch['obj_positions'][0].cpu().numpy()  # (T0, Nobj, Nv, 3)
            if gt_pos.shape[0] >= 3:
                # Compute GT accelerations: a = (x_{t+1} - 2*x_t + x_{t-1}) / dt^2
                # For frames 0,1,2,...,T0-1, we can compute accel at centers 1,2,...,T0-2
                gt_accel_seq = gt_pos[2:] - 2*gt_pos[1:-1] + gt_pos[:-2]  # (T0-2, Nobj, Nv, 3)

                # Subsample GT accelerations (same logic as predicted accelerations)
                gt_accel_indices_render = []
                for idx in render_frame_indices:
                    if idx >= 1:
                        accel_idx = idx - 1
                        if accel_idx < gt_accel_seq.shape[0]:
                            gt_accel_indices_render.append(accel_idx)
                if len(gt_accel_indices_render) > 0:
                    gt_accel_seq_render = gt_accel_seq[gt_accel_indices_render]
                    logger.info(f"[Acceleration Sampling] GT accel: {gt_accel_seq.shape} -> {gt_accel_seq_render.shape}")
                else:
                    gt_accel_seq_render = None
            else:
                pass
        except Exception as e:
            logger.error(f"Failed to prepare GT accelerations for visualization: {e}")
            gt_accel_seq_render = None

        # Masks for valid vertices per frame/object from the sample batch
        try:
            vmask_all = batch['valid_vtx_mask'][0].cpu().numpy()  # (T0, Nobj, Nv)
        except Exception:
            vmask_all = None

        # Load faces for mesh rendering from batch (required; no fallback).
        # This ensures faces match the permuted object order in training data.
        if 'faces_list' not in batch or batch['faces_list'] is None:
            raise RuntimeError(
                f"Missing 'faces_list' in batch! DataLoader must provide faces. "
                f"Check that faces_list is included in sample dict."
            )
        # Use permuted faces from batch (handles object permutation augmentation)
        faces_list = batch['faces_list'][0]  # First sample in batch
        if len(faces_list) != Nobj:
            raise RuntimeError(f"Faces count mismatch: batch has {len(faces_list)} faces but {Nobj} objects")
        logger.info(f"[Faces] Using permuted faces from batch: {len(faces_list)} objects")

        frames_dir = self.log_dir / f'temp_frames_epoch_{self.epoch}_mobj_pred'
        if video_renderer == 'polyscope':
            frames_dir.mkdir(exist_ok=True)
            try:
                if hasattr(ps, 'set_allow_headless'):
                    ps.set_allow_headless(True)
                if hasattr(ps, 'set_backend'):
                    ps.set_backend('openGL3_egl')
                ps.init('openGL3_egl')
            except Exception as e:
                raise RuntimeError(f"Failed to initialize Polyscope headless backend: {e}")
            ps.set_autocenter_structures(False)
            ps.set_ground_plane_mode("none")
            ps.set_ground_plane_height(0.0)
            ps.set_navigation_style('turntable')
            ps.set_up_dir('z_up')

        # Colors per object
        rng = np.random.default_rng(123)
        colors = rng.random((Nobj, 3)) * 0.6 + 0.2

        # Visualization parameters - READ FROM CONFIG
        checkpoint_config = self.config['training']['checkpoint']

        # Camera settings from config
        camera_config = checkpoint_config.get('camera', {})
        camera_mode = camera_config.get('mode', 'auto')

        first = batch['obj_positions'][0, 0].cpu().numpy().reshape(-1, 3)

        if camera_mode == 'fixed':
            # Use fixed camera from config
            center = np.array(camera_config.get('center', [0.0, 0.0, 0.5]), dtype=np.float64)
            eye = np.array(camera_config.get('eye', [3.0, -3.0, 3.0]), dtype=np.float64)
        else:
            # Auto camera: compute from scene bounds
            center = first.mean(axis=0)
            bbox_min = first.min(axis=0)
            bbox_max = first.max(axis=0)
            radius = 0.5 * np.linalg.norm(bbox_max - bbox_min)
            auto_direction = np.array(camera_config.get('auto_direction', [2.0, -2.0, 2.0]), dtype=np.float64)
            auto_scale = camera_config.get('auto_scale', 3.0)
            eye = center + auto_direction * max(radius, 1e-3) * auto_scale

        # Ground plane settings from config
        ground_config = checkpoint_config.get('ground', {})
        ground_size = ground_config.get('size', 10.0)
        ground_color = ground_config.get('color', [0.8, 0.8, 0.8])
        ground_vertices = np.array([
            [-ground_size, -ground_size, 0.0],
            [ ground_size, -ground_size, 0.0],
            [ ground_size,  ground_size, 0.0],
            [-ground_size,  ground_size, 0.0]
        ])
        ground_faces = np.array([[0, 1, 2], [0, 2, 3]])

        total_frames = actual_render_frames

        # GT overlay settings (from config, can be overridden by override_if_overlay parameter)
        if override_if_overlay is not None:
            if_overlay = int(override_if_overlay)  # Use override value
        else:
            if_overlay = 1 if checkpoint_config.get('if_overlay', True) else 0  # Config -> bool, convert to int
        gt_transparency = checkpoint_config.get('gt_transparency', 0.1)  # Default: very transparent (0.1 = 90% transparent)

        # Arrow visualization settings (from config)
        arrow_enable = 1 if checkpoint_config.get('arrow_enable', False) else 0  # Default: disabled
        arrow_scale = checkpoint_config.get('arrow_scale', 1.0)
        arrow_subsample = checkpoint_config.get('arrow_subsample', 1)
        arrow_thickness = checkpoint_config.get('arrow_thickness', 0.0001)
        point_size = checkpoint_config.get('point_size', 0.0001)

        logger.info(f"[Visualization Config] if_overlay={if_overlay}, gt_transparency={gt_transparency}, "
                    f"arrow_enable={arrow_enable}, arrow_scale={arrow_scale}, arrow_thickness={arrow_thickness}, "
                    f"arrow_subsample={arrow_subsample}, point_size={point_size}")
        
        # Prepare GT positions if overlay is enabled
        # Need to load full GT sequence from disk because batch only contains sequence_length frames
        # (e.g., batch has 4 frames but rollout predicts 48 frames)
        gt_full = None
        if if_overlay == 1:
            logger.info(f"[GT Overlay] if_overlay=1, starting GT loading...")
            try:
                # Check batch GT length vs rollout length
                batch_gt = batch['obj_positions'][0].cpu().numpy()  # (T_batch, Nobj, Nv, 3)
                T_batch = batch_gt.shape[0]
                T_rollout = pred_full.shape[0]

                logger.info(f"GT overlay: Batch has {T_batch} GT frames, Rollout has {T_rollout} frames")

                # If rollout is longer than batch (common case), load full GT from disk
                if T_rollout > T_batch:
                    logger.info(f"Loading full GT sequence from disk to match rollout length...")

                    # Find all frame directories
                    frame_dirs = sorted([d for d in ds_dir.glob('frame_*') if d.is_dir()])
                    if len(frame_dirs) == 0:
                        raise RuntimeError(f"No frame_* directories found in {ds_dir}")

                    # Same frames as the rollout: from the sequence's first physical frame at skip_step_size
                    _frame_numbers = [int(d.name.split('_')[-1]) for d in frame_dirs]
                    if video_start_frame not in _frame_numbers:
                        raise RuntimeError(f"frame_{video_start_frame:05d} not found in {ds_dir}")
                    frame_dirs_sampled = frame_dirs[_frame_numbers.index(video_start_frame)::skip_step_size]
                    T_gt = min(len(frame_dirs_sampled), T_rollout)  # Don't exceed rollout length

                    # Load first frame to get dimensions
                    # Use exact file pattern obj_{i}_pos.npy to avoid matching point cloud versions
                    # (e.g., obj_0_pc1024_pos.npy, obj_0_pc512_pos.npy, etc.)
                    first_frame = frame_dirs_sampled[0]

                    # Count objects by checking which obj_{i}_pos.npy files exist
                    obj_files_first = []
                    obj_idx = 0
                    while True:
                        obj_file = first_frame / f"obj_{obj_idx}_pos.npy"
                        if obj_file.exists():
                            obj_files_first.append(obj_file)
                            obj_idx += 1
                        else:
                            break

                    Nobj_gt = len(obj_files_first)
                    if Nobj_gt == 0:
                        raise RuntimeError(f"No obj_{{i}}_pos.npy files found in {first_frame}")

                    # Get max vertex count across all objects (different objects may have different Nv)
                    Nv_per_obj = []
                    for obj_file in obj_files_first:
                        obj_pos = np.load(obj_file)
                        Nv_per_obj.append(obj_pos.shape[0])
                    Nv_max = max(Nv_per_obj)

                    logger.info(f"GT loading: Nobj={Nobj_gt}, Nv per object={Nv_per_obj}, Nv_max={Nv_max}")

                    # Allocate array with max vertex count (pad smaller objects with zeros)
                    gt_full = np.zeros((T_gt, Nobj_gt, Nv_max, 3), dtype=np.float32)

                    # Load all sampled frames using exact file pattern obj_{i}_pos.npy
                    for t_idx in range(T_gt):
                        frame_dir = frame_dirs_sampled[t_idx]
                        for obj_idx in range(Nobj_gt):
                            obj_file = frame_dir / f"obj_{obj_idx}_pos.npy"
                            if obj_file.exists():
                                obj_pos = np.load(obj_file).astype(np.float32)
                                if self.config['data'].get('object_scale_file'):
                                    if not hasattr(self, '_video_object_scale'):
                                        from src.data.object_scale import from_config as _osf
                                        self._video_object_scale = _osf(self.config)
                                    obj_pos = self._video_object_scale.apply(obj_file, obj_pos)
                                Nv_curr = obj_pos.shape[0]
                                gt_full[t_idx, obj_idx, :Nv_curr, :] = obj_pos

                    logger.info(f"Loaded full GT from disk: shape {gt_full.shape}")
                    gt_from_disk = True

                    # Apply permutation to GT to match batch order (if augmentation is enabled)
                    # This ensures GT meshes match the permuted faces_list from batch
                    if 'permutation' in batch and batch['permutation'] is not None:
                        perm = batch['permutation'][0]  # Get permutation from batch (batch size 1)
                        if isinstance(perm, torch.Tensor):
                            perm = perm.cpu().numpy()
                        logger.info(f"[GT Permutation] Applying object permutation to GT: {perm.tolist()}")
                        gt_full = gt_full[:, perm]  # Permute object dimension: (T, Nobj, Nv, 3)[:, perm]
                        logger.info(f"[GT Permutation] GT permuted successfully")
                else:
                    # Batch GT is sufficient (training mode with short rollout); it already carries the batch's
                    # rotation augmentation
                    gt_full = batch_gt
                    gt_from_disk = False
                    logger.info(f"Using batch GT: shape {gt_full.shape}")
                    # Batch GT is already permuted by DataLoader, no need to permute again

                # Assert GT length matches rollout length for proper overlay
                assert gt_full.shape[0] == T_rollout, (
                    f"GT length mismatch! GT has {gt_full.shape[0]} frames, "
                    f"but rollout has {T_rollout} frames (num_steps={num_frames_rollout}). "
                    f"GT and rollout must have same length for proper overlay visualization."
                )
                logger.info(f"GT-Rollout length check passed: {gt_full.shape[0]} frames")
                logger.info(f"[GT Overlay] GT loaded successfully: gt_full.shape={gt_full.shape}")

                # Apply same rotation to GT as was applied to batch (for augmentation verification); ground truth
                # taken from the batch is already rotated
                if rotation_angle_deg != 0.0 and gt_from_disk:
                    import math
                    logger.info(f"[GT Overlay] Applying {rotation_angle_deg:.0f}° rotation to GT for consistency")
                    angle_rad = math.radians(rotation_angle_deg)
                    cos_a = math.cos(angle_rad)
                    sin_a = math.sin(angle_rad)
                    # Z-axis rotation matrix (numpy)
                    R_np = np.array([
                        [cos_a, -sin_a, 0],
                        [sin_a,  cos_a, 0],
                        [0,      0,     1]
                    ], dtype=np.float32)

                    # Rotate all GT positions: gt_full shape (T, Nobj, Nv, 3)
                    T_gt, Nobj_gt, Nv_gt, _ = gt_full.shape
                    gt_flat = gt_full.reshape(-1, 3)  # (T*Nobj*Nv, 3)
                    gt_rotated = np.matmul(gt_flat, R_np.T)  # (T*Nobj*Nv, 3)
                    gt_full = gt_rotated.reshape(T_gt, Nobj_gt, Nv_gt, 3)
                    logger.info(f"[GT Overlay] GT rotation applied successfully")

                # Apply same render subsampling to GT for video efficiency
                gt_full_render = gt_full[render_frame_indices]  # (num_frames_render, Nobj, Nv, 3)
                logger.info(f"[GT Overlay] Applied render subsampling: {gt_full.shape} -> {gt_full_render.shape}")
                gt_full = gt_full_render  # Use subsampled GT for rendering

            except Exception as e:
                if (self.config.get('inference', {}) or {}).get('mode', False):
                    raise RuntimeError(f"[GT Overlay] requested overlay, ground truth could not be loaded: {e}") from e
                logger.warning(f"[GT Overlay] Failed to load GT positions for overlay: {e}")
                import traceback
                traceback.print_exc()
                gt_full = None
                if_overlay = 0  # Disable overlay if GT is unavailable (training-time diagnostics only)
                logger.warning(f"[GT Overlay] GT loading failed, disabling overlay. if_overlay set to 0")

        if video_renderer == 'filament':
            # real time: key frame k is k * skip_step_size physical frames after the first
            times = render_frame_indices * s / fps_orig
            self._render_video_filament(ds_dir, batch, pred_full_render, gt_full if if_overlay == 1 else None,
                                        vmask_all, output_path, fps_out, times)
            return inference_fps

        for t in range(total_frames):
            # Register ground
            ps.register_surface_mesh("ground", ground_vertices, ground_faces, enabled=True)
            ps.get_surface_mesh("ground").set_color(ground_color)
            for i in range(Nobj):
                # Render predicted mesh (opaque) - using subsampled predictions
                V_pred = pred_full_render[t, i]
                F = faces_list[i]
                ps.register_surface_mesh(f"obj_{i}_pred", V_pred, F, enabled=True, smooth_shade=True)
                ps.get_surface_mesh(f"obj_{i}_pred").set_color(colors[i])
                ps.get_surface_mesh(f"obj_{i}_pred").set_transparency(1.0)  # Fully opaque

                # Render GT mesh (transparent) if overlay is enabled
                if if_overlay == 1 and gt_full is not None and t < gt_full.shape[0] and i < gt_full.shape[1]:
                    V_gt_full = gt_full[t, i]  # (Nv_max, 3) - may contain zero-padding

                    # Find actual number of non-zero vertices for this object
                    # (GT array is padded with zeros for objects with fewer vertices)
                    # Check if vertex has any non-zero coordinates
                    vertex_norms = np.linalg.norm(V_gt_full, axis=1)
                    # Count consecutive non-zero vertices from the start
                    # (assumes padding is at the end)
                    Nv_actual = V_pred.shape[0]  # Use predicted vertex count as reference

                    # Trim GT vertices to match predicted vertex count
                    V_gt = V_gt_full[:Nv_actual]

                    # Use faces as-is (they should match the actual vertex count)
                    F_gt = F

                    # Filter faces to only use valid vertex indices
                    F_gt = F_gt[np.all(F_gt < Nv_actual, axis=1)]

                    # GT mesh with same faces, rendered highly transparent so it does not obscure the prediction
                    if F_gt.shape[0] > 0:  # Only render if we have valid faces
                        ps.register_surface_mesh(f"obj_{i}_gt", V_gt, F_gt, enabled=True, smooth_shade=True)
                        ps.get_surface_mesh(f"obj_{i}_gt").set_color([0.9, 0.9, 0.9])  # Light gray
                        ps.get_surface_mesh(f"obj_{i}_gt").set_transparency(gt_transparency)  # From config (default: 0.1 = 90% transparent)

                # For backwards compatibility: also register without _pred suffix
                # (arrows reference "obj_{i}" positions)
                V = V_pred  # Use predicted positions for arrow attachment

                # Visualize predicted (red) and GT (green) acceleration arrows
                if arrow_enable == 1:
                    try:
                        # Optional subsampling for clarity/performance
                        arrow_k = max(1, arrow_subsample)

                        # Anchor indices for this frame/object
                        anchor_idx_np = batch['anchor_idx'][0, i].cpu().numpy()  # (Nq,)
                        # Build a unified index into the anchor set to apply mask and subsample consistently
                        sel_idx = np.arange(anchor_idx_np.shape[0])  # base 0..Nq-1
                        if vmask_all is not None:
                            mask_idx = min(t, vmask_all.shape[0] - 1)
                            m_full = vmask_all[mask_idx, i].astype(bool)  # (Nv,)
                            anchor_mask = m_full[anchor_idx_np]           # (Nq,)
                            sel_idx = sel_idx[anchor_mask]
                        if arrow_k > 1:
                            sel_idx = sel_idx[::arrow_k]
                        if sel_idx.size == 0:
                            continue
                        # Sampled anchor positions (for vectors attachment)
                        V_sampled = V[anchor_idx_np[sel_idx]]

                        # Predicted acceleration on anchors (RED arrows) - using subsampled accelerations
                        if pred_accel_anchor_seq_render is not None and t >= 1 and (t-1) < pred_accel_anchor_seq_render.shape[0]:
                            # render_t=0 has no accel (frame 0), render_t=1..N correspond to accel_idx=0..N-1
                            pred_accel_anchor_t = pred_accel_anchor_seq_render[t-1, i]  # (Nq, 3)
                            pred_accel_sampled = pred_accel_anchor_t[sel_idx]

                            # Register point cloud with red arrows (points invisible, only arrows visible)
                            pc_pred = ps.register_point_cloud(f"obj_{i}_pred_accel", V_sampled,
                                                              radius=0.0,
                                                              color=(1.0, 0.0, 0.0))
                            pc_pred.add_vector_quantity("pred_accel", pred_accel_sampled,
                                                        enabled=True,
                                                        length=arrow_scale,
                                                        radius=arrow_thickness,
                                                        color=(1.0, 0.0, 0.0))
                        else:
                            pass  # Frame 0 or accel index invalid

                        # GT acceleration (green arrows) at the same anchor vertices - using subsampled accelerations
                        if gt_accel_seq_render is not None and t >= 1 and (t-1) < gt_accel_seq_render.shape[0]:
                            gt_accel_t = gt_accel_seq_render[t-1, i]  # (Nv, 3)
                            gt_accel_sampled = gt_accel_t[anchor_idx_np[sel_idx]]

                            # Register point cloud with green arrows (points invisible, only arrows visible)
                            pc_gt = ps.register_point_cloud(f"obj_{i}_gt_accel", V_sampled,
                                                            radius=0.0,
                                                            color=(0.0, 1.0, 0.0))
                            pc_gt.add_vector_quantity("gt_accel", gt_accel_sampled,
                                                      enabled=True,
                                                      length=arrow_scale,
                                                      radius=arrow_thickness,
                                                      color=(0.0, 1.0, 0.0))
                        else:
                            pass

                    except Exception as e:
                        logger.debug(f"Failed to visualize acceleration arrows: {e}")
                        pass
            ps.look_at(eye, center)  # ps.look_at(camera_location, target)
            ps.screenshot(str(frames_dir / f'frame_{t:05d}.png'))
            ps.remove_all_structures()

        # Console log synthesis parameters
        try:
            expected_duration = float(total_frames) / max(1e-6, fps_out)
        except Exception:
            expected_duration = -1.0
        logger.info(
            f"[video] dataset={dataset_name} split={split} fps_orig={fps_orig} skip={skip_step_size} "
            f"fps_out={fps_out:.6f} base_num_frames={base_num_frames} "
            f"rollout_frames={num_frames_rollout} rendered_frames={total_frames} expected_duration_s={expected_duration:.3f}"
        )

        # Encode video
        ffmpeg_cmd = [
            'ffmpeg', '-hide_banner', '-loglevel', 'error',
            '-r', f"{fps_out:.6f}",
            '-i', str(frames_dir / 'frame_%05d.png'),
            '-vf', 'scale=trunc(iw/2)*2:trunc(ih/2)*2',
            '-c:v', 'mpeg4', '-q:v', '2', '-pix_fmt', 'yuv420p', '-y',
            str(output_path)
        ]
        subprocess.run(ffmpeg_cmd, check=True)

        # Clean up temporary frames
        import shutil
        shutil.rmtree(frames_dir, ignore_errors=True)
        logger.info(f"  [multiobj_pred] Cleaned up temp frames: {frames_dir}")

        # Return inference FPS for video file naming
        return inference_fps

    def _visualize_predictions_vs_gt(self, predictions, batch, batch_idx=0):
        """Visualize predictions vs ground truth (multi-object)."""
        import matplotlib.pyplot as plt
        import numpy as np
        from mpl_toolkits.mplot3d import Axes3D
        
        vis_dir = self.log_dir / 'per_frame_prediction'
        vis_dir.mkdir(parents=True, exist_ok=True)
        
        if 'obj_positions_denorm' not in predictions or 'obj_positions' not in batch:
            logger.warning("Multi-object tensors missing for visualization; skipping.")
            return
        
        pred = predictions['obj_positions_denorm'][0].detach().cpu().numpy()    # (T-2,Nobj,Nv,3)
        gt_all = batch['obj_positions'][0].detach().cpu().numpy()               # (T,Nobj,Nv,3)
        vmask_all = batch['valid_vtx_mask'][0].detach().cpu().numpy()           # (T,Nobj,Nv)
        obj_valid = batch['valid_obj_mask'][0].detach().cpu().numpy()           # (Nobj,)
        Tm2, Nobj, Nv, _ = pred.shape
        num_pred_frames = Tm2
        
        rng = np.random.default_rng(123)
        colors_pred = rng.random((Nobj, 3)) * 0.6 + 0.2
        colors_gt = np.clip(colors_pred + 0.4, 0.0, 1.0)
        
        fig = plt.figure(figsize=(20, 5 * ((num_pred_frames + 1) // 4 + 1)))
        for t in range(num_pred_frames):
            ax = fig.add_subplot((num_pred_frames + 3) // 4, 4, t + 1, projection='3d')
            frame_pred = pred[t]                # (Nobj,Nv,3)
            frame_gt = gt_all[t + 2]           # (Nobj,Nv,3)
            frame_mask = vmask_all[t + 2]      # (Nobj,Nv)
            any_points = False
            for j in range(Nobj):
                if not obj_valid[j]:
                    continue
                m = frame_mask[j]
                if not np.any(m):
                    continue
                any_points = True
                P = frame_pred[j][m]
                G = frame_gt[j][m]
                ax.scatter(P[:, 0], P[:, 2], P[:, 1], c=[colors_pred[j]], marker='o', s=18, alpha=0.7)
                ax.scatter(G[:, 0], G[:, 2], G[:, 1], c=[colors_gt[j]], marker='^', s=14, alpha=0.8)
            ax.set_xlabel('X')
            ax.set_ylabel('Z')
            ax.set_zlabel('Y')
            ax.set_title(f'Frame {t+2}')
            if not any_points:
                ax.text2D(0.1, 0.5, "no valid points", transform=ax.transAxes)
        
        plot_path = vis_dir / f'pred_vs_gt_multi_epoch_{self.epoch:04d}_batch_{batch_idx:03d}.png'
        plt.savefig(plot_path, dpi=150, bbox_inches='tight')
        plt.close()
    
    
    
    @staticmethod
    def _atomic_save(obj, path):
        """torch.save to <path>.part, then rename: a reader (the EMA watcher, a resume) never sees a partly written file,
        and an interrupted write leaves the previous file intact."""
        path = Path(path)
        tmp = path.with_name(path.name + '.part')
        torch.save(obj, tmp)
        os.replace(tmp, path)

    def save_checkpoint(self, is_best=False):
        """Save model checkpoint (rank0 only in DDP mode)."""
        # Only rank 0 saves checkpoints to prevent file conflicts in DDP
        if not self.is_main:
            return

        checkpoint = {
            'epoch': self.epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict() if self.optimizer is not None else None,
            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler is not None else None,
            'best_val_loss': float(self.best_val_loss),
            'config': self.config,
            # The scheduler state above was captured after the end-of-epoch step (see train()).
            'scheduler_stepped': True,
        }
        if self.ema is not None:
            checkpoint['ema_state_dict'] = self.ema.state_dict()
            checkpoint['ema_num_updates'] = self.ema.num_updates
            # Weights-only file of the averaged model, loadable by the standard evaluation driver.
            _ema_every = int(self.config['training'].get('ema', {}).get('save_every_n_epochs', 1))
            if self.epoch % _ema_every == 0 or self.epoch == int(self.config['training']['max_epochs']) - 1:
                self._atomic_save({'epoch': self.epoch, 'model_state_dict': checkpoint['ema_state_dict'],
                                   'config': self.config, 'ema_decay': self.ema.decay,
                                   'ema_num_updates': self.ema.num_updates},
                                  self.checkpoint_dir / f'epoch_{self.epoch:04d}_ema.pth')

        # Save the best checkpoint before the latest one: the best-model threshold that latest.pth carries (and a resume
        # continues from) always belongs to a best.pth that was written
        if is_best:
            best_path = self.checkpoint_dir / 'best.pth'
            self._atomic_save(checkpoint, best_path)
            logger.info(f"Saved best model with val_loss: {self.best_val_loss:.4f}")

        # Save latest checkpoint
        latest_path = self.checkpoint_dir / 'latest.pth'
        self._atomic_save(checkpoint, latest_path)

        # Generate visualization video if enabled (do this first, independent of checkpoint saving)
        self.generate_visualization_video()

        # Save periodic checkpoint
        save_every = self.config['training']['checkpoint'].get('save_every_n_epochs', 10)
        if self.epoch % save_every == 0:
            epoch_path = self.checkpoint_dir / f'epoch_{self.epoch:04d}.pth'
            self._atomic_save(checkpoint, epoch_path)
    
    def run_inference(self):
        """
        Inference-only mode: Generate videos and compute RMSE metrics.
        Called when --inference flag is passed or inference.mode=true in config.
        """
        logger.info("=" * 80)
        logger.info("INFERENCE MODE STARTED")
        logger.info("=" * 80)

        self.model.eval()
        self.epoch = 0  # Set epoch for logging purposes

        # Get inference config
        inference_config = self.config.get('inference', {})
        generate_videos = inference_config.get('generate_videos', True)
        compute_metrics = inference_config.get('compute_metrics', True)

        # Step 1: Compute RMSE metrics
        if compute_metrics:
            logger.info("-" * 40)
            logger.info("Step 1: Computing RMSE Metrics")
            logger.info("-" * 40)
            try:
                test_metrics = self.evaluator.evaluate_rollout_rmse(epoch=0)
                if test_metrics:
                    logger.info("=" * 60)
                    logger.info("INFERENCE RESULTS - RMSE Metrics:")
                    logger.info("=" * 60)
                    for key, value in sorted(test_metrics.items()):
                        if isinstance(value, float):
                            logger.info(f"  {key}: {value:.6f}")
                        else:
                            logger.info(f"  {key}: {value}")
                    logger.info("=" * 60)
                else:
                    logger.warning("No metrics returned from evaluation")
            except Exception as e:
                # an evaluation run must fail when its metrics fail (a partial eval.txt must not pass as a result)
                logger.error(f"RMSE evaluation failed: {e}")
                raise
        else:
            logger.info("Metrics computation skipped (inference.compute_metrics=false)")

        # Step 2: Generate visualization videos (after the metrics, which they cannot affect)
        if generate_videos:
            logger.info("-" * 40)
            logger.info("Step 2: Generating Visualization Videos")
            logger.info("-" * 40)
            # a failed requested video fails the run (generate_visualization_video raises in inference mode)
            self.generate_visualization_video()
            logger.info("Video generation completed")
        else:
            logger.info("Video generation skipped (inference.generate_videos=false)")

        # Step 3: Optional validation pass (if enabled)
        if self.enable_validation and self.val_loader is not None:
            logger.info("-" * 40)
            logger.info("Step 3: Running Validation Pass")
            logger.info("-" * 40)
            try:
                val_loss = self.validate()
                logger.info(f"Validation Loss: {val_loss:.6f}")
            except Exception as e:
                logger.error(f"Validation failed: {e}")
                import traceback
                traceback.print_exc()

        logger.info("=" * 80)
        logger.info("INFERENCE MODE COMPLETED")
        logger.info("=" * 80)

        # Close tensorboard writer
        if hasattr(self, 'writer') and self.writer is not None:
            self.writer.close()

    def train(self):
        """Main training loop (or inference-only for empty train set)."""
        if self.inference_only:
            logger.info("Inference-only mode: running single validation pass and generating videos")
            self.epoch = 0
            if self.enable_validation:
                val_loss = self.validate()
                logger.info(f"Validation Loss: {val_loss:.4f}")
            else:
                logger.info("Validation skipped (disabled by config)")

            # Try to generate videos, but gracefully handle GPU memory errors in CPU mode
            try:
                self.generate_visualization_video()
            except (RuntimeError, Exception) as e:
                if 'out of memory' in str(e).lower() or 'opengl' in str(e).lower():
                    logger.error(f"Video generation failed (GPU memory/OpenGL error): {e}")
                    logger.warning("Skipping video generation - GPU memory insufficient for Polyscope rendering")
                    logger.warning("To generate videos, free GPU memory or run on a machine with an available GPU")
                else:
                    raise

            logger.info("Inference-only mode completed")
            return

        logger.info(f"Starting training for {self.config['training']['max_epochs']} epochs")

        # Initial evaluation BEFORE training starts (epoch=-1 to distinguish from post-epoch evals)
        # Controlled by training.metrics_test.initial_eval config flag
        initial_eval_enabled = self.config['training'].get('metrics_test', {}).get('initial_eval', False)
        if self.evaluator.metrics_enabled and initial_eval_enabled:
            logger.info("=" * 80)
            logger.info("[Initial Evaluation] Running RMSE evaluation BEFORE training starts")
            logger.info("=" * 80)
            # Use epoch=-1 to indicate pre-training evaluation; a failure stops the run (as the per-epoch one does)
            initial_metrics = self.evaluator.evaluate_rollout_rmse(epoch=-1)
            if initial_metrics:
                logger.info("[Initial Evaluation] Completed successfully")
        elif self.evaluator.metrics_enabled:
            logger.info("[Initial Evaluation] Skipped (initial_eval=false in config)")

        max_epochs = self.config['training']['max_epochs']
        if self.start_epoch > 0:
            if self.start_epoch >= max_epochs:
                raise RuntimeError(
                    f"Resumed checkpoint already completed epoch {self.start_epoch - 1} but "
                    f"training.max_epochs={max_epochs}: raise max_epochs to continue, or fine-tune "
                    f"from a weights-only checkpoint to start a new schedule")
            logger.info(f"[Resume] Continuing training at epoch {self.start_epoch} of {max_epochs}")
        for epoch in range(self.start_epoch, max_epochs):
            self.epoch = epoch

            # Update DistributedSampler epoch for proper shuffling in multi-GPU training;
            # otherwise every epoch uses the same random seed and the same data order.
            try:
                if hasattr(self.train_loader, 'sampler') and hasattr(self.train_loader.sampler, 'set_epoch'):
                    self.train_loader.sampler.set_epoch(epoch)
                if hasattr(self.val_loader, 'sampler') and hasattr(self.val_loader.sampler, 'set_epoch'):
                    self.val_loader.sampler.set_epoch(epoch)
            except Exception:
                pass

            # Expose current epoch to model for warmup/ramp schedules (handle DDP)
            try:
                if isinstance(self.model, DDP):
                    setattr(self.model.module, 'current_epoch', epoch)
                else:
                    setattr(self.model, 'current_epoch', epoch)
            except Exception:
                pass

            # Update dataset epoch for noise scheduling
            if hasattr(self.train_loader.dataset, 'set_epoch'):
                self.train_loader.dataset.set_epoch(epoch)
                # Log noise status
                noise_config = self.config['data'].get('augmentation', {}).get('random_walk', {})
                if noise_config.get('enable', False):
                    start_epoch = noise_config.get('start_epoch', 0)
                    noise_std = noise_config.get('noise_std', 0)
                    if epoch >= start_epoch:
                        logger.info(f"Epoch {epoch}: Random walk noise ACTIVE (std={noise_std})")
                    else:
                        logger.info(f"Epoch {epoch}: Random walk noise INACTIVE (starts at epoch {start_epoch})")

            # Training
            train_stats = self.train_epoch()
            logger.info(
                "Epoch %d: Train Loss = %.4g | pos_before=%.4g, pos_after=%.4g, accel_before=%.4g, accel_after=%.4g, "
                "edge=%.4g, pen_loss=%.4g, pen_acc=%.4g",
                epoch,
                train_stats['loss'],
                train_stats['pos_before_icp'],
                train_stats['pos_after_icp'],
                train_stats['accel_before_icp'],
                train_stats['accel_after_icp'],
                train_stats['edge'],
                train_stats['pen_sign_loss'],
                train_stats['pen_sign_acc'],
            )

            # Synchronize all ranks before validation in DDP mode
            # Ensures all GPUs finish the training epoch before starting validation
            if dist.is_available() and dist.is_initialized():
                dist.barrier()

            # Validation
            val_loss = None
            if self.enable_validation:
                val_loss = self.validate()
                # visualization is inside the self.validate() function
                logger.info(f"Epoch {epoch}: Val Loss = {val_loss:.4f}")
            else:
                logger.info(f"Epoch {epoch}: Validation skipped (disabled by config)")

            # RMSE evaluation on test/val sets (HopNet-style checkpoints). A failure stops the run (with the
            # traceback): training on without evaluations would also skip best-checkpoint selection, and one failing
            # rank would leave the others waiting in the evaluator's collectives.
            test_metrics = {}
            evaluated = self.evaluator.should_evaluate(epoch)
            if evaluated:
                test_metrics = self.evaluator.evaluate_rollout_rmse(epoch) or {}

            # One-time consistency check after epoch 1 (only on main process)
            if epoch == 1 and self.is_main:
                if not hasattr(self, '_consistency_check_done'):
                    try:
                        self.check_forward_rollout_consistency()
                        self._consistency_check_done = True
                    except Exception as e:
                        logger.error(f"Consistency check failed: {e}")
                        raise

            # Check if best model
            # When validation is disabled, use test metrics (position_rmse_step050) instead
            if self.enable_validation and val_loss is not None:
                # Use validation loss as best criterion
                is_best = val_loss < self.best_val_loss
                if is_best:
                    self.best_val_loss = float(val_loss)
                    logger.info(f"New best model: Val Loss = {val_loss:.4f}")
            else:
                # Use a rollout position-RMSE metric as best criterion (lower is better). The
                # evaluator reports keys of the form 'p_rmse_@<frame>/s=<step size>'.
                # training.checkpoint.best_metric: one key, or a list of keys whose mean is the criterion (e.g. the
                # step-size-1 long-horizon checkpoints @200/@300/@400)
                best_key = self.config['training'].get('checkpoint', {}).get('best_metric', 'p_rmse_@050/s=1')
                _best_keys = list(best_key) if isinstance(best_key, (list, tuple)) else [best_key]
                if evaluated and self.is_main and self.evaluator.metrics_enabled and not (
                        all(k in test_metrics for k in _best_keys) and all(np.isfinite(float(test_metrics[k])) for k in _best_keys)):
                    # the evaluator returns the metrics on the main rank; missing or non-finite selection metrics
                    # would silently stop best-checkpoint selection
                    raise RuntimeError(f"epoch {epoch}: checkpoint-selection metrics {_best_keys} missing or non-finite "
                                       f"in the evaluation ({ {k: test_metrics.get(k) for k in _best_keys} })")
                if test_metrics and all(k in test_metrics for k in _best_keys):
                    # float(): the evaluator's aggregates are numpy scalars, and a numpy scalar
                    # saved in the checkpoint breaks torch.load's default weights_only=True path.
                    current_rmse = float(sum(float(test_metrics[k]) for k in _best_keys) / len(_best_keys))
                    is_best = current_rmse < self.best_val_loss
                    if is_best:
                        self.best_val_loss = current_rmse
                        logger.info(f"New best model: {best_key} = {current_rmse:.6f} m")
                else:
                    # No metrics available, skip best check
                    is_best = False
                    logger.info("No test metrics available, skipping best checkpoint check")

            # Advance the LR schedule (cosine/constant schedulers) BEFORE saving, so that the
            # scheduler state stored in the checkpoint is the one a continuous run uses for the
            # next epoch; a resumed run then reproduces the same LR sequence.
            self.scheduler.step()

            # Save checkpoint and save the visualization video
            self.save_checkpoint(is_best=is_best)

            # Log per-parameter-group learning rates
            for i, pg in enumerate(self.optimizer.param_groups):
                group_name = 'base' if i == 0 else f'group{i}'
                actual_lr = pg.get('lr', 0.0)
                tag = f'train/lr_{group_name}'
                self.writer.add_scalar(tag, actual_lr, epoch)
                if epoch % 10 == 0:
                    logger.info(f"[LR Tracking] {group_name}: {actual_lr:.6e}")
            
            # Log teacher forcing ratio if using scheduled sampling
            if self.use_scheduled_sampling:
                teacher_ratio = self.get_teacher_forcing_ratio(epoch)
                self.writer.add_scalar('train/teacher_forcing_ratio', teacher_ratio, epoch)
                logger.info(f"Epoch {epoch}: Teacher Forcing Ratio = {teacher_ratio:.3f}")
        
        logger.info("Training completed!")
        self.writer.close()


def _adapt_state_dict_for_height_input(model, state_dict):
    """Zero-pad the encoder's first convolution for the optional height input channel so a
    checkpoint trained without it warm-starts with an identical function."""
    msd = model.state_dict()
    for key, w in list(state_dict.items()):
        if key.endswith('shared_obj_encoder.conv1.weight') and key in msd and w.dim() == 3:
            want = msd[key].shape[1]
            if w.shape[1] == want - 1:
                pad = torch.zeros(w.shape[0], 1, w.shape[2], dtype=w.dtype, device=w.device)
                state_dict[key] = torch.cat([w, pad], dim=1)
                print(f"[Checkpoint] height_input: {key} zero-padded {w.shape[1]}->{want} input channels (exact warm start)", flush=True)
    return state_dict


def main():
    parser = argparse.ArgumentParser(description='Train the RigidFormer multi-body dynamics model')
    parser.add_argument('--config', type=str, default='configs/movi_b.yaml',
                       help='Path to config file')
    parser.add_argument('--resume', type=str, default=None,
                       help='Path to checkpoint to resume from')
    parser.add_argument('--checkpoint', type=str, default=None,
                       help='Path to checkpoint to resume from (alias for --resume)')
    parser.add_argument('--device', type=str, default='cuda',
                       help='Device to use (cuda or cpu)')
    parser.add_argument('--use_multi_dataset', action='store_true',
                       help='Use multi-dataset loader for training on multiple datasets')
    parser.add_argument('--seed', type=int, default=42,
                       help='Random seed for reproducibility')
    parser.add_argument('--inference', action='store_true',
                       help='Run inference only (no training). Generates videos and computes metrics.')
    parser.add_argument('--compile', action='store_true',
                       help='Compile the encoder, decoder and predictor with torch.compile for faster '
                            'inference. One-time compilation cost on startup; results may differ from '
                            'the eager path at floating-point rounding level.')
    parser.add_argument('--cuda-graph', action='store_true',
                       help='Capture the rollout step as a CUDA graph and replay it, roughly halving '
                            'per-step inference time. Trajectories are identical to the standard path.')
    args = parser.parse_args()
    if args.cuda_graph:
        os.environ['RIGIDFORMER_CUDA_GRAPH'] = '1'

    # Handle both --resume and --checkpoint arguments
    checkpoint_path = args.resume or args.checkpoint

    # Load config
    with open(args.config, 'r') as f:
        config = yaml.safe_load(f)

    # Create dataloaders with split ratios from config
    train_ratio = config['data'].get('train_split', 0.8)
    val_ratio = config['data'].get('val_split', 0.1)

    # Check if validation is enabled
    enable_validation = config['training'].get('enable_validation', True)
    if not enable_validation:
        logger.info("Validation disabled by config - val_loader will not be created")
        val_ratio = 0.0  # Allocate all non-train data to test set

    # DDP setup (optional)
    ddp_world_size = int(os.environ.get('WORLD_SIZE', '1'))
    ddp_rank = int(os.environ.get('RANK', '0'))
    ddp_local_rank = int(os.environ.get('LOCAL_RANK', '0'))
    use_ddp = ddp_world_size > 1
    if use_ddp:
        if args.device == 'cuda':
            torch.cuda.set_device(ddp_local_rank)
        # Set 8-hour timeout for NCCL operations (evaluation can be slow)
        from datetime import timedelta
        nccl_timeout = timedelta(hours=8)
        dist.init_process_group(
            backend='nccl' if args.device == 'cuda' else 'gloo',
            timeout=nccl_timeout
        )
        logger.info(f"Initialized DDP: rank={ddp_rank}, world_size={ddp_world_size}, local_rank={ddp_local_rank}, timeout=8h")

    # Set random seeds for reproducibility.
    # In DDP, each rank gets a different seed for data augmentation
    # but the same seed for model initialization.
    base_seed = args.seed
    seed = base_seed + (ddp_rank if use_ddp else 0)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(base_seed)  # Same model init across all ranks
    torch.cuda.manual_seed_all(base_seed)  # Same model init across all GPUs
    logger.info(f"Set random seed: base={base_seed}, rank_seed={seed}")

    # Check for inference mode early (before dataloader creation)
    inference_mode = args.inference or config.get('inference', {}).get('mode', False)
    if inference_mode:
        config.setdefault('inference', {})['mode'] = True   # one source of truth for --inference and inference.mode
    if inference_mode:
        logger.info("=" * 60)
        logger.info("INFERENCE MODE DETECTED - Skipping train dataset indexing")
        logger.info("=" * 60)

    # Pure multi-object mode (MOVi datasets only)
    logger.info("Using multi-dataset loader")
    logger.info(f"Data path: {config['data']['dataset_path']}")
    logger.info(f"Split file: {config['data']['split_file']} (required)")
    from src.data.multi_dataloader import check_scene_set
    logger.info(f"Scene set: {check_scene_set(config)}")
    if ((config.get('data', {}) or {}).get('if_point_cloud') or {}).get('enable', False):
        # never forwarded to the datasets: the run would silently use mesh vertices
        raise NotImplementedError('data.if_point_cloud.enable: point-cloud inputs are not supported by train.py')
    if not inference_mode:
        logger.info("Creating dataloaders... this may take a few minutes on first run")
    # Get noise configuration from config
    noise_config = config['data'].get('augmentation', {})
    if noise_config:
        # Add dt for physics calculations
        noise_config['dt'] = config['data'].get('dt', 1.0)
        logger.info(f"Noise augmentation config: {noise_config}")

    train_loader, val_loader, test_loader = create_multi_dataloaders(
        data_path=config['data']['dataset_path'],
        train_batch_size=config['training']['train_batch_size'],
        val_batch_size=config['training']['val_batch_size'],
        eval_batch_size=config['training']['eval_batch_size'],
        sequence_length=config['data']['sequence_length'],
        num_workers=config['data']['num_workers'],
        train_ratio=train_ratio,
        val_ratio=val_ratio,
        split_file=config.get('data', {}).get('split_file', 'dataset_split.json'),
        skip_step_size=config['data'].get('skip_step_size', 1),
        skip_first_frames=config['data'].get('skip_first_frames', 2),
        load_gt_motion=config['data'].get('load_gt_motion', False),
        distributed=use_ddp,
        rank=ddp_rank,
        world_size=ddp_world_size,
        noise_config=noise_config,
        augmentation=config['data'].get('augmentation', {}),
        anchor_fps_start=config['data'].get('anchor_fps_start', 'index0'),
        fps_order=config['data'].get('fps_order', False),
        object_scale_file=config['data'].get('object_scale_file'),
        anchor_k=config['data'].get('anchor_k', None),
        max_vertices_per_object=config['data'].get('max_vertices_per_object', None),
        max_objects_per_frame=config['data'].get('max_objects_per_frame', None),
        if_phys_param=config.get('model', {}).get('if_phys_param', False),
        random_step_size=config['data'].get('random_step_size', {'enable': False}),  # Multi-timescale training config
        partial_view=config['data'].get('partial_view'),  # partial point clouds of a simulated depth camera
        stride=config['data'].get('stride', 1),  # windows every `stride` frames; with random_start_frame: the epoch length
        inference_mode=inference_mode  # Skip train dataset indexing in inference mode
    )
    try:
        if enable_validation:
            logger.info(f"Dataloaders ready. train_batches={len(train_loader)}, val_batches={len(val_loader)}")
        else:
            logger.info(f"Dataloaders ready. train_batches={len(train_loader)} (val disabled)")
    except Exception:
        logger.info("Dataloaders ready.")
    
    # Create model
    model = CUT3RDynamicsModel(config)

    # Print parameter counts for each module
    logger.info("=" * 80)
    logger.info("Model Parameter Breakdown:")
    logger.info("=" * 80)

    # First pass: collect all parameter counts
    module_stats = []
    total_params = 0

    if hasattr(model, 'enable_multi_object') and model.enable_multi_object:
        # Multi-object modules
        if hasattr(model, 'shared_obj_encoder'):
            encoder_params = sum(p.numel() for p in model.shared_obj_encoder.parameters())
            total_params += encoder_params
            module_stats.append(("PointNet Encoder (shared)", encoder_params))

        if getattr(model, 'encoder_film', None) is not None:
            film_params = sum(p.numel() for p in model.encoder_film.parameters())
            total_params += film_params
            module_stats.append(("Encoder Step-Size FiLM", film_params))

        if getattr(model, 'avp', None) is not None:
            avp_params = sum(p.numel() for p in model.avp.parameters())
            total_params += avp_params
            module_stats.append(("Anchor-Vertex Pooling", avp_params))

        if hasattr(model, 'obj_feature_proj'):
            obj_proj_params = sum(p.numel() for p in model.obj_feature_proj.parameters())
            total_params += obj_proj_params
            module_stats.append(("Object Feature Projection", obj_proj_params))

        if hasattr(model, 'anchor_query_proj'):
            anchor_proj_params = sum(p.numel() for p in model.anchor_query_proj.parameters())
            total_params += anchor_proj_params
            module_stats.append(("Anchor Query Projection", anchor_proj_params))

    # State tokens (learnable parameters)
    if hasattr(model, 'state_tokens'):
        state_params = model.state_tokens.numel()
        total_params += state_params
        module_stats.append(("State Tokens (learnable)", state_params))

    if hasattr(model, 'state_pos'):
        state_pos_params = model.state_pos.numel()
        total_params += state_pos_params
        module_stats.append(("State Positions (learnable)", state_pos_params))

    # Decoder
    if hasattr(model, 'decoder'):
        decoder_params = sum(p.numel() for p in model.decoder.parameters())
        total_params += decoder_params
        module_stats.append(("Bidirectional Decoder", decoder_params))

    # Predictor
    if hasattr(model, 'predictor'):
        predictor_params = sum(p.numel() for p in model.predictor.parameters())
        total_params += predictor_params
        module_stats.append(("Multi-scale Predictor", predictor_params))

    # Transform layer
    if hasattr(model, 'transform_layer'):
        transform_params = sum(p.numel() for p in model.transform_layer.parameters())
        total_params += transform_params
        module_stats.append(("Transform Layer", transform_params))

    # The totals count every parameter of the model, so a part not listed above is shown, not left out
    all_params = sum(p.numel() for p in model.parameters())
    if all_params > total_params:
        module_stats.append(("Other parameters", all_params - total_params))
    total_params = all_params
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    # Second pass: print with percentages
    for name, param_count in module_stats:
        percentage = (param_count / total_params * 100) if total_params > 0 else 0
        logger.info(f"  {name:35s} {param_count:>12,} params  ({percentage:>5.2f}%)")



    logger.info("=" * 80)
    total_params_millions = total_params / 1_000_000
    model_size_mb = total_params * 4 / 1024 / 1024
    logger.info(f"  TOTAL PARAMETERS:                 {total_params:>12,} params ({total_params_millions:.1f}M)")
    logger.info(f"  TOTAL TRAINABLE PARAMETERS:       {trainable_params:>12,} params ({trainable_params / 1_000_000:.1f}M)")
    logger.info(f"  TOTAL MODEL SIZE:                 {model_size_mb:>12.2f} MB (FP32)")
    logger.info("=" * 80)

    # (model.predictor.rmsnorm_eps is applied in the model constructor)

    # Optional torch.compile of fixed-shape submodules for TRAINING (training.compile). The decoder and the
    # predictor see (B, 16 object slots, 4 anchors) tensors, so they compile to a few graphs whose fused kernels
    # replace thousands of small launches per step. nn.Module.compile() compiles in place: parameter names and
    # checkpoints are unchanged. Results differ from eager at floating-point rounding level (fused kernels and
    # the compiled dropout RNG), so this is a per-run choice, recorded in the log.
    _tc = config.get('training', {}).get('compile', {}) or {}
    if _tc.get('enable', False) and not inference_mode:
        _tc_modules = list(_tc.get('modules', ['decoder', 'predictor']))
        _tc_mode = _tc.get('mode', 'default')
        _tc_dynamic = _tc.get('dynamic', False)
        # modules whose inputs change size from batch to batch (e.g. the encoder: vertex count with no padding)
        _tc_dyn_modules = set(_tc.get('dynamic_modules', []))
        for _name in _tc_modules:
            if not hasattr(model, _name):
                raise ValueError(f"training.compile.modules: model has no submodule '{_name}'")
            getattr(model, _name).compile(mode=_tc_mode, dynamic=(True if _name in _tc_dyn_modules else _tc_dynamic))
        logger.info(f"[Compile] training torch.compile ACTIVE: modules={_tc_modules} mode={_tc_mode} "
                    f"dynamic={_tc_dynamic} dynamic_modules={sorted(_tc_dyn_modules)}")

    # training.optimizer.frozen_params: parameters (or module prefixes) that keep their loaded values when fine-tuning.
    # Adam moves a parameter by about lr per step whatever its scale: even at lr 4e-7, the acceleration head's output
    # bias (predictor.fusion_out.bias, a constant acceleration in m/frame^2; gravity is 1.7e-4 at 240 fps) drifts by
    # several percent of gravity within one epoch.
    _frozen = [str(f) for f in ((config.get('training', {}).get('optimizer') or {}).get('frozen_params') or [])]
    if _frozen:
        _hit = []
        for _n, _p in model.named_parameters():
            _k = _n.replace('module.', '').replace('_orig_mod.', '')
            if any(_k == f or _k.startswith(f + '.') for f in _frozen):
                _p.requires_grad_(False)
                _hit.append(_k)
        _miss = [f for f in _frozen if not any(h == f or h.startswith(f + '.') for h in _hit)]
        if _miss:
            raise ValueError(f"training.optimizer.frozen_params: no parameter named {_miss}")
        logger.info(f"[Frozen] training.optimizer.frozen_params: {_hit} keep their loaded values")

    if use_ddp:
        # Convert BatchNorm to SyncBatchNorm for multi-GPU training;
        # otherwise each GPU keeps different BN statistics.
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
        logger.info("Converted BatchNorm to SyncBatchNorm for multi-GPU training")

        device_id = ddp_local_rank if args.device == 'cuda' else None
        # training.ddp (optional): static_graph lets DDP record which parameters are used once and skip the
        # per-iteration autograd-graph search that find_unused_parameters=True otherwise performs (valid when
        # every iteration uses the same parameters, as the teacher-forced training step does);
        # gradient_as_bucket_view makes gradients views into the communication buckets (no copy).
        _ddp_cfg = config.get('training', {}).get('ddp', {}) or {}
        model = DDP(
            model.to(args.device),
            device_ids=[device_id] if device_id is not None else None,
            output_device=device_id if device_id is not None else None,
            find_unused_parameters=True,
            static_graph=bool(_ddp_cfg.get('static_graph', False)),
            gradient_as_bucket_view=bool(_ddp_cfg.get('gradient_as_bucket_view', False))
        )
        if _ddp_cfg:
            logger.info(f"[DDP] options: static_graph={bool(_ddp_cfg.get('static_graph', False))} "
                        f"gradient_as_bucket_view={bool(_ddp_cfg.get('gradient_as_bucket_view', False))}")
    # training.rank_seeded_rng: after the (identical, rank-0-broadcast) model initialisation, each rank's torch RNG
    # gets its own seed, so the per-batch rotation augmentation and dropout / stochastic-depth masks are drawn
    # independently per rank; with the shared seed every rank drew the same ones at every iteration
    if use_ddp and bool(config.get('training', {}).get('rank_seeded_rng', False)):
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        logger.info(f"[Seed] torch RNG re-seeded per rank after model setup: {seed}")
    
    # Create trainer
    # Visualization CLI params
    viz_params = {
        'arrow_scale': float(os.environ.get('ARROW_SCALE', '1.0')),
        'point_size': float(os.environ.get('POINT_SIZE', '0.001')),
        'arrow_thickness': float(os.environ.get('ARROW_THICKNESS', '0.002')),
        'arrow_subsample': int(os.environ.get('ARROW_SUBSAMPLE', '1')),
        'arrow_enable': int(os.environ.get('ARROW_ENABLE', '1')),
    }

    trainer = Trainer(config, model, train_loader, val_loader, device=args.device, viz_params=viz_params)
    
    # Resume from checkpoint if specified
    if checkpoint_path:
        # Load checkpoint to CPU to avoid GPU0 memory spike in multi-GPU resume
        checkpoint = torch.load(checkpoint_path, map_location='cpu')

        # Handle DDP/non-DDP checkpoint compatibility:
        # DDP models have a "module." prefix in state_dict keys, non-DDP models don't
        state_dict = checkpoint['model_state_dict']
        is_model_ddp = isinstance(model, DDP)

        # The clean-arm switches add no parameters, so a checkpoint loads under either setting
        # without a shape error. Make any difference between the checkpoint's config and the
        # current config explicit (fine-tuning across a switch change is intended; evaluation
        # with mismatched switches is not).
        _ckpt_cfg = checkpoint.get('config')
        if isinstance(_ckpt_cfg, dict):
            def _flat(d, prefix):
                out = {}
                for _k, _v in (d or {}).items():
                    if isinstance(_v, dict):
                        out.update(_flat(_v, f'{prefix}{_k}.'))
                    else:
                        out[f'{prefix}{_k}'] = _v
                return out
            # Every model.* setting is compared, except the ones that do not change the function a checkpoint
            # computes in evaluation: speed paths equal up to rounding, training-only settings, the rollout's
            # numerics-only Kabsch reference (equal in exact arithmetic), and rmsnorm_eps (its own rule below).
            _eval_neutral = ('model.knn_packed', 'model.encoder.packed', 'model.encoder.packed_pooling',
                             'model.kabsch_reference', 'model.intersection_handling.after_kabsch_loss',
                             'model.predictor.rmsnorm_eps', 'model.predictor.fp32_output')
            _numerics_only = ('model.kabsch_reference',)
            # Settings absent from a checkpoint's configuration take these defaults
            _defaults = {'model.decoder.block_attn_res.final_aggregate': False, 'model.kabsch_reference': 'advancing',
                         'model.encoder.subset_small_objects': 'all_below_8'}
            # data.* settings that change what the model sees in evaluation (the object-scale sidecar may be swapped
            # for the target dataset in cross-dataset evaluation: only whether one is used must agree)
            _data_switches = {
                'data.center_scene_by_t0': (lambda c: bool(c.get('data', {}).get('center_scene_by_t0', False))),
                'data.anchor_fps_start': (lambda c: c.get('data', {}).get('anchor_fps_start', 'index0')),
                'data.fps_order': (lambda c: bool(c.get('data', {}).get('fps_order', False))),
                'data.anchor_k': (lambda c: c.get('data', {}).get('anchor_k')),
                'data.object_scale_file (set)': (lambda c: bool(c.get('data', {}).get('object_scale_file'))),
            }
            _defaults['model.predictor.use_rope'] = True
            _fk, _fc = _flat(_ckpt_cfg.get('model'), 'model.'), _flat(config.get('model'), 'model.')

            def _inactive(name, cfg):
                # settings without effect in evaluation: dropout / stochastic depth rates, frozen-BN training mode,
                # the object-object handler's parameters while it is disabled
                leaf = name.rsplit('.', 1)[-1]
                if 'dropout' in leaf or 'drop_path' in leaf or leaf == 'freeze_bn_stats':
                    return True
                return (name.startswith('model.intersection_handling.object.') and name != 'model.intersection_handling.object.enable'
                        and not cfg.get('model.intersection_handling.object.enable', False))
            _mismatch = []
            for _name in sorted(set(_fk) | set(_fc)):
                _a = _fk.get(_name, _defaults.get(_name, '<absent>'))
                _b = _fc.get(_name, _defaults.get(_name, '<absent>'))
                if _a == _b or (_inactive(_name, _fk) and _inactive(_name, _fc)):
                    continue
                if _name.startswith(_eval_neutral) and not _name.startswith(_numerics_only):
                    continue
                logger.warning(f"[Checkpoint] switch mismatch: {_name} checkpoint={_a} -> current={_b}")
                print(f"[Checkpoint] switch mismatch: {_name} checkpoint={_a} -> current={_b}", flush=True)
                if not _name.startswith(_eval_neutral):
                    _mismatch.append(_name)
            for _name, _get in _data_switches.items():
                _a, _b = _get(_ckpt_cfg), _get(config)
                if _a != _b:
                    logger.warning(f"[Checkpoint] switch mismatch: {_name} checkpoint={_a} -> current={_b}")
                    print(f"[Checkpoint] switch mismatch: {_name} checkpoint={_a} -> current={_b}", flush=True)
                    _mismatch.append(_name)
            # The predictor RMSNorm eps pin reproduces the training function; a checkpoint trained with it pinned
            # must be evaluated with it pinned
            _eps = lambda c: c.get('model', {}).get('predictor', {}).get('rmsnorm_eps', 'default') or 'default'
            if _eps(_ckpt_cfg) != 'default' and _eps(config) != _eps(_ckpt_cfg):
                print(f"[Checkpoint] switch mismatch: model.predictor.rmsnorm_eps checkpoint={_eps(_ckpt_cfg)} -> "
                      f"current={_eps(config)}", flush=True)
                _mismatch.append('model.predictor.rmsnorm_eps')
            # Evaluation must run the function the checkpoint was trained as: refuse behaviour-setting mismatches
            # unless a name is listed in inference.allow_switch_mismatch (deliberate knob experiments)
            _allowed = set(config.get('inference', {}).get('allow_switch_mismatch', []) or [])
            _refused = [n for n in _mismatch if n not in _allowed]
            if (args.inference or config.get('inference', {}).get('mode', False)) and _refused:
                raise ValueError(f"[Checkpoint] evaluation config differs from the checkpoint's training config in "
                                 f"{_refused}; pass the checkpoint's own config (BASE_CONFIG) or list the names in "
                                 f"inference.allow_switch_mismatch")
        else:
            print("[Checkpoint] no config stored in checkpoint; clean-arm switches not cross-checked", flush=True)

        # Check if checkpoint is from DDP (has "module." prefix)
        is_checkpoint_ddp = any(k.startswith('module.') for k in state_dict.keys())

        if is_model_ddp and not is_checkpoint_ddp:
            # Loading non-DDP checkpoint into DDP model: add "module." prefix
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                new_state_dict[f'module.{k}'] = v
            state_dict = new_state_dict
            logger.info("Converted non-DDP checkpoint to DDP format (added 'module.' prefix)")
        elif not is_model_ddp and is_checkpoint_ddp:
            # Loading DDP checkpoint into non-DDP model: remove "module." prefix
            from collections import OrderedDict
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:] if k.startswith('module.') else k  # Remove 'module.'
                new_state_dict[name] = v
            state_dict = new_state_dict
            logger.info("Converted DDP checkpoint to non-DDP format (removed 'module.' prefix)")

        # ============================================================
        # Adapt state tokens if size changed (e.g., 48 -> 27)
        # Uses 3D interpolation to preserve learned patterns
        # ============================================================
        state_dict = _adapt_state_tokens_if_needed(state_dict, model, logger)
        # Height input channel: zero-pad conv1 before the size-mismatch filter (exact warm start).
        state_dict = _adapt_state_dict_for_height_input(model, state_dict)

        model_state = model.state_dict()

        # Filter out size-mismatched keys before loading (these layers are re-initialized)
        keys_to_remove = []
        for k, v in state_dict.items():
            if k in model_state and v.shape != model_state[k].shape:
                logger.warning(f"[Checkpoint] Size mismatch for '{k}': checkpoint {v.shape} vs model {model_state[k].shape}. "
                              f"Will SKIP (re-initialize from scratch).")
                keys_to_remove.append(k)
        for k in keys_to_remove:
            del state_dict[k]
        if keys_to_remove:
            logger.warning(f"[Checkpoint] Removed {len(keys_to_remove)} size-mismatched keys. These layers will be randomly initialized.")
        model.load_state_dict(state_dict, strict=True)

        if args.compile:
            if inference_mode:
                for name in ('shared_obj_encoder', 'decoder', 'predictor'):
                    if hasattr(model, name):
                        setattr(model, name, torch.compile(getattr(model, name),
                                                           mode='max-autotune-no-cudagraphs',
                                                           dynamic=False))
                logger.info("torch.compile enabled for encoder, decoder and predictor")
            else:
                logger.warning("--compile is supported for inference only; ignoring it for training")

        # Restore optimizer/scheduler state.
        # A weights-only checkpoint (no optimizer state) starts the optimizer fresh.
        optimizer_state = checkpoint.get('optimizer_state_dict')
        scheduler_state = checkpoint.get('scheduler_state_dict')

        if optimizer_state is None:
            logger.info("Checkpoint contains model weights only; optimizer starts fresh")
        elif trainer.optimizer is not None:
            try:
                trainer.optimizer.load_state_dict(optimizer_state)
            except (ValueError, RuntimeError) as e:
                # A different parameter-group layout (e.g. training.optimizer.decay_matrices_only) cannot resume the
                # optimizer, and restoring the scheduler and epoch without it would mis-assign the group learning rates
                raise ValueError(f"[Checkpoint] optimizer state does not match the configured optimizer ({e}); resume "
                                 f"with the checkpoint's own optimizer settings, or start from its weights with a "
                                 f"weights-only checkpoint file") from e
            # The loaded param groups carry the eps they were saved with; the config decides (default 1e-8).
            _cfg_eps = float(config['training'].get('optimizer', {}).get('eps', 1e-8))
            for _g in trainer.optimizer.param_groups:
                if _g.get('eps') != _cfg_eps:
                    logger.info(f"[Optimizer] eps {_g.get('eps')} from the checkpoint replaced by the configured {_cfg_eps:g}")
                    _g['eps'] = _cfg_eps
        else:
            logger.info("Inference-only mode: skipping optimizer state loading")

        scheduler_loaded = False
        if scheduler_state is not None and trainer.scheduler is not None:
            try:
                trainer.scheduler.load_state_dict(scheduler_state)
                scheduler_loaded = True
                logger.info("Loaded scheduler state from checkpoint")
            except (KeyError, TypeError) as e:
                logger.warning(f"Could not load scheduler state (likely different scheduler type): {e}")
                logger.info("Scheduler will start fresh with current configuration")
        elif trainer.scheduler is None:
            logger.info("Inference-only mode: skipping scheduler state loading")

        trainer.best_val_loss = checkpoint.get('best_val_loss', float('inf'))
        ckpt_epoch = checkpoint.get('epoch', None)
        if optimizer_state is not None and ckpt_epoch is not None and not inference_mode:
            # A resumed run writes into a new directory: carry the source run's best.pth (and its metric) over, so the
            # new directory always holds the best checkpoint of the whole trajectory, also when the source run was
            # interrupted between publishing best.pth and latest.pth (best then holds the lower metric).
            _src_best = Path(checkpoint_path).parent / 'best.pth'
            if _src_best.exists():
                _b = torch.load(_src_best, map_location='cpu', weights_only=False)
                _b_metric = float(_b.get('best_val_loss', float('inf')))
                if _b_metric < trainer.best_val_loss:
                    logger.info(f"[Resume] {_src_best} (epoch {_b.get('epoch')}) holds a lower best metric {_b_metric:.6f} "
                                f"than the resumed checkpoint ({trainer.best_val_loss:.6f}); using it")
                    trainer.best_val_loss = _b_metric
                if trainer.is_main:
                    Trainer._atomic_save(_b, trainer.checkpoint_dir / 'best.pth')
                    logger.info(f"[Resume] best checkpoint of the source run copied to {trainer.checkpoint_dir / 'best.pth'}")
                del _b
        if optimizer_state is not None and ckpt_epoch is not None and not inference_mode:
            # Full training checkpoint (weights + optimizer + epoch): resume, i.e. continue the
            # epoch count where the run stopped. A weights-only checkpoint instead starts a new
            # schedule at epoch 0 (fine-tuning / warm start).
            trainer.start_epoch = int(ckpt_epoch) + 1
            if scheduler_loaded and not checkpoint.get('scheduler_stepped', False):
                # A checkpoint without 'scheduler_stepped' stored the scheduler state from before the
                # end-of-epoch step; advance once so the LR matches a continuous run.
                trainer.scheduler.step()
                logger.info("Checkpoint without scheduler_stepped: advanced the LR scheduler by one epoch to match a continuous run")
            logger.info(f"Resuming training after completed epoch {ckpt_epoch} (next epoch: {trainer.start_epoch})")
        else:
            trainer.start_epoch = 0
            if ckpt_epoch is not None and not inference_mode:
                logger.info(f"Loaded weights from epoch {ckpt_epoch}; training restarts at epoch 0 with a fresh schedule")
        trainer.epoch = trainer.start_epoch
        if getattr(trainer, 'ema', None) is not None:
            if trainer.start_epoch > 0 and 'ema_state_dict' in checkpoint:
                trainer.ema.load_state_dict(checkpoint['ema_state_dict'], checkpoint.get('ema_num_updates'))
                logger.info(f"[EMA] restored averaged weights from the checkpoint (updates={trainer.ema.num_updates})")
            else:
                trainer.ema.reset()
                logger.info("[EMA] average restarted from the loaded weights")
    
    # inference_mode already determined earlier (before dataloader creation)
    if inference_mode:
        # ============================================================
        # INFERENCE MODE: Only generate videos and compute metrics
        # ============================================================
        logger.info("=" * 80)
        logger.info("INFERENCE MODE: Generating videos and computing metrics")
        logger.info("=" * 80)

        if not checkpoint_path:
            raise ValueError("Inference mode requires --checkpoint argument!")

        try:
            trainer.run_inference()
        except KeyboardInterrupt:
            logger.info("Inference interrupted by user (Ctrl+C)")
            raise
        except Exception as e:
            logger.error(f"Inference failed with error: {e}")
            raise
        finally:
            if use_ddp:
                try:
                    dist.destroy_process_group()
                    logger.info("DDP process group destroyed successfully")
                except Exception as cleanup_error:
                    logger.warning(f"Failed to destroy DDP process group: {cleanup_error}")
    else:
        # ============================================================
        # TRAINING MODE: Normal training loop
        # ============================================================
        try:
            trainer.train()
        except KeyboardInterrupt:
            logger.info("Training interrupted by user (Ctrl+C)")
            raise
        except Exception as e:
            logger.error(f"Training failed with error: {e}")
            raise
        finally:
            # Always clean up the DDP process group, even on error
            if use_ddp:
                try:
                    dist.destroy_process_group()
                    logger.info("DDP process group destroyed successfully")
                except Exception as cleanup_error:
                    logger.warning(f"Failed to destroy DDP process group: {cleanup_error}")


if __name__ == '__main__':
    main()