physics_audio/training_evaluation/trainer.py

# ~/mlpa/training_evaluation/trainer.py
# Full universal version with all fixes and improvements:
# - Modern torch.amp API throughout
# - Robust post-jump visualization and loss computation (fresh no-grad forward on selected best state)
# - Improved selection logic: always accept the best candidate if it improves over pre_jump_loss
#   (also accepts entropic if best doesn't but entropic does; otherwise continues current state)
# - Parallel rollouts with robust try/except sequential fallback
# - CPU transfer for state_dict/data to ensure safe multiprocessing pickling
# - Early stopping in rollouts
# - Random uniform noise stds for better exploration + fixed zero-noise continuation
# - All previous features preserved where compatible
# - Integrated rich.progress for enhanced visual feedback (main loop + nested rollout bars)
# - Live per-candidate printing during rollouts for immediate feedback

import time
import random
import copy
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp

from rich.progress import Progress, BarColumn, MofNCompleteColumn, TextColumn, TimeRemainingColumn
from torch.amp import autocast, GradScaler
from torch.nn.utils import clip_grad_norm_
from geoopt.optim import RiemannianAdam
from geoopt.manifolds import Stiefel
from concurrent.futures import ProcessPoolExecutor, as_completed
from . import config as training_config
from .config import *
from .model import StiefelDampedCoupledInharmGR
from .utils import stiefel_dist, safe_proj, align_and_compute_freq, get_pca_initial_basis, manifold
from .losses import total_loss
from .fr_utils import (
    mode_amplitudes_from_stiefel,
    modal_spectral_envelope,
    fr_loss_kwargs_from_batch,
    resolve_target_mode_amps,
)
from .audio_utils import modal_synthesis_torch, extract_coupling_skew
from .viz import save_trajectory_frame, save_detailed_pyramid_plot, plot_smith_chart

# Force spawn for CUDA safety
mp.set_start_method('spawn', force=True)


def _invariant_loss_kwargs(
    model: nn.Module,
    reference_coupling_skew: torch.Tensor | None,
    fr_invariant_weight: float,
    fr_invariant_coupling: float | None = None,
    fr_invariant_speed: float | None = None,
    fr_invariant_inharm: float | None = None,
    fr_invariant_modal: float | None = None,
    fr_replace_mse_priors: bool | None = None,
    reference_speed_scalars: torch.Tensor | None = None,
    reference_inharm_b: torch.Tensor | None = None,
) -> dict:
    kw = {
        "log_base_rate": model.log_base_rate,
        "log_slope": model.log_slope,
        "coupling_skew": extract_coupling_skew(model),
        "target_coupling_skew": reference_coupling_skew,
        "target_speed_scalars": reference_speed_scalars,
        "target_inharm_b": reference_inharm_b,
        "fr_invariant_weight_override": fr_invariant_weight,
    }
    if fr_invariant_coupling is not None:
        kw["fr_invariant_coupling_override"] = fr_invariant_coupling
    if fr_invariant_speed is not None:
        kw["fr_invariant_speed_override"] = fr_invariant_speed
    if fr_invariant_inharm is not None:
        kw["fr_invariant_inharm_override"] = fr_invariant_inharm
    if fr_invariant_modal is not None:
        kw["fr_invariant_modal_override"] = fr_invariant_modal
    if fr_replace_mse_priors is not None:
        kw["fr_replace_mse_priors_override"] = fr_replace_mse_priors
    return kw


def loss_to_gamma(loss: float) -> float:
    if loss <= SWR_BEST_LOSS_EST:
        return 0.0
    swr = loss / SWR_BEST_LOSS_EST
    return (swr - 1.0) / (swr + 1.0)


def add_euclidean_noise(model: nn.Module, std: float):
    with torch.no_grad():
        for name, param in model.named_parameters():
            if not param.requires_grad or name == "base":
                continue
            if RECIPROCAL_NOISE and name in ["log_speed", "log_base_rate", "log_slope",
                                             "raw_lin_b", "raw_quad_b", "raw_coupling_strength"]:
                exp_param = torch.exp(param)
                noise = torch.randn_like(exp_param) * std * RECIPROCAL_STRENGTH
                perturbed = 1.0 / (1.0 / exp_param + noise)
                perturbed = perturbed.clamp(min=exp_param * 0.1, max=exp_param * 10.0)
                param.copy_(torch.log(perturbed.clamp(min=1e-8)))
            else:
                param.add_(torch.randn_like(param) * std)

        if TANGENT_PROJECT_AFTER_NOISE:
            model.vel_dir_raw.data = manifold.proju(model.base, model.vel_dir_raw.data)


def _rollout_worker(arg):
    (idx, pre_state_dict_cpu, noise_std, data_points_cpu, times_cpu,
     initial_basis_cpu, worker_seed, lr_geo, lr_slow, prior_targets_cpu,
     rollout_horizon, fr_invariant_weight, fr_invariant_coupling, fr_invariant_speed,
     fr_invariant_inharm, fr_invariant_modal, fr_replace_mse_priors, fr_mode_weight, fr_spectral_weight,
     target_mode_amps_cpu, target_spectrum_cpu, reference_coupling_skew_cpu,
     reference_speed_scalars_cpu, reference_inharm_b_cpu) = arg

    torch.manual_seed(worker_seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    from .model import StiefelDampedCoupledInharmGR
    from .config import DIM, K_MODES, USE_AMP
    from .losses import total_loss
    from torch.nn.utils import clip_grad_norm_

    model = StiefelDampedCoupledInharmGR(DIM, K_MODES, initial_basis_cpu.to(device))
    model.load_state_dict(pre_state_dict_cpu)
    model.to(device)

    add_euclidean_noise(model, noise_std)

    optimizer = RiemannianAdam([
        {'params': model.base, 'lr': lr_geo},
        {'params': [p for n, p in model.named_parameters() if p is not model.base], 'lr': lr_slow}
    ], stabilize=10)

    scaler = GradScaler('cuda', enabled=USE_AMP)

    data_points = data_points_cpu.to(device)
    times = times_cpu.to(device)
    target_mode_amps = target_mode_amps_cpu.to(device) if target_mode_amps_cpu is not None else None
    target_spectrum = target_spectrum_cpu.to(device) if target_spectrum_cpu is not None else None
    reference_coupling_skew = (
        reference_coupling_skew_cpu.to(device) if reference_coupling_skew_cpu is not None else None
    )
    reference_speed_scalars = (
        reference_speed_scalars_cpu.to(device) if reference_speed_scalars_cpu is not None else None
    )
    reference_inharm_b = (
        reference_inharm_b_cpu.to(device) if reference_inharm_b_cpu is not None else None
    )

    best_local_loss = float('inf')
    best_state = None
    steps_no_improve = 0
    patience = 4000
    min_steps = 2000
    steps_performed = 0

    for local_step in range(rollout_horizon):
        optimizer.zero_grad(set_to_none=True)

        with autocast('cuda', enabled=USE_AMP):
            preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq = model(times)
            fr_kw = fr_loss_kwargs_from_batch(
                preds, data_points,
                fr_mode_weight=fr_mode_weight,
                fr_spectral_weight=fr_spectral_weight,
                fr_invariant_weight=fr_invariant_weight,
                fr_invariant_modal=fr_invariant_modal,
                target_mode_amps=target_mode_amps,
                target_spectrum=target_spectrum,
            )
            loss = total_loss(
                preds, data_points, damping_rates, coupling_strength, inharm_b, speed_scalars,
                prior_targets=prior_targets_cpu,
                **_invariant_loss_kwargs(
                    model, reference_coupling_skew, fr_invariant_weight,
                    fr_invariant_coupling, fr_invariant_speed, fr_invariant_inharm, fr_invariant_modal,
                    fr_replace_mse_priors,
                    reference_speed_scalars, reference_inharm_b,
                ),
                **fr_kw,
            )

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

        steps_performed = local_step + 1
        current_local_loss = loss.item()

        if current_local_loss < best_local_loss - 1e-6:
            best_local_loss = current_local_loss
            best_state = copy.deepcopy(model.state_dict())
            steps_no_improve = 0
        else:
            steps_no_improve += 1

        if steps_no_improve >= patience and steps_performed >= min_steps:
            break

    # Load best state and compute accurate final loss
    if best_state is not None:
        model.load_state_dict(best_state)

    with torch.no_grad():
        preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq = model(times)
        fr_kw = fr_loss_kwargs_from_batch(
            preds, data_points,
            fr_mode_weight=fr_mode_weight,
            fr_spectral_weight=fr_spectral_weight,
            fr_invariant_weight=fr_invariant_weight,
            fr_invariant_modal=fr_invariant_modal,
            target_mode_amps=target_mode_amps,
            target_spectrum=target_spectrum,
        )
        final_loss = total_loss(
            preds, data_points, damping_rates, coupling_strength, inharm_b, speed_scalars,
            prior_targets=prior_targets_cpu,
            **_invariant_loss_kwargs(
                model, reference_coupling_skew, fr_invariant_weight,
                fr_invariant_coupling, fr_invariant_speed, fr_invariant_inharm, fr_invariant_modal,
                fr_replace_mse_priors,
                reference_speed_scalars, reference_inharm_b,
            ),
            **fr_kw,
        ).item()

    # Return CPU state_dict for safe pickling
    best_state_cpu = {k: v.cpu().clone().detach() if torch.is_tensor(v) else v for k, v in model.state_dict().items()}

    return {
        'idx': idx,
        'noise_std': noise_std,
        'final_loss': final_loss,
        'best_state': best_state_cpu,
        'steps': steps_performed
    }


def run_single_seed(
    seed: int,
    noise_amp: float = NOISE_AMP,
    force_punctuated: bool | None = FORCE_PUNCTUATED,
    *,
    real_audio_data: torch.Tensor | None = None,
    real_audio_times: torch.Tensor | None = None,
    real_audio_initial_basis: torch.Tensor | None = None,
    target_waveform: torch.Tensor | None = None,
    audio_sr: int = REAL_AUDIO_SR,
    audio_duration: float | None = None,
    prior_targets: dict | None = None,
    max_steps: int | None = None,
    stft_weight: float = 0.0,
    fr_invariant_weight: float | None = None,
    fr_invariant_coupling: float | None = None,
    fr_invariant_speed: float | None = None,
    fr_invariant_inharm: float | None = None,
    fr_invariant_modal: float | None = None,
    fr_replace_mse_priors: bool | None = None,
    fr_mode_weight: float | None = None,
    fr_spectral_weight: float | None = None,
    preinitialized_model: StiefelDampedCoupledInharmGR | None = None,
    jump_test: dict | None = None,
):
    start_time = time.time()
    training_steps = max_steps if max_steps is not None else MAX_STEPS
    use_real_audio = real_audio_data is not None
    if fr_invariant_weight is None:
        fr_invariant_weight = training_config.fr_invariant_weight
    if fr_invariant_coupling is None:
        fr_invariant_coupling = training_config.fr_invariant_coupling
    if fr_invariant_speed is None:
        fr_invariant_speed = training_config.fr_invariant_speed
    if fr_invariant_inharm is None:
        fr_invariant_inharm = training_config.fr_invariant_inharm
    if fr_invariant_modal is None:
        fr_invariant_modal = training_config.fr_invariant_modal
    if fr_replace_mse_priors is None:
        fr_replace_mse_priors = training_config.fr_replace_mse_priors
    if fr_mode_weight is None:
        fr_mode_weight = training_config.fr_mode_weight
    if fr_spectral_weight is None:
        fr_spectral_weight = training_config.fr_spectral_weight

    torch.manual_seed(DATA_SEED + seed)
    np.random.seed(DATA_SEED + seed)
    random.seed(DATA_SEED + seed)

    if use_real_audio:
        data_points = real_audio_data.to(device)
        times = real_audio_times.to(device) if real_audio_times is not None else TIMES
        initial_basis = real_audio_initial_basis if real_audio_initial_basis is not None else get_pca_initial_basis(data_points, K_MODES)
        exact_points = data_points
        true_base = None
        true_vel_dir = None
        true_freq = None
        true_coupling_skew = None
    else:
        # === True physics generation ===
        raw = torch.randn(DIM, K_MODES, device=device)
        true_base = safe_proj(raw)

        true_vel_dir_raw = torch.randn(DIM, K_MODES, device=device)
        true_vel_dir = manifold.proju(true_base, true_vel_dir_raw)
        true_vel_dir = true_vel_dir / (true_vel_dir.norm(dim=0, keepdim=True) + 1e-8)

        true_freq = IDEAL_HARMONICS * torch.sqrt(1 + TRUE_INHARM_B * IDEAL_HARMONICS.pow(2))
        true_vel = true_vel_dir * VELOCITY_SCALE_BASE * true_freq

        true_coupling_raw = torch.randn(K_MODES, K_MODES, device=device) * 0.05
        true_coupling_skew = true_coupling_raw.tril(diagonal=-1) - true_coupling_raw.triu(diagonal=1)
        true_coupling_vel = manifold.proju(true_base, true_base @ true_coupling_skew)

        true_vel_total = true_vel + TRUE_COUPLING_STRENGTH * true_coupling_vel

        abs_times = torch.abs(TIMES).view(-1, 1, 1)
        envelope = torch.exp(-TRUE_DAMPING_RATES * abs_times)

        base_batch = true_base.unsqueeze(0).expand(N_POINTS, -1, -1)
        vel_batch = TIMES.view(-1, 1, 1) * true_vel_total.unsqueeze(0) * envelope
        exact_points = manifold.expmap(base_batch, vel_batch)

        data_points = exact_points + noise_amp * torch.randn_like(exact_points)
        times = TIMES
        initial_basis = get_pca_initial_basis(data_points, K_MODES)

    reference_coupling_skew = None
    reference_speed_scalars = None
    reference_inharm_b = None
    if not use_real_audio and true_coupling_skew is not None:
        reference_coupling_skew = true_coupling_skew.detach()
        reference_speed_scalars = torch.ones(K_MODES, device=device) * VELOCITY_SCALE_BASE
        reference_inharm_b = TRUE_INHARM_B
    elif prior_targets is not None:
        reference_coupling_skew = prior_targets.get("coupling_skew")
        reference_speed_scalars = prior_targets.get("speed_scalars")
        reference_inharm_b = prior_targets.get("inharm_b")

    target_mode_amps = resolve_target_mode_amps(data_points, prior_targets)
    target_spectrum = modal_spectral_envelope(data_points)

    # === Visualization setup ===
    _, _, proj_matrix = torch.pca_lowrank(data_points.reshape(-1, DIM), q=3, center=True, niter=6)
    proj_matrix = proj_matrix.cpu()
    project_to_3d = lambda pts: (pts.reshape(-1, DIM).cpu() @ proj_matrix).reshape(pts.shape[0], -1, 3).numpy()
    true_proj = project_to_3d(exact_points)
    data_proj = project_to_3d(data_points)
    time_norm = ((times - times.min()) / (times.max() - times.min() + 1e-8)).cpu().numpy()

    # Real-audio uses shorter stagnation patience for punctuated jumps
    if use_real_audio:
        stagnation_patience = REAL_AUDIO_STAGNATION_PATIENCE
        min_step_for_jump = REAL_AUDIO_MIN_STEP_FOR_JUMP
        max_jumps = REAL_AUDIO_MAX_JUMPS
        pop_size = REAL_AUDIO_POP_SIZE
        jump_std_min = REAL_AUDIO_JUMP_STD_MIN
        jump_std_max = REAL_AUDIO_JUMP_STD_MAX
        rollout_horizon = REAL_AUDIO_ROLLOUT_HORIZON
    else:
        stagnation_patience = STAGNATION_PATIENCE
        min_step_for_jump = MIN_STEP_FOR_JUMP_CHECK
        max_jumps = MAX_JUMPS
        pop_size = POP_SIZE
        jump_std_min = JUMP_STD_MIN
        jump_std_max = JUMP_STD_MAX
        rollout_horizon = ROLLOUT_HORIZON

    use_parallel_rollouts = USE_PARALLEL_ROLLOUTS

    # Jump-test overrides (low patience, forced jumps, artificial plateau)
    artificial_plateau_at = None
    force_jump_every = None
    if jump_test:
        if jump_test.get('low_patience') is not None:
            stagnation_patience = jump_test['low_patience']
        if jump_test.get('min_step_for_jump') is not None:
            min_step_for_jump = jump_test['min_step_for_jump']
        if jump_test.get('max_jumps') is not None:
            max_jumps = jump_test['max_jumps']
        artificial_plateau_at = jump_test.get('artificial_plateau_at')
        force_jump_every = jump_test.get('force_jump_every')
        if jump_test.get('pop_size') is not None:
            pop_size = jump_test['pop_size']
        if jump_test.get('rollout_horizon') is not None:
            rollout_horizon = jump_test['rollout_horizon']
        use_parallel_rollouts = not jump_test.get('sequential_rollouts', True)
        if jump_test.get('verbose'):
            print(f"  [jump_test] patience={stagnation_patience} min_step={min_step_for_jump} "
                  f"force_every={force_jump_every} plateau_at={artificial_plateau_at}")

    # === Model & optimizer ===
    if preinitialized_model is not None:
        model = preinitialized_model.to(device)
    else:
        model = StiefelDampedCoupledInharmGR(DIM, K_MODES, initial_basis).to(device)

    def get_curriculum_lrs(step):
        if step < STAGE1_STEPS:
            return STAGE1_LR_GEO, STAGE1_LR_SLOW
        elif step < STAGE2_STEPS:
            return STAGE2_LR_GEO, STAGE2_LR_SLOW
        else:
            return STAGE3_LR_GEO, STAGE3_LR_SLOW

    optimizer = RiemannianAdam([
        {'params': model.base, 'lr': STAGE1_LR_GEO},
        {'params': [p for n, p in model.named_parameters() if p is not model.base], 'lr': STAGE1_LR_SLOW}
    ], stabilize=10)

    scaler = GradScaler('cuda', enabled=USE_AMP)

    global_step = 0
    jumps_performed = 0
    stagnation_steps = 0
    best_loss = float('inf')

    print("\nStarting training loop...\n")

    # Main progress bar with rich
    main_progress = Progress(
        TextColumn("[bold magenta]{task.description}"),
        BarColumn(),
        MofNCompleteColumn(),
        TextColumn("{task.percentage:>3.0f}%"),
        TimeRemainingColumn(),
        TextColumn("• [red]loss:[/red] {task.fields[loss]} • [green]jumps:[/green] {task.fields[jumps]} • [blue]best:[/blue] {task.fields[best]}"),
    )

    with main_progress:
        task_label = f"Real-audio seed {seed}" if use_real_audio else f"Seed {seed}  noise={noise_amp:.3f}"
        main_task = main_progress.add_task(
            task_label,
            total=training_steps,
            loss="--.------",
            jumps=0,
            best="--.------",
        )

        while global_step < training_steps:
            lr_geo, lr_slow = get_curriculum_lrs(global_step)
            if artificial_plateau_at is not None and global_step >= artificial_plateau_at:
                lr_geo, lr_slow = 1e-9, 1e-9
            optimizer.param_groups[0]['lr'] = lr_geo
            optimizer.param_groups[1]['lr'] = lr_slow

            if force_jump_every and global_step >= min_step_for_jump and jumps_performed < max_jumps:
                if global_step % force_jump_every == 0:
                    stagnation_steps = stagnation_patience

            optimizer.zero_grad(set_to_none=True)

            with autocast('cuda', enabled=USE_AMP):
                preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq = model(times)
                synth_waveform = None
                if stft_weight > 0.0 and target_waveform is not None and audio_duration is not None:
                    coupling_skew = extract_coupling_skew(model)
                    synth_waveform = modal_synthesis_torch(
                        full_freq, damping_rates, audio_duration, audio_sr,
                        coupling_strength=coupling_strength,
                        coupling_skew=coupling_skew,
                    )
                fr_kw = fr_loss_kwargs_from_batch(
                    preds, data_points,
                    fr_mode_weight=fr_mode_weight,
                    fr_spectral_weight=fr_spectral_weight,
                    fr_invariant_weight=fr_invariant_weight,
                    fr_invariant_modal=fr_invariant_modal,
                    target_mode_amps=target_mode_amps,
                    target_spectrum=target_spectrum,
                )
                loss = total_loss(
                    preds, data_points, damping_rates, coupling_strength, inharm_b, speed_scalars,
                    prior_targets=prior_targets,
                    synth_waveform=synth_waveform,
                    target_waveform=target_waveform,
                    stft_weight=stft_weight,
                    **_invariant_loss_kwargs(
                        model, reference_coupling_skew, fr_invariant_weight,
                        fr_invariant_coupling, fr_invariant_speed, fr_invariant_inharm, fr_invariant_modal,
                        fr_replace_mse_priors,
                        reference_speed_scalars, reference_inharm_b,
                    ),
                    **fr_kw,
                )

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()

            current_loss = loss.item()
            global_step += 1

            if current_loss < best_loss - 1e-6:
                best_loss = current_loss
                stagnation_steps = 0
            else:
                stagnation_steps += 1

            main_progress.update(
                main_task,
                advance=1,
                loss=f"{current_loss:.6f}",
                jumps=jumps_performed,
                best=f"{best_loss:.6f}",
            )

            if global_step % 5000 == 0:
                save_trajectory_frame(global_step, preds.detach(), current_loss, seed,
                                      project_to_3d, time_norm, true_proj, data_proj, is_jump=False)

            # === Punctuated jump on stagnation ===
            if (stagnation_steps >= stagnation_patience and
                global_step >= min_step_for_jump and
                jumps_performed < max_jumps):

                print(f"\n>>> PUNCTUATED JUMP {jumps_performed + 1} at step {global_step} | loss {current_loss:.8f} <<<")

                pre_jump_loss = current_loss

                # CPU transfer for safe multiprocessing pickling
                pre_state_dict_cpu = {k: v.cpu().clone().detach() if torch.is_tensor(v) else v
                                      for k, v in model.state_dict().items()}
                data_points_cpu = data_points.cpu()
                times_cpu = times.cpu()
                initial_basis_cpu = initial_basis.cpu()
                prior_targets_cpu = prior_targets

                current_jump_std = jump_std_min + (jump_std_max - jump_std_min) * (jumps_performed / max(1, max_jumps - 1))

                args = []
                for i in range(pop_size):
                    std = 0.0 if INCLUDE_ZERO_JUMP and i == 0 else random.uniform(jump_std_min, current_jump_std)
                    worker_seed = random.randint(0, 2**32 - 1)
                    ref_skew_cpu = (
                        reference_coupling_skew.cpu() if reference_coupling_skew is not None else None
                    )
                    ref_speed_cpu = (
                        reference_speed_scalars.cpu() if reference_speed_scalars is not None else None
                    )
                    ref_inharm_cpu = (
                        reference_inharm_b.cpu() if reference_inharm_b is not None else None
                    )
                    arg = (i, pre_state_dict_cpu, std, data_points_cpu, times_cpu,
                           initial_basis_cpu, worker_seed, lr_geo, lr_slow, prior_targets_cpu,
                           rollout_horizon, fr_invariant_weight, fr_invariant_coupling,
                           fr_invariant_speed, fr_invariant_inharm, fr_invariant_modal,
                           fr_replace_mse_priors,
                           fr_mode_weight, fr_spectral_weight,
                           target_mode_amps.cpu(), target_spectrum.cpu(), ref_skew_cpu,
                           ref_speed_cpu, ref_inharm_cpu)
                    args.append(arg)

                candidates = []
                rollout_progress = Progress(
                    TextColumn("[bold yellow]{task.description}"),
                    BarColumn(),
                    MofNCompleteColumn(),
                    TextColumn("{task.percentage:>3.0f}%"),
                    TimeRemainingColumn(),
                )

                if use_parallel_rollouts:
                    print(f"   Parallel rollouts with {min(PARALLEL_MAX_WORKERS, pop_size)} workers...")
                    try:
                        with ProcessPoolExecutor(max_workers=min(PARALLEL_MAX_WORKERS, pop_size)) as executor:
                            futures = [executor.submit(_rollout_worker, a) for a in args]
                            with rollout_progress:
                                rollout_task = rollout_progress.add_task("Parallel rollouts", total=pop_size)
                                for future in as_completed(futures):
                                    try:
                                        res = future.result()
                                        candidates.append(res)
                                        print(f"   Candidate {res['idx']:2d} | std {res['noise_std']:.3f} | steps {res['steps']:5d} | loss {res['final_loss']:.8f}")
                                    except Exception as e:
                                        print(f"   Worker failed: {e}")
                                    rollout_progress.advance(rollout_task)
                    except Exception as e:
                        print(f"   Parallel failed ({e}), falling back to sequential")
                        with rollout_progress:
                            rollout_task = rollout_progress.add_task("Sequential rollouts", total=pop_size)
                            for a in args:
                                try:
                                    res = _rollout_worker(a)
                                    candidates.append(res)
                                    print(f"   Candidate {res['idx']:2d} | std {res['noise_std']:.3f} | steps {res['steps']:5d} | loss {res['final_loss']:.8f}")
                                except Exception as e:
                                    print(f"   Sequential worker failed: {e}")
                                rollout_progress.advance(rollout_task)
                else:
                    print("   Sequential rollouts...")
                    with rollout_progress:
                        rollout_task = rollout_progress.add_task("Sequential rollouts", total=pop_size)
                        for a in args:
                            try:
                                res = _rollout_worker(a)
                                candidates.append(res)
                                print(f"   Candidate {res['idx']:2d} | std {res['noise_std']:.3f} | steps {res['steps']:5d} | loss {res['final_loss']:.8f}")
                            except Exception as e:
                                print(f"   Sequential worker failed: {e}")
                            rollout_progress.advance(rollout_task)

                # Summary (sorted by idx)
                for c in sorted(candidates, key=lambda x: x['idx']):
                    print(f"   Candidate {c['idx']:2d} | std {c['noise_std']:.3f} | steps {c['steps']:5d} | loss {c['final_loss']:.8f}")

                # === Selection logic ===
                if USE_SWR_SELECTION:
                    scores = [loss_to_gamma(c['final_loss']) for c in candidates]
                else:
                    scores = [-c['final_loss'] for c in candidates]

                temp = max(ENTROPIC_MIN_TEMP, ENTROPIC_START_TEMP * (ENTROPIC_DECAY ** jumps_performed))
                probs = torch.softmax(torch.tensor(scores) / temp, dim=0).numpy()
                entropic_idx = np.random.choice(len(candidates), p=probs)
                entropic_candidate = candidates[entropic_idx]

                best_candidate = min(candidates, key=lambda c: c['final_loss'])

                selected_candidate = None
                if best_candidate['final_loss'] < pre_jump_loss - 1e-6:
                    selected_candidate = best_candidate
                    print(f"   Selected best candidate {best_candidate['idx']} (std {best_candidate['noise_std']:.3f}) | loss {best_candidate['final_loss']:.8f} (improvement)")
                elif entropic_candidate['final_loss'] < pre_jump_loss - 1e-6:
                    selected_candidate = entropic_candidate
                    print(f"   Selected entropic candidate {entropic_idx} (std {entropic_candidate['noise_std']:.3f}) | loss {entropic_candidate['final_loss']:.8f} (improvement)")
                else:
                    print("   No improvement found → continuing from current state")

                if selected_candidate is not None:
                    model.load_state_dict(selected_candidate['best_state'])
                    jumps_performed += 1

                # Robust post-jump visualization
                with torch.no_grad():
                    preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq = model(times)
                    fr_kw = fr_loss_kwargs_from_batch(
                        preds, data_points,
                        fr_mode_weight=fr_mode_weight,
                        fr_spectral_weight=fr_spectral_weight,
                        fr_invariant_weight=fr_invariant_weight,
                        fr_invariant_modal=fr_invariant_modal,
                        target_mode_amps=target_mode_amps,
                        target_spectrum=target_spectrum,
                    )
                    current_loss = total_loss(
                        preds, data_points, damping_rates, coupling_strength, inharm_b, speed_scalars,
                        prior_targets=prior_targets,
                        **_invariant_loss_kwargs(
                            model, reference_coupling_skew, fr_invariant_weight,
                            fr_invariant_coupling, fr_invariant_speed, fr_invariant_inharm, fr_invariant_modal,
                            fr_replace_mse_priors,
                            reference_speed_scalars, reference_inharm_b,
                        ),
                        **fr_kw,
                    ).item()

                save_trajectory_frame(global_step, preds.detach(), current_loss, seed,
                                      project_to_3d, time_norm, true_proj, data_proj, is_jump=(selected_candidate is not None))

                stagnation_steps = 0
                best_loss = current_loss

                # Update main progress with new loss/best after jump
                main_progress.update(
                    main_task,
                    loss=f"{current_loss:.6f}",
                    jumps=jumps_performed,
                    best=f"{best_loss:.6f}",
                )

    wall_time = time.time() - start_time

    # === Final evaluation ===
    with torch.no_grad():
        preds, damping_rates, coupling_strength, inharm_b, speed_scalars, full_freq = model(times)

    learned_vel_dir = manifold.proju(model.base, model.vel_dir_raw)
    learned_vel_dir = learned_vel_dir / (learned_vel_dir.norm(dim=0, keepdim=True) + 1e-8)

    coupling_val = coupling_strength.item()
    speed_rel_std = (speed_scalars.std(unbiased=False) / (speed_scalars.mean() + 1e-8)).item()
    max_inharm_b = inharm_b.max().item()
    per_mode_mse = (preds - exact_points).pow(2).mean(dim=[0, 1])
    true_var_per_mode = exact_points.var(dim=0).mean(dim=0).cpu().numpy()
    relative_mse_pred = (per_mode_mse / (per_mode_mse.detach().new_tensor(true_var_per_mode) + 1e-8)).cpu().numpy()
    total_recon_mse_pred = per_mode_mse.mean().item()

    if use_real_audio:
        target_damping = prior_targets['damping_rates'] if prior_targets else TRUE_DAMPING_RATES
        target_coupling = prior_targets.get('coupling_strength', TRUE_COUPLING_STRENGTH) if prior_targets else TRUE_COUPLING_STRENGTH
        true_geo_dist = 0.0
        coupling_err = abs(coupling_val - target_coupling)
        damping_rmse = torch.sqrt(F.mse_loss(damping_rates, target_damping)).item()
        damping_corr = float(np.corrcoef(damping_rates.cpu().numpy(), target_damping.cpu().numpy())[0, 1])
        freq_rmse = 0.0
        freq_corr = 1.0
        aligned_rates = damping_rates
        aligned_speed_scalars = speed_scalars
        aligned_learned_freq = full_freq
        aligned_inharm_b = inharm_b
        true_full_freq_np = full_freq.cpu().numpy()
        strict_success = damping_rmse < STRICT_DAMPING_RMSE and total_recon_mse_pred < 0.5
        loose_success = strict_success or total_recon_mse_pred < 1.0
    else:
        aligned_rates, aligned_speed_scalars, aligned_learned_freq, aligned_inharm_b = align_and_compute_freq(
            true_vel_dir, learned_vel_dir, damping_rates, speed_scalars, inharm_b, full_freq
        )
        true_geo_dist = stiefel_dist(model.base, true_base).item()
        coupling_err = abs(coupling_val - TRUE_COUPLING_STRENGTH)
        damping_rmse = torch.sqrt(F.mse_loss(aligned_rates, TRUE_DAMPING_RATES)).item()
        damping_corr = np.corrcoef(aligned_rates.cpu().numpy(), TRUE_DAMPING_RATES.cpu().numpy())[0, 1]
        freq_rmse = torch.sqrt(F.mse_loss(aligned_learned_freq, VELOCITY_SCALE_BASE * true_freq)).item()
        freq_corr = np.corrcoef(aligned_learned_freq.cpu().numpy(),
                                (VELOCITY_SCALE_BASE * true_freq).cpu().numpy())[0, 1]
        max_inharm_b = aligned_inharm_b.max().item()
        true_full_freq_np = (VELOCITY_SCALE_BASE * true_freq).cpu().numpy()
        strict_success = (true_geo_dist < STRICT_GEO_DIST and
                          max_inharm_b < STRICT_MAX_INHARM and
                          speed_rel_std < STRICT_SPEED_STD and
                          coupling_err < STRICT_COUPLING_ERR and
                          damping_rmse < STRICT_DAMPING_RMSE and
                          damping_corr > STRICT_DAMPING_CORR and
                          freq_corr > STRICT_FREQ_CORR and
                          freq_rmse < STRICT_FREQ_RMSE)
        loose_success = strict_success or (max_inharm_b < LOOSE_MAX_INHARM and speed_rel_std < LOOSE_SPEED_STD)

    save_detailed_pyramid_plot(
        seed=seed,
        status="FINAL",
        true_geo_dist=true_geo_dist,
        total_recon_mse_pred=total_recon_mse_pred,
        coupling_val=coupling_val,
        coupling_err=coupling_err,
        wall_time=wall_time,
        damping_rmse=damping_rmse,
        damping_corr=damping_corr,
        speed_rel_std=speed_rel_std,
        freq_rmse=freq_rmse,
        freq_corr=freq_corr,
        max_inharm_b=max_inharm_b,
        aligned_rates=aligned_rates,
        aligned_speed_scalars=aligned_speed_scalars,
        aligned_learned_freq=aligned_learned_freq,
        aligned_inharm_b=aligned_inharm_b,
        relative_mse_pred=relative_mse_pred,
        true_var_per_mode=true_var_per_mode,
        true_full_freq_np=true_full_freq_np,
        strict_success=strict_success,
        loose_success=loose_success,
    )

    result = {
        'seed': seed,
        'jumps': jumps_performed,
        'final_geo_dist': true_geo_dist,
        'max_inharm_b': max_inharm_b,
        'speed_rel_std': speed_rel_std,
        'coupling_err': coupling_err,
        'damping_rmse': damping_rmse,
        'damping_corr': damping_corr,
        'freq_rmse': freq_rmse,
        'freq_corr': freq_corr,
        'strict_success': strict_success,
        'loose_success': loose_success,
        'wall_time': wall_time,
        'total_recon_mse_pred': total_recon_mse_pred,
        'model': model,
        'full_freq': full_freq.detach(),
        'damping_rates': damping_rates.detach(),
        'coupling_strength': coupling_val,
        'inharm_b': inharm_b.detach(),
        'coupling_skew': extract_coupling_skew(model).detach(),
    }
    return result