train.py

from __future__ import annotations

import argparse
import json
import math
import os
import random
import secrets
import shutil
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn.functional as F
from accelerate import Accelerator
from huggingface_hub import HfApi
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import get_cosine_schedule_with_warmup

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from config import ensure_dir, flatten_for_wandb, load_config, save_json
from data import (
    OBS_ROLE_NONE,
    OBS_ROLE_USER,
    ThoughtLoopConversationDataset,
    build_chunk_batch,
    identity_collate,
    resolve_bucket_horizon_ticks,
)
from model import ThoughtLoopT5Gemma

try:
    import wandb
except ImportError:  # pragma: no cover
    wandb = None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train the recurrent T5Gemma thought-loop model.")
    parser.add_argument("--config", default="configs/sft_b200.yaml")
    parser.add_argument(
        "--resume-from-checkpoint",
        default=None,
        help="Local checkpoint path to resume from, or 'latest'/'auto' to use the newest checkpoint in the run dir.",
    )
    parser.add_argument(
        "--wandb-run-id",
        default=None,
        help="Optional W&B run id to resume. If omitted, train.py reuses a stored/local run id when available.",
    )
    return parser.parse_args()


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def build_optimizer(model: torch.nn.Module, training_cfg: dict[str, Any]) -> torch.optim.Optimizer:
    no_decay_terms = ("bias", "norm", "Norm")
    decay_params = []
    no_decay_params = []
    for name, parameter in model.named_parameters():
        target = no_decay_params if any(term in name for term in no_decay_terms) else decay_params
        target.append(parameter)

    fused = bool(training_cfg.get("fused_adamw", True) and torch.cuda.is_available())
    return torch.optim.AdamW(
        [
            {"params": decay_params, "weight_decay": float(training_cfg["weight_decay"])},
            {"params": no_decay_params, "weight_decay": 0.0},
        ],
        lr=float(training_cfg["learning_rate"]),
        betas=(float(training_cfg["adam_beta1"]), float(training_cfg["adam_beta2"])),
        eps=float(training_cfg["adam_epsilon"]),
        fused=fused,
    )


def resolve_bucket_tbptt_steps(training_cfg: dict[str, Any], duration_bucket: str) -> int | None:
    bucket_map = training_cfg.get("tbptt_steps_by_bucket", {})
    if duration_bucket in bucket_map:
        return int(bucket_map[duration_bucket])

    start_steps = int(training_cfg.get("tbptt_start_steps", 0))
    end_steps = int(training_cfg.get("tbptt_end_steps", start_steps))
    if start_steps <= 0 or end_steps <= 0:
        return None
    return end_steps


def resolve_linear_schedule_value(
    *,
    enabled: bool,
    global_step: int,
    total_update_steps: int,
    start_fraction: float,
    end_fraction: float,
    start_value: float,
    end_value: float,
) -> float | None:
    if not enabled:
        return None

    clamped_total_steps = max(total_update_steps, 1)
    start_step = int(round(clamped_total_steps * start_fraction))
    end_step = int(round(clamped_total_steps * end_fraction))
    if end_step <= start_step:
        end_step = start_step + 1

    if global_step <= start_step:
        return start_value
    if global_step >= end_step:
        return end_value

    progress = (global_step - start_step) / max(end_step - start_step, 1)
    return start_value + progress * (end_value - start_value)


def should_freeze_gate_head(
    training_cfg: dict[str, Any],
    global_step: int,
    total_update_steps: int,
) -> bool:
    if not bool(training_cfg.get("freeze_gate_head_first_half", False)):
        return False
    freeze_fraction = float(training_cfg.get("gate_head_freeze_fraction", 0.5))
    return global_step < int(round(max(total_update_steps, 1) * freeze_fraction))


def resolve_decoder_trainable_fraction(
    training_cfg: dict[str, Any],
    global_step: int,
    total_update_steps: int,
) -> float | None:
    return resolve_linear_schedule_value(
        enabled=bool(training_cfg.get("gradual_unfreeze_decoder", False)),
        global_step=global_step,
        total_update_steps=total_update_steps,
        start_fraction=float(training_cfg.get("decoder_unfreeze_start_fraction", 0.0)),
        end_fraction=float(training_cfg.get("decoder_unfreeze_end_fraction", 0.5)),
        start_value=float(training_cfg.get("decoder_initial_trainable_fraction", 0.0)),
        end_value=1.0,
    )


def resolve_batch_duration_bucket(conversations: list[dict[str, Any]]) -> str:
    if not conversations:
        return "short"
    buckets = {str(conversation["duration_bucket"]) for conversation in conversations}
    if len(buckets) == 1:
        return next(iter(buckets))
    return "mixed"


def ensure_single_chunk_dataset(dataset: ThoughtLoopConversationDataset, split: str) -> None:
    offending_examples = [
        example for example in dataset.examples if int(example.get("chunk_count", 0)) != 1
    ]
    if offending_examples:
        raise ValueError(
            f"{split} dataset expected exactly one chunk per conversation, but found "
            f"{len(offending_examples)} multi-chunk examples. Check chunk_ticks/max_horizon_ticks."
        )


def maybe_sort_dataset_by_length(
    dataset: ThoughtLoopConversationDataset,
    *,
    enabled: bool,
) -> None:
    if not enabled:
        return
    dataset.examples.sort(
        key=lambda item: (
            int(item.get("effective_total_ticks", item.get("total_ticks", 0))),
            int(item.get("total_ticks", 0)),
            str(item.get("row_id", "")),
        )
    )


def move_chunk_batch_to_device(batch: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]:
    return {key: value.to(device) for key, value in batch.items()}


def apply_runtime_model_config(model: ThoughtLoopT5Gemma, config: dict[str, Any]) -> None:
    model.config = config
    model_cfg = config["model"]
    expected_z_slots = int(model_cfg["z_slots"])
    if int(model.z_slots) != expected_z_slots:
        raise ValueError(
            f"Loaded model has z_slots={model.z_slots}, but config requested z_slots={expected_z_slots}."
        )
    model.thought_loop_proposal_mode = str(
        model_cfg.get("thought_loop_proposal_mode", model.thought_loop_proposal_mode)
    ).strip().lower()
    model.preserve_observation_encoder_manifold = bool(
        model_cfg.get(
            "preserve_observation_encoder_manifold",
            model.thought_loop_proposal_mode == "observation_hidden_compression",
        )
    )
    model.observation_encoder_use_state_context = bool(
        model_cfg.get("observation_encoder_use_state_context", False)
    )
    model.latent_attention_mask_mode = str(
        model_cfg.get("latent_attention_mask_mode", getattr(model, "latent_attention_mask_mode", "observed"))
    ).strip().lower()
    if model.latent_attention_mask_mode not in {"observed", "full"}:
        raise ValueError(f"Unsupported latent_attention_mask_mode: {model.latent_attention_mask_mode}")


def build_initial_model(config: dict[str, Any]) -> ThoughtLoopT5Gemma:
    initial_model_path = config["model"].get("initial_model_path")
    if initial_model_path:
        model = ThoughtLoopT5Gemma.from_pretrained(str(initial_model_path), device="cpu", map_location="cpu")
        apply_runtime_model_config(model, config)
        return model
    return ThoughtLoopT5Gemma(config)


def rollout_active_rows_only(
    *,
    model: ThoughtLoopT5Gemma,
    batch: dict[str, torch.Tensor],
    z_state: torch.Tensor,
    state_mask: torch.Tensor,
    step_index: int,
    active_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    active_indices = torch.nonzero(active_mask, as_tuple=False).flatten()
    active_next_z, active_gate_logits, active_next_mask = model.rollout_step_with_mask(
        z_state=z_state[active_indices],
        state_mask=state_mask[active_indices],
        observation_input_ids=batch["observation_input_ids"][active_indices, step_index, :],
        observation_attention_mask=batch["observation_attention_mask"][active_indices, step_index, :],
        observation_role_ids=batch["observation_role"][active_indices, step_index],
        delta_seconds=batch["delta_seconds"][active_indices, step_index],
        elapsed_seconds=batch["elapsed_seconds"][active_indices, step_index],
        since_last_user_seconds=torch.zeros_like(batch["elapsed_seconds"][active_indices, step_index]),
        since_last_agent_seconds=torch.zeros_like(batch["elapsed_seconds"][active_indices, step_index]),
    )

    next_z = z_state.clone()
    next_z[active_indices] = active_next_z
    next_state_mask = state_mask.clone()
    next_state_mask[active_indices] = active_next_mask
    gate_logits = torch.zeros(batch["tick_mask"].shape[0], device=z_state.device, dtype=active_gate_logits.dtype)
    gate_logits[active_indices] = active_gate_logits
    return next_z, next_state_mask, gate_logits


def compute_gate_positive_weight(batch: dict[str, torch.Tensor], training_cfg: dict[str, Any]) -> torch.Tensor:
    if not bool(training_cfg.get("use_dynamic_gate_positive_weight", True)):
        return torch.tensor(float(training_cfg["gate_positive_weight"]), device=batch["tick_mask"].device)

    active_targets = batch["gate_target"][batch["tick_mask"]]
    positive_count = active_targets.sum()
    negative_count = active_targets.numel() - positive_count
    if positive_count.item() <= 0.0:
        fallback = float(training_cfg["gate_positive_weight"])
        return torch.tensor(fallback, device=batch["tick_mask"].device)

    dynamic_weight = negative_count / positive_count.clamp_min(1.0)
    min_weight = float(training_cfg.get("dynamic_gate_positive_weight_min", 1.0))
    max_weight = float(training_cfg.get("dynamic_gate_positive_weight_max", 256.0))
    dynamic_weight = dynamic_weight.clamp(min=min_weight, max=max_weight)
    return dynamic_weight.to(device=batch["tick_mask"].device)


def compute_distance_aware_hazard_gate_loss(
    gate_logits: torch.Tensor,
    gate_targets: torch.Tensor,
    gate_loss_mask: torch.Tensor,
    training_cfg: dict[str, Any],
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    device = gate_logits.device
    batch_size, max_steps = gate_logits.shape
    support_window_ticks = int(training_cfg.get("hazard_soft_target_window_ticks", 16))
    target_tau_ticks = float(training_cfg.get("hazard_soft_target_tau_ticks", 4.0))
    censor_weight = float(training_cfg.get("hazard_censor_weight", 0.25))

    log_hazard = F.logsigmoid(gate_logits)
    log_survival = F.logsigmoid(-gate_logits)

    gate_loss_sum = torch.zeros((), device=device)
    gate_loss_count = torch.zeros((), device=device)
    hazard_event_count = torch.zeros((), device=device)
    hazard_censor_count = torch.zeros((), device=device)

    for batch_index in range(batch_size):
        active_len = int(gate_loss_mask[batch_index].sum().item())
        if active_len <= 0:
            continue

        active_targets = gate_targets[batch_index, :active_len]
        active_log_hazard = log_hazard[batch_index, :active_len]
        active_log_survival = log_survival[batch_index, :active_len]
        positive_indices = torch.nonzero(active_targets > 0.5, as_tuple=False).flatten().tolist()

        segment_start = 0
        for positive_index in positive_indices:
            segment_log_hazard = active_log_hazard[segment_start : positive_index + 1]
            segment_log_survival = active_log_survival[segment_start : positive_index + 1]
            if segment_log_hazard.numel() <= 0:
                segment_start = positive_index + 1
                continue

            prefix_log_survival = torch.cumsum(segment_log_survival, dim=0) - segment_log_survival
            segment_log_event = prefix_log_survival + segment_log_hazard

            if support_window_ticks > 0:
                support_start = max(0, segment_log_event.numel() - support_window_ticks)
            else:
                support_start = 0

            support_log_event = segment_log_event[support_start:]
            support_indices = torch.arange(
                support_start,
                segment_log_event.numel(),
                device=device,
                dtype=torch.float32,
            )
            distances = (segment_log_event.numel() - 1) - support_indices

            if target_tau_ticks > 0.0:
                target_weights = torch.exp(-distances / target_tau_ticks)
            else:
                target_weights = torch.zeros_like(distances)
                target_weights[-1] = 1.0
            target_weights = target_weights / target_weights.sum().clamp_min(1e-8)

            segment_loss = -(target_weights * support_log_event).sum()
            gate_loss_sum = gate_loss_sum + segment_loss
            gate_loss_count = gate_loss_count + 1.0
            hazard_event_count = hazard_event_count + 1.0
            segment_start = positive_index + 1

        if segment_start < active_len:
            tail_log_survival = active_log_survival[segment_start:active_len]
            if tail_log_survival.numel() > 0:
                censored_loss = -tail_log_survival.mean()
                gate_loss_sum = gate_loss_sum + censor_weight * censored_loss
                gate_loss_count = gate_loss_count + censor_weight
                hazard_censor_count = hazard_censor_count + 1.0

    gate_loss = gate_loss_sum / gate_loss_count.clamp_min(1.0)
    metrics = {
        "gate_loss_sum": gate_loss_sum.detach(),
        "gate_count": gate_loss_count.detach(),
        "hazard_event_count": hazard_event_count.detach(),
        "hazard_censor_count": hazard_censor_count.detach(),
    }
    return gate_loss, metrics


def build_temporally_soft_gate_targets(
    *,
    gate_targets: torch.Tensor,
    gate_loss_mask: torch.Tensor,
    observation_roles: torch.Tensor,
    training_cfg: dict[str, Any],
) -> torch.Tensor:
    early_window = max(0, int(training_cfg.get("gate_soft_target_early_window_ticks", 3)))
    late_window = max(0, int(training_cfg.get("gate_soft_target_late_window_ticks", 3)))
    tau_ticks = max(1e-6, float(training_cfg.get("gate_soft_target_tau_ticks", 1.5)))
    peak = float(training_cfg.get("gate_soft_target_peak", 1.0))
    floor = max(0.0, float(training_cfg.get("gate_soft_target_floor", 0.0)))
    shape = str(training_cfg.get("gate_soft_target_shape", "gaussian")).strip().lower()
    block_user_ticks = bool(training_cfg.get("gate_soft_target_block_user_ticks", True))

    soft_targets = torch.zeros_like(gate_targets)
    candidate_mask = gate_loss_mask > 0
    if block_user_ticks:
        candidate_mask = candidate_mask & (observation_roles != OBS_ROLE_USER)

    batch_size, _ = gate_targets.shape
    for batch_index in range(batch_size):
        active_len = int(gate_loss_mask[batch_index].sum().item())
        if active_len <= 0:
            continue

        positive_indices = torch.nonzero(
            gate_targets[batch_index, :active_len] > 0.5,
            as_tuple=False,
        ).flatten().tolist()
        for positive_index in positive_indices:
            start_index = max(0, positive_index - early_window)
            end_index = min(active_len, positive_index + late_window + 1)
            if end_index <= start_index:
                continue

            indices = torch.arange(start_index, end_index, device=gate_targets.device)
            distances = (indices - positive_index).abs().to(dtype=gate_targets.dtype)
            if shape == "exponential":
                values = torch.exp(-distances / tau_ticks)
            elif shape == "linear":
                side_windows = torch.where(
                    indices <= positive_index,
                    torch.full_like(indices, max(early_window, 1)),
                    torch.full_like(indices, max(late_window, 1)),
                ).to(dtype=gate_targets.dtype)
                values = 1.0 - distances / (side_windows + 1.0)
            else:
                values = torch.exp(-(distances.pow(2)) / (2.0 * tau_ticks * tau_ticks))

            values = (values * peak).clamp(0.0, 1.0)
            if floor > 0.0:
                values = torch.where(values > 0.0, values.clamp_min(floor), values)

            valid = candidate_mask[batch_index, indices]
            if not torch.any(valid):
                continue

            valid_indices = indices[valid]
            soft_targets[batch_index, valid_indices] = torch.maximum(
                soft_targets[batch_index, valid_indices],
                values[valid],
            )

    return torch.maximum(soft_targets, gate_targets).clamp(0.0, 1.0)


def compute_soft_bce_gate_loss(
    *,
    gate_logits: torch.Tensor,
    hard_gate_targets: torch.Tensor,
    gate_loss_mask: torch.Tensor,
    observation_roles: torch.Tensor,
    positive_weight: torch.Tensor,
    training_cfg: dict[str, Any],
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    soft_targets = build_temporally_soft_gate_targets(
        gate_targets=hard_gate_targets,
        gate_loss_mask=gate_loss_mask,
        observation_roles=observation_roles,
        training_cfg=training_cfg,
    )
    loss_raw = F.binary_cross_entropy_with_logits(gate_logits, soft_targets, reduction="none")
    soft_weight = 1.0 + (positive_weight.to(device=gate_logits.device) - 1.0) * soft_targets
    masked_loss = loss_raw * soft_weight * gate_loss_mask
    gate_loss_sum = masked_loss.sum()
    gate_count = gate_loss_mask.sum()
    gate_loss = gate_loss_sum / gate_count.clamp_min(1.0)
    metrics = {
        "gate_loss_sum": gate_loss_sum.detach(),
        "gate_count": gate_count.detach(),
        "soft_target_mass": (soft_targets * gate_loss_mask).sum().detach(),
        "soft_target_positive_count": ((soft_targets >= 0.5) & (gate_loss_mask > 0)).float().sum().detach(),
    }
    return gate_loss, metrics


def compute_gate_ranking_loss(
    gate_logits: torch.Tensor,
    gate_targets: torch.Tensor,
    gate_loss_mask: torch.Tensor,
    ranking_margin: float,
    ranking_lookback: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    device = gate_logits.device
    ranking_loss_sum = torch.zeros((), device=device)
    ranking_pair_count = torch.zeros((), device=device)

    batch_size, max_steps = gate_logits.shape
    for batch_index in range(batch_size):
        for step_index in range(max_steps):
            if float(gate_loss_mask[batch_index, step_index]) <= 0.0:
                continue
            if float(gate_targets[batch_index, step_index]) <= 0.5:
                continue

            positive_logit = gate_logits[batch_index, step_index]
            search_start = max(0, step_index - ranking_lookback)
            for previous_index in range(search_start, step_index):
                if float(gate_loss_mask[batch_index, previous_index]) <= 0.0:
                    continue
                if float(gate_targets[batch_index, previous_index]) >= 0.5:
                    continue

                negative_logit = gate_logits[batch_index, previous_index]
                ranking_loss_sum = ranking_loss_sum + F.relu(ranking_margin - positive_logit + negative_logit)
                ranking_pair_count = ranking_pair_count + 1

    return ranking_loss_sum, ranking_pair_count


def compute_chunk_losses(
    model: ThoughtLoopT5Gemma,
    batch: dict[str, torch.Tensor],
    training_cfg: dict[str, Any],
    z_state: torch.Tensor,
    state_mask: torch.Tensor,
    tbptt_steps: int | None = None,
    freeze_gate_head_to_targets: bool = False,
) -> tuple[torch.Tensor, dict[str, torch.Tensor], torch.Tensor, torch.Tensor]:
    device = batch["tick_mask"].device
    batch_size, max_steps = batch["tick_mask"].shape
    gate_loss_mode = str(training_cfg.get("gate_loss_mode", "hazard")).strip().lower()
    use_hazard_gate_loss = gate_loss_mode == "hazard"
    use_soft_bce_gate_loss = gate_loss_mode == "soft_bce"
    pos_weight = (
        torch.zeros((), device=device)
        if use_hazard_gate_loss
        else compute_gate_positive_weight(batch, training_cfg)
    )

    gate_loss_sum = torch.zeros((), device=device)
    gate_count = torch.zeros((), device=device)
    decoder_loss_sum = torch.zeros((), device=device)
    decoder_token_count = torch.zeros((), device=device)
    state_penalty_sum = torch.zeros((), device=device)
    state_penalty_count = torch.zeros((), device=device)

    tp = torch.zeros((), device=device)
    fp = torch.zeros((), device=device)
    fn = torch.zeros((), device=device)
    correct = torch.zeros((), device=device)
    counted = torch.zeros((), device=device)
    target_positive_count = torch.zeros((), device=device)
    predicted_positive_count = torch.zeros((), device=device)

    state_delta_weight = float(training_cfg["state_delta_weight"])
    steps_since_detach = 0
    gate_logits_history: list[torch.Tensor] = []
    gate_target_history: list[torch.Tensor] = []
    gate_loss_mask_history: list[torch.Tensor] = []
    observation_role_history: list[torch.Tensor] = []
    hazard_event_count = torch.zeros((), device=device)
    hazard_censor_count = torch.zeros((), device=device)
    soft_target_mass = torch.zeros((), device=device)
    soft_target_positive_count = torch.zeros((), device=device)
    teacher_logit_scale = float(training_cfg.get("gate_teacher_logit_scale", 20.0))

    for step_index in range(max_steps):
        active_mask = batch["tick_mask"][:, step_index]
        if not torch.any(active_mask):
            break

        previous_z = z_state
        next_z, next_state_mask, gate_logits = rollout_active_rows_only(
            model=model,
            batch=batch,
            z_state=z_state,
            state_mask=state_mask,
            step_index=step_index,
            active_mask=active_mask,
        )
        z_state = next_z
        state_mask = next_state_mask

        step_loss_mask = active_mask.float()
        step_targets = batch["gate_target"][:, step_index]
        if freeze_gate_head_to_targets:
            gate_logits = torch.where(
                step_targets > 0.5,
                torch.full_like(step_targets, teacher_logit_scale),
                torch.full_like(step_targets, -teacher_logit_scale),
            )

        gate_logits_history.append(gate_logits)
        gate_target_history.append(step_targets)
        gate_loss_mask_history.append(step_loss_mask)
        observation_role_history.append(batch["observation_role"][:, step_index])

        if torch.any(step_loss_mask > 0):
            step_predictions = (torch.sigmoid(gate_logits) >= 0.5).float()
            valid = step_loss_mask > 0
            correct = correct + ((step_predictions == step_targets) & valid).float().sum()
            counted = counted + valid.float().sum()
            tp = tp + ((step_predictions == 1) & (step_targets == 1) & valid).float().sum()
            fp = fp + ((step_predictions == 1) & (step_targets == 0) & valid).float().sum()
            fn = fn + ((step_predictions == 0) & (step_targets == 1) & valid).float().sum()
            target_positive_count = target_positive_count + ((step_targets == 1) & valid).float().sum()
            predicted_positive_count = predicted_positive_count + ((step_predictions == 1) & valid).float().sum()

        if (
            torch.any(step_loss_mask > 0)
            and not use_hazard_gate_loss
            and not use_soft_bce_gate_loss
            and not freeze_gate_head_to_targets
        ):
            gate_loss_raw = F.binary_cross_entropy_with_logits(gate_logits, step_targets, reduction="none")
            gate_weight = torch.where(step_targets > 0.5, pos_weight, torch.ones_like(step_targets))
            gate_loss_sum = gate_loss_sum + (gate_loss_raw * gate_weight * step_loss_mask).sum()
            gate_count = gate_count + step_loss_mask.sum()

            if state_delta_weight > 0.0:
                state_delta = (next_z - previous_z).pow(2).mean(dim=(1, 2))
                state_penalty_sum = state_penalty_sum + (state_delta * step_loss_mask).sum()
                state_penalty_count = state_penalty_count + step_loss_mask.sum()
        elif torch.any(step_loss_mask > 0) and state_delta_weight > 0.0:
            state_delta = (next_z - previous_z).pow(2).mean(dim=(1, 2))
            state_penalty_sum = state_penalty_sum + (state_delta * step_loss_mask).sum()
            state_penalty_count = state_penalty_count + step_loss_mask.sum()

        step_labels = batch["decoder_labels"][:, step_index, :]
        decoder_rows = active_mask & torch.any(step_labels != -100, dim=-1)
        if torch.any(decoder_rows):
            step_loss_sum, step_token_count = model.decoder_loss(
                z_state[decoder_rows],
                step_labels[decoder_rows],
                encoder_attention_mask=state_mask[decoder_rows],
            )
            decoder_loss_sum = decoder_loss_sum + step_loss_sum
            decoder_token_count = decoder_token_count + step_token_count

        if tbptt_steps is not None:
            steps_since_detach += 1
            if steps_since_detach >= tbptt_steps:
                z_state = z_state.detach()
                state_mask = state_mask.detach()
                steps_since_detach = 0

    if gate_logits_history:
        stacked_gate_logits = torch.stack(gate_logits_history, dim=1)
        stacked_gate_targets = torch.stack(gate_target_history, dim=1)
        stacked_gate_loss_mask = torch.stack(gate_loss_mask_history, dim=1)
        stacked_observation_roles = torch.stack(observation_role_history, dim=1)
    else:
        stacked_gate_logits = torch.zeros((batch_size, 0), device=device)
        stacked_gate_targets = torch.zeros((batch_size, 0), device=device)
        stacked_gate_loss_mask = torch.zeros((batch_size, 0), device=device)
        stacked_observation_roles = torch.zeros((batch_size, 0), device=device, dtype=torch.long)

    gate_ranking_loss = torch.zeros((), device=device)
    gate_ranking_loss_sum = torch.zeros((), device=device)
    gate_ranking_pair_count = torch.zeros((), device=device)
    if freeze_gate_head_to_targets:
        gate_loss = torch.zeros((), device=device)
        gate_loss_sum = torch.zeros((), device=device)
        gate_count = torch.zeros((), device=device)
        hazard_event_count = torch.zeros((), device=device)
        hazard_censor_count = torch.zeros((), device=device)
    elif use_hazard_gate_loss and gate_logits_history:
        gate_loss, hazard_metrics = compute_distance_aware_hazard_gate_loss(
            gate_logits=stacked_gate_logits,
            gate_targets=stacked_gate_targets,
            gate_loss_mask=stacked_gate_loss_mask,
            training_cfg=training_cfg,
        )
        gate_loss_sum = hazard_metrics["gate_loss_sum"]
        gate_count = hazard_metrics["gate_count"]
        hazard_event_count = hazard_metrics["hazard_event_count"]
        hazard_censor_count = hazard_metrics["hazard_censor_count"]
    elif use_soft_bce_gate_loss and gate_logits_history:
        gate_loss, soft_bce_metrics = compute_soft_bce_gate_loss(
            gate_logits=stacked_gate_logits,
            hard_gate_targets=stacked_gate_targets,
            gate_loss_mask=stacked_gate_loss_mask,
            observation_roles=stacked_observation_roles,
            positive_weight=pos_weight,
            training_cfg=training_cfg,
        )
        gate_loss_sum = soft_bce_metrics["gate_loss_sum"]
        gate_count = soft_bce_metrics["gate_count"]
        soft_target_mass = soft_bce_metrics["soft_target_mass"]
        soft_target_positive_count = soft_bce_metrics["soft_target_positive_count"]
    else:
        gate_loss = gate_loss_sum / gate_count.clamp_min(1.0)

    if (
        not use_hazard_gate_loss
        and bool(training_cfg.get("use_ranked_gate_objective", False))
        and gate_logits_history
        and not freeze_gate_head_to_targets
    ):
        gate_ranking_loss_sum, gate_ranking_pair_count = compute_gate_ranking_loss(
            gate_logits=stacked_gate_logits,
            gate_targets=stacked_gate_targets,
            gate_loss_mask=stacked_gate_loss_mask,
            ranking_margin=float(training_cfg.get("gate_ranking_margin", 0.5)),
            ranking_lookback=int(training_cfg.get("gate_ranking_lookback", 64)),
        )
        gate_ranking_loss = gate_ranking_loss_sum / gate_ranking_pair_count.clamp_min(1.0)
        gate_loss = (
            float(training_cfg.get("gate_bce_weight", 1.0)) * gate_loss
            + float(training_cfg.get("gate_ranking_weight", 1.0)) * gate_ranking_loss
        )

    decoder_loss = decoder_loss_sum / decoder_token_count.clamp_min(1.0)
    state_delta_loss = state_penalty_sum / state_penalty_count.clamp_min(1.0)
    total_loss = (
        float(training_cfg["gate_loss_weight"]) * gate_loss
        + float(training_cfg["decoder_loss_weight"]) * decoder_loss
        + state_delta_weight * state_delta_loss
    )

    metrics = {
        "loss_total": total_loss.detach(),
        "gate_loss_sum": gate_loss_sum.detach(),
        "gate_count": gate_count.detach(),
        "decoder_loss_sum": decoder_loss_sum.detach(),
        "decoder_token_count": decoder_token_count.detach(),
        "state_penalty_sum": state_penalty_sum.detach(),
        "state_penalty_count": state_penalty_count.detach(),
        "gate_tp": tp.detach(),
        "gate_fp": fp.detach(),
        "gate_fn": fn.detach(),
        "gate_correct": correct.detach(),
        "gate_eval_count": counted.detach(),
        "gate_target_positive_count": target_positive_count.detach(),
        "gate_pred_positive_count": predicted_positive_count.detach(),
        "gate_ranking_loss_sum": gate_ranking_loss_sum.detach(),
        "gate_ranking_pair_count": gate_ranking_pair_count.detach(),
        "hazard_event_count": hazard_event_count.detach(),
        "hazard_censor_count": hazard_censor_count.detach(),
        "soft_target_mass": soft_target_mass.detach(),
        "soft_target_positive_count": soft_target_positive_count.detach(),
        "active_ticks": batch["tick_mask"].float().sum().detach(),
        "gate_positive_weight": pos_weight.detach(),
        "metrics_count": torch.ones((), device=device),
    }
    return total_loss, metrics, z_state, state_mask


def reduce_metrics(
    accelerator: Accelerator,
    metric_list: list[dict[str, torch.Tensor]],
    training_cfg: dict[str, Any],
) -> dict[str, float]:
    totals: dict[str, torch.Tensor] = {}
    for metrics in metric_list:
        for key, value in metrics.items():
            totals[key] = totals.get(key, torch.zeros_like(value)) + value

    reduced = {key: accelerator.reduce(value, reduction="sum") for key, value in totals.items()}

    gate_loss_mode = str(training_cfg.get("gate_loss_mode", "hazard")).strip().lower()
    gate_loss = reduced["gate_loss_sum"] / reduced["gate_count"].clamp_min(1.0)
    gate_ranking_loss = reduced["gate_ranking_loss_sum"] / reduced["gate_ranking_pair_count"].clamp_min(1.0)
    decoder_loss = reduced["decoder_loss_sum"] / reduced["decoder_token_count"].clamp_min(1.0)
    state_delta_loss = reduced["state_penalty_sum"] / reduced["state_penalty_count"].clamp_min(1.0)
    gate_precision = reduced["gate_tp"] / (reduced["gate_tp"] + reduced["gate_fp"]).clamp_min(1.0)
    gate_recall = reduced["gate_tp"] / (reduced["gate_tp"] + reduced["gate_fn"]).clamp_min(1.0)
    gate_f1 = 2 * gate_precision * gate_recall / (gate_precision + gate_recall).clamp_min(1e-8)
    if gate_loss_mode != "hazard" and bool(training_cfg.get("use_ranked_gate_objective", False)):
        gate_loss = (
            float(training_cfg.get("gate_bce_weight", 1.0)) * gate_loss
            + float(training_cfg.get("gate_ranking_weight", 1.0)) * gate_ranking_loss
        )
    total_loss = (
        float(training_cfg["gate_loss_weight"]) * gate_loss
        + float(training_cfg["decoder_loss_weight"]) * decoder_loss
        + float(training_cfg["state_delta_weight"]) * state_delta_loss
    )

    return {
        "loss_total": float(total_loss),
        "loss_gate": float(gate_loss),
        "loss_gate_ranking": float(gate_ranking_loss),
        "loss_decoder": float(decoder_loss),
        "loss_state_delta": float(state_delta_loss),
        "gate_accuracy": float(reduced["gate_correct"] / reduced["gate_eval_count"].clamp_min(1.0)),
        "gate_precision": float(gate_precision),
        "gate_recall": float(gate_recall),
        "gate_f1": float(gate_f1),
        "gate_target_positive_count": float(reduced["gate_target_positive_count"]),
        "gate_pred_positive_count": float(reduced["gate_pred_positive_count"]),
        "gate_positive_weight": float(reduced["gate_positive_weight"] / reduced["metrics_count"].clamp_min(1.0)),
        "hazard_event_count": float(reduced["hazard_event_count"]),
        "hazard_censor_count": float(reduced["hazard_censor_count"]),
        "soft_target_mass": float(reduced["soft_target_mass"]),
        "soft_target_positive_count": float(reduced["soft_target_positive_count"]),
        "decoder_token_count": float(reduced["decoder_token_count"]),
        "active_ticks": float(reduced["active_ticks"]),
    }


@torch.no_grad()
def evaluate(
    accelerator: Accelerator,
    model: ThoughtLoopT5Gemma,
    dataloader: DataLoader,
    training_cfg: dict[str, Any],
    rollout_cfg: dict[str, Any],
    tokenizer_pad_token_id: int,
    max_batches: int | None = None,
) -> dict[str, float]:
    model.eval()
    all_metrics: list[dict[str, torch.Tensor]] = []
    chunk_ticks = int(rollout_cfg["chunk_ticks"])
    tick_seconds = float(rollout_cfg["tick_seconds"])

    for batch_index, conversations in enumerate(dataloader):
        if max_batches is not None and batch_index >= max_batches:
            break

        batch_size = len(conversations)
        z_state = model.initial_state(batch_size=batch_size, device=accelerator.device)
        state_mask = model.initial_state_mask(batch_size=batch_size, device=accelerator.device)
        max_chunks = max(int(conversation["chunk_count"]) for conversation in conversations)
        for chunk_index in range(max_chunks):
            chunk_batch = build_chunk_batch(
                conversations=conversations,
                chunk_index=chunk_index,
                pad_token_id=tokenizer_pad_token_id,
                chunk_ticks=chunk_ticks,
                tick_seconds=tick_seconds,
            )
            if chunk_batch is None:
                break

            chunk_batch = move_chunk_batch_to_device(chunk_batch, accelerator.device)
            _, metrics, z_state, state_mask = compute_chunk_losses(
                model=model,
                batch=chunk_batch,
                training_cfg=training_cfg,
                z_state=z_state,
                state_mask=state_mask,
                tbptt_steps=None,
                freeze_gate_head_to_targets=False,
            )
            z_state = z_state.detach()
            state_mask = state_mask.detach()
            all_metrics.append(metrics)

    if not all_metrics:
        model.train()
        return {}

    reduced = reduce_metrics(accelerator, all_metrics, training_cfg)
    model.train()
    return reduced


def copy_supporting_files(output_dir: Path, config_path: Path) -> None:
    for filename in ("config.py", "data.py", "model.py", "inference.py", "train.py"):
        shutil.copy2(SCRIPT_DIR / filename, output_dir / filename)
    shutil.copy2(config_path, output_dir / config_path.name)


def build_model_card(config: dict[str, Any], best_metrics: dict[str, float]) -> str:
    return "\n".join(
        [
            "# Samantha Thought-Loop SFT",
            "",
            f"- Base model: `{config['model']['base_model_name']}`",
            f"- Cleaned dataset repo: `{config['dataset'].get('cleaned_repo_id')}`",
            f"- Raw dataset repo: `{config['dataset'].get('raw_repo_id')}`",
            f"- Latent slots: `{config['model']['z_slots']}`",
            f"- Fixed tick seconds: `{config['rollout']['tick_seconds']}`",
            f"- Chunk ticks: `{config['rollout']['chunk_ticks']}`",
            f"- Max horizon ticks: `{config['rollout'].get('max_horizon_ticks', 36000)}`",
            f"- Explicit time features: `{config['model'].get('use_explicit_time_features', False)}`",
            f"- Gate loss mode: `{config['training'].get('gate_loss_mode', 'hazard')}`",
            f"- Hazard soft target window: `{config['training'].get('hazard_soft_target_window_ticks', 16)}`",
            f"- Hazard soft target tau ticks: `{config['training'].get('hazard_soft_target_tau_ticks', 4.0)}`",
            f"- Hazard censor weight: `{config['training'].get('hazard_censor_weight', 0.25)}`",
            f"- TBPTT by bucket: `{config['training'].get('tbptt_steps_by_bucket', {})}`",
            "",
            "## Best Validation Metrics",
            "",
            *[f"- `{key}`: `{value:.6f}`" for key, value in sorted(best_metrics.items())],
            "",
        ]
    )


def export_model(
    accelerator: Accelerator,
    model: ThoughtLoopT5Gemma,
    output_dir: Path,
    config: dict[str, Any],
    config_path: Path,
    best_metrics: dict[str, float],
) -> None:
    if not accelerator.is_main_process:
        return
    output_dir.mkdir(parents=True, exist_ok=True)
    unwrapped = accelerator.unwrap_model(model)
    unwrapped.save_pretrained(output_dir)
    copy_supporting_files(output_dir, config_path)
    (output_dir / "README.md").write_text(build_model_card(config, best_metrics), encoding="utf-8")
    save_json(output_dir / "best_metrics.json", best_metrics)


def maybe_push_to_hub(config: dict[str, Any], export_dir: Path) -> None:
    model_repo_id = config["hub"].get("model_repo_id")
    if not model_repo_id:
        return
    api = HfApi()
    api.create_repo(repo_id=model_repo_id, repo_type="model", exist_ok=True, private=bool(config["hub"].get("private")))
    api.upload_folder(folder_path=str(export_dir), repo_id=model_repo_id, repo_type="model")


def maybe_push_checkpoint_to_hub(
    accelerator: Accelerator,
    config: dict[str, Any],
    checkpoint_dir: Path,
) -> None:
    model_repo_id = config["hub"].get("model_repo_id")
    if not accelerator.is_main_process or not model_repo_id:
        return
    if not bool(config["hub"].get("push_checkpoints", True)):
        return

    run_name = str(config["wandb"].get("run_name", "run"))
    checkpoint_prefix = str(config["hub"].get("checkpoint_path_prefix", "checkpoints")).strip("/")
    path_parts = [part for part in (checkpoint_prefix, SCRIPT_DIR.name, run_name, checkpoint_dir.name) if part]
    path_in_repo = "/".join(path_parts)

    api = HfApi()
    api.create_repo(repo_id=model_repo_id, repo_type="model", exist_ok=True, private=bool(config["hub"].get("private")))
    api.upload_folder(
        folder_path=str(checkpoint_dir),
        repo_id=model_repo_id,
        repo_type="model",
        path_in_repo=path_in_repo,
        commit_message=f"Add {SCRIPT_DIR.name} {checkpoint_dir.name}",
    )


def read_json(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    with path.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise ValueError(f"Expected a JSON object in {path}.")
    return payload


def checkpoint_global_step(checkpoint_dir: Path) -> int | None:
    state = read_json(checkpoint_dir / "trainer_state.json")
    if "global_step" in state:
        return int(state["global_step"])

    prefix = "checkpoint-"
    if checkpoint_dir.name.startswith(prefix):
        suffix = checkpoint_dir.name[len(prefix) :]
        if suffix.isdigit():
            return int(suffix)
    return None


def find_latest_checkpoint(run_dir: Path) -> Path | None:
    checkpoints = []
    for candidate in run_dir.glob("checkpoint-*"):
        if not candidate.is_dir():
            continue
        step = checkpoint_global_step(candidate)
        if step is not None:
            checkpoints.append((step, candidate))
    if not checkpoints:
        return None
    return max(checkpoints, key=lambda item: item[0])[1]


def resolve_resume_checkpoint(requested: str | None, run_dir: Path) -> Path | None:
    if requested is None:
        return None

    value = str(requested).strip()
    if not value:
        return None

    if value.lower() in {"1", "true", "yes", "auto", "latest"}:
        checkpoint_dir = find_latest_checkpoint(run_dir)
        if checkpoint_dir is None:
            raise FileNotFoundError(f"No checkpoint-* directory found under {run_dir}.")
        return checkpoint_dir

    checkpoint_dir = Path(value).expanduser()
    if not checkpoint_dir.is_absolute():
        checkpoint_dir = SCRIPT_DIR / checkpoint_dir
    checkpoint_dir = checkpoint_dir.resolve()
    if not checkpoint_dir.is_dir():
        raise FileNotFoundError(f"Resume checkpoint does not exist: {checkpoint_dir}")
    return checkpoint_dir


def parse_wandb_run_id_from_dir(path: Path) -> str | None:
    name = path.resolve(strict=False).name if path.is_symlink() else path.name
    for prefix in ("offline-run-", "run-"):
        if name.startswith(prefix):
            parts = name.split("-")
            if parts:
                return parts[-1]
    return None


def discover_local_wandb_run_id(run_name: str) -> str | None:
    wandb_dir = SCRIPT_DIR / "wandb"
    if not wandb_dir.exists():
        return None

    candidates: list[Path] = []
    latest = wandb_dir / "latest-run"
    if latest.exists() or latest.is_symlink():
        candidates.append(latest)
    candidates.extend(sorted(wandb_dir.glob("*run-*"), key=lambda path: path.stat().st_mtime, reverse=True))

    for candidate in candidates:
        run_id = parse_wandb_run_id_from_dir(candidate)
        if not run_id:
            continue
        resolved = candidate.resolve(strict=False)
        config_path = resolved / "files" / "config.yaml"
        debug_log_path = resolved / "logs" / "debug.log"
        if config_path.exists() and run_name in config_path.read_text(encoding="utf-8", errors="ignore"):
            return run_id
        if debug_log_path.exists() and run_name in debug_log_path.read_text(encoding="utf-8", errors="ignore"):
            return run_id
    return None


def resolve_wandb_run_id(args: argparse.Namespace, config: dict[str, Any], run_dir: Path) -> str:
    run_id_path = run_dir / "wandb_run_id.txt"
    explicit = args.wandb_run_id or os.environ.get("WANDB_RUN_ID") or config["wandb"].get("run_id")
    if explicit:
        run_id = str(explicit).strip()
    elif run_id_path.exists():
        run_id = run_id_path.read_text(encoding="utf-8").strip()
    else:
        run_id = discover_local_wandb_run_id(str(config["wandb"].get("run_name", ""))) or secrets.token_hex(4)

    if not run_id:
        run_id = secrets.token_hex(4)
    run_id_path.write_text(run_id + "\n", encoding="utf-8")
    return run_id


def save_trainer_state(
    checkpoint_dir: Path,
    *,
    global_step: int,
    epoch: int,
    next_batch_index: int,
    best_validation_loss: float,
    best_metrics: dict[str, float],
    wandb_run_id: str | None,
) -> None:
    payload: dict[str, Any] = {
        "global_step": int(global_step),
        "epoch": int(epoch),
        "next_batch_index": int(next_batch_index),
        "best_validation_loss": best_validation_loss if math.isfinite(best_validation_loss) else None,
        "best_metrics": best_metrics,
        "wandb_run_id": wandb_run_id,
    }
    save_json(checkpoint_dir / "trainer_state.json", payload)


def load_resume_progress(
    checkpoint_dir: Path,
    *,
    total_chunk_microsteps: int,
    gradient_accumulation_steps: int,
    num_train_epochs: int,
) -> tuple[int, int, int, dict[str, Any]]:
    state = read_json(checkpoint_dir / "trainer_state.json")
    global_step = int(state.get("global_step") or checkpoint_global_step(checkpoint_dir) or 0)
    if "epoch" in state and "next_batch_index" in state:
        start_epoch = int(state["epoch"])
        batches_to_skip = int(state["next_batch_index"])
    else:
        updates_per_epoch = max(1, math.ceil(total_chunk_microsteps / max(gradient_accumulation_steps, 1)))
        start_epoch = min(global_step // updates_per_epoch, max(num_train_epochs - 1, 0))
        updates_into_epoch = global_step - start_epoch * updates_per_epoch
        batches_to_skip = updates_into_epoch * max(gradient_accumulation_steps, 1)

    while batches_to_skip >= total_chunk_microsteps and start_epoch < num_train_epochs:
        batches_to_skip -= total_chunk_microsteps
        start_epoch += 1

    return global_step, start_epoch, batches_to_skip, state


def load_best_metrics(export_dir: Path, resume_state: dict[str, Any]) -> tuple[float, dict[str, float]]:
    best_metrics = resume_state.get("best_metrics")
    if not isinstance(best_metrics, dict):
        best_metrics = read_json(export_dir / "best_metrics.json")

    loss = resume_state.get("best_validation_loss")
    if loss is None and isinstance(best_metrics, dict):
        loss = best_metrics.get("loss_total")
    best_validation_loss = float(loss) if loss is not None else float("inf")
    return best_validation_loss, {str(key): float(value) for key, value in best_metrics.items()}


def main() -> None:
    args = parse_args()
    config_path = Path(args.config)
    config = load_config(config_path)
    training_cfg = config["training"]
    rollout_cfg = config["rollout"]

    set_seed(int(training_cfg["seed"]))
    wandb_enabled = bool(config["wandb"].get("enabled", True) and wandb is not None)

    accelerator = Accelerator(
        gradient_accumulation_steps=int(training_cfg["gradient_accumulation_steps"]),
        mixed_precision=training_cfg.get("mixed_precision", "bf16"),
        log_with="wandb" if wandb_enabled else None,
        project_dir=config["paths"]["run_root"],
    )

    run_root = ensure_dir(config["paths"]["run_root"])
    export_root = ensure_dir(config["paths"]["export_root"])
    run_name = config["wandb"]["run_name"]
    run_dir = ensure_dir(run_root / run_name)
    export_dir = ensure_dir(export_root / run_name)
    resume_request = args.resume_from_checkpoint or training_cfg.get("resume_from_checkpoint")
    resume_checkpoint = resolve_resume_checkpoint(resume_request, run_dir)
    wandb_run_id = resolve_wandb_run_id(args, config, run_dir) if wandb_enabled else None

    model = build_initial_model(config)
    if model.use_explicit_time_features:
        raise ValueError("This trainer expects model.use_explicit_time_features=false for fixed-cadence training.")
    tokenizer = model.tokenizer

    train_dataset = ThoughtLoopConversationDataset(config=config, tokenizer=tokenizer, split="train")
    validation_dataset = ThoughtLoopConversationDataset(config=config, tokenizer=tokenizer, split="validation")
    sort_dataset_by_length = bool(training_cfg.get("sort_dataset_by_length", False))

    ensure_single_chunk_dataset(train_dataset, "train")
    ensure_single_chunk_dataset(validation_dataset, "validation")
    maybe_sort_dataset_by_length(train_dataset, enabled=sort_dataset_by_length)
    maybe_sort_dataset_by_length(validation_dataset, enabled=sort_dataset_by_length)

    train_loader = DataLoader(
        train_dataset,
        batch_size=int(training_cfg["micro_batch_size"]),
        shuffle=bool(training_cfg.get("shuffle_train", True)) and not sort_dataset_by_length,
        num_workers=int(training_cfg["num_workers"]),
        pin_memory=True,
        persistent_workers=bool(int(training_cfg["num_workers"]) > 0),
        collate_fn=identity_collate,
    )
    validation_loader = DataLoader(
        validation_dataset,
        batch_size=int(training_cfg["eval_batch_size"]),
        shuffle=False,
        num_workers=int(training_cfg["num_workers"]),
        pin_memory=True,
        persistent_workers=bool(int(training_cfg["num_workers"]) > 0),
        collate_fn=identity_collate,
    )

    optimizer = build_optimizer(model, training_cfg)
    total_chunk_microsteps = len(train_loader)
    total_update_steps = max(
        1,
        math.ceil(total_chunk_microsteps / max(int(training_cfg["gradient_accumulation_steps"]), 1))
        * int(training_cfg["num_train_epochs"]),
    )
    warmup_steps = max(1, int(total_update_steps * float(training_cfg["warmup_ratio"])))
    scheduler = get_cosine_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=warmup_steps,
        num_training_steps=max(1, total_update_steps),
    )

    model, optimizer, train_loader, validation_loader, scheduler = accelerator.prepare(
        model, optimizer, train_loader, validation_loader, scheduler
    )

    resume_state: dict[str, Any] = {}
    global_step = 0
    start_epoch = 0
    batches_to_skip = 0
    if resume_checkpoint is not None:
        accelerator.print(f"Loading local checkpoint from {resume_checkpoint}")
        accelerator.load_state(str(resume_checkpoint))
        global_step, start_epoch, batches_to_skip, resume_state = load_resume_progress(
            resume_checkpoint,
            total_chunk_microsteps=total_chunk_microsteps,
            gradient_accumulation_steps=int(training_cfg["gradient_accumulation_steps"]),
            num_train_epochs=int(training_cfg["num_train_epochs"]),
        )
        accelerator.print(
            f"Resuming at global_step={global_step}, epoch={start_epoch + 1}, "
            f"skipping {batches_to_skip} batch(es)."
        )

    if accelerator.is_main_process and wandb_enabled:
        accelerator.print(f"W&B run id: {wandb_run_id}")
        accelerator.init_trackers(
            project_name=config["wandb"]["project"],
            config=flatten_for_wandb(config),
            init_kwargs={
                "wandb": {
                    "name": run_name,
                    "id": wandb_run_id,
                    "resume": str(config["wandb"].get("resume", "allow")),
                }
            },
        )

    unwrapped_model = accelerator.unwrap_model(model)
    current_decoder_fraction = resolve_decoder_trainable_fraction(training_cfg, global_step, total_update_steps)
    if current_decoder_fraction is None:
        current_decoder_fraction = 1.0
    current_decoder_trainable_layers = unwrapped_model.set_decoder_trainable_fraction(current_decoder_fraction)

    current_gate_head_frozen = should_freeze_gate_head(training_cfg, global_step, total_update_steps)
    unwrapped_model.set_gate_trainable(not current_gate_head_frozen)

    progress_bar = tqdm(
        total=total_update_steps,
        initial=global_step,
        disable=not accelerator.is_local_main_process,
        desc="Training",
        dynamic_ncols=True,
    )

    best_validation_loss, best_metrics = load_best_metrics(export_dir, resume_state)
    model.train()

    chunk_ticks = int(rollout_cfg["chunk_ticks"])
    tick_seconds = float(rollout_cfg["tick_seconds"])

    for epoch in range(start_epoch, int(training_cfg["num_train_epochs"])):
        epoch_train_loader = train_loader
        batch_index_offset = 0
        if epoch == start_epoch and batches_to_skip > 0:
            epoch_train_loader = accelerator.skip_first_batches(train_loader, batches_to_skip)
            batch_index_offset = batches_to_skip

        for batch_index, conversations in enumerate(epoch_train_loader, start=batch_index_offset):
            batch_bucket = resolve_batch_duration_bucket(conversations)
            current_tbptt_steps = resolve_bucket_tbptt_steps(training_cfg, batch_bucket)
            current_horizon_ticks = resolve_bucket_horizon_ticks(batch_bucket, rollout_cfg)

            z_state = model.initial_state(batch_size=len(conversations), device=accelerator.device)
            state_mask = model.initial_state_mask(batch_size=len(conversations), device=accelerator.device)
            max_chunks = max(int(conversation["chunk_count"]) for conversation in conversations)

            for chunk_index in range(max_chunks):
                scheduled_decoder_fraction = resolve_decoder_trainable_fraction(training_cfg, global_step, total_update_steps)
                if scheduled_decoder_fraction is None:
                    scheduled_decoder_fraction = 1.0
                if abs(scheduled_decoder_fraction - current_decoder_fraction) > 1e-6:
                    current_decoder_fraction = scheduled_decoder_fraction
                    current_decoder_trainable_layers = unwrapped_model.set_decoder_trainable_fraction(
                        current_decoder_fraction
                    )

                scheduled_gate_head_frozen = should_freeze_gate_head(training_cfg, global_step, total_update_steps)
                if scheduled_gate_head_frozen != current_gate_head_frozen:
                    current_gate_head_frozen = scheduled_gate_head_frozen
                    unwrapped_model.set_gate_trainable(not current_gate_head_frozen)

                chunk_batch = build_chunk_batch(
                    conversations=conversations,
                    chunk_index=chunk_index,
                    pad_token_id=tokenizer.pad_token_id,
                    chunk_ticks=chunk_ticks,
                    tick_seconds=tick_seconds,
                )
                if chunk_batch is None:
                    break
                chunk_batch = move_chunk_batch_to_device(chunk_batch, accelerator.device)

                with accelerator.accumulate(model):
                    total_loss, metrics, z_state, state_mask = compute_chunk_losses(
                        model=model,
                        batch=chunk_batch,
                        training_cfg=training_cfg,
                        z_state=z_state,
                        state_mask=state_mask,
                        tbptt_steps=current_tbptt_steps,
                        freeze_gate_head_to_targets=current_gate_head_frozen,
                    )
                    accelerator.backward(total_loss)
                    accelerator.clip_grad_norm_(model.parameters(), float(training_cfg["max_grad_norm"]))
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad(set_to_none=True)

                z_state = z_state.detach()
                state_mask = state_mask.detach()

                if accelerator.sync_gradients:
                    global_step += 1
                    progress_bar.update(1)
                    progress_bar.set_postfix(
                        loss=f"{float(total_loss.detach()):.4f}",
                        bucket=batch_bucket,
                        tbptt=int(current_tbptt_steps or 0),
                        horizon=int(current_horizon_ticks),
                        decoder=current_decoder_trainable_layers,
                        gate="locked" if current_gate_head_frozen else "train",
                        refresh=False,
                    )

                    if global_step % int(training_cfg["logging_steps"]) == 0:
                        reduced = reduce_metrics(accelerator, [metrics], training_cfg)
                        reduced["lr"] = scheduler.get_last_lr()[0]
                        reduced["tbptt_steps"] = float(current_tbptt_steps or 0)
                        reduced["horizon_ticks"] = float(current_horizon_ticks)
                        reduced["chunk_ticks"] = float(chunk_ticks)
                        reduced["decoder_trainable_layers"] = float(current_decoder_trainable_layers)
                        reduced["gate_head_frozen"] = float(current_gate_head_frozen)
                        reduced["epoch"] = epoch + 1
                        reduced["global_step"] = global_step
                        accelerator.log({f"train/{key}": value for key, value in reduced.items()}, step=global_step)

                    if global_step % int(training_cfg["eval_steps"]) == 0:
                        validation_metrics = evaluate(
                            accelerator=accelerator,
                            model=model,
                            dataloader=validation_loader,
                            training_cfg=training_cfg,
                            rollout_cfg=rollout_cfg,
                            tokenizer_pad_token_id=tokenizer.pad_token_id,
                            max_batches=training_cfg.get("eval_max_batches"),
                        )
                        accelerator.log(
                            {f"validation/{key}": value for key, value in validation_metrics.items()},
                            step=global_step,
                        )

                        if validation_metrics and validation_metrics["loss_total"] < best_validation_loss:
                            best_validation_loss = validation_metrics["loss_total"]
                            best_metrics = validation_metrics
                            export_model(
                                accelerator=accelerator,
                                model=model,
                                output_dir=export_dir,
                                config=config,
                                config_path=config_path,
                                best_metrics=best_metrics,
                            )

                    if global_step % int(training_cfg["checkpoint_steps"]) == 0:
                        checkpoint_dir = ensure_dir(run_dir / f"checkpoint-{global_step}")
                        accelerator.save_state(str(checkpoint_dir))
                        if accelerator.is_main_process:
                            save_trainer_state(
                                checkpoint_dir,
                                global_step=global_step,
                                epoch=epoch,
                                next_batch_index=batch_index + 1,
                                best_validation_loss=best_validation_loss,
                                best_metrics=best_metrics,
                                wandb_run_id=wandb_run_id,
                            )
                        accelerator.wait_for_everyone()
                        maybe_push_checkpoint_to_hub(
                            accelerator=accelerator,
                            config=config,
                            checkpoint_dir=checkpoint_dir,
                        )
                        accelerator.wait_for_everyone()

    if not best_metrics:
        best_metrics = evaluate(
            accelerator=accelerator,
            model=model,
            dataloader=validation_loader,
            training_cfg=training_cfg,
            rollout_cfg=rollout_cfg,
            tokenizer_pad_token_id=tokenizer.pad_token_id,
            max_batches=training_cfg.get("eval_max_batches"),
        )
        export_model(
            accelerator=accelerator,
            model=model,
            output_dir=export_dir,
            config=config,
            config_path=config_path,
            best_metrics=best_metrics,
        )

    if accelerator.is_main_process:
        training_summary = {
            "train_conversations": len(train_dataset),
            "validation_conversations": len(validation_dataset),
            "train_chunk_microsteps": total_chunk_microsteps,
            "trainable_parameters": accelerator.unwrap_model(model).trainable_parameter_count(),
            "best_metrics": best_metrics,
            "global_step": global_step,
            "resumed_from_checkpoint": str(resume_checkpoint) if resume_checkpoint is not None else None,
            "wandb_run_id": wandb_run_id,
            "decoder_layer_count": unwrapped_model.decoder_layer_count(),
            "tick_seconds": tick_seconds,
            "chunk_ticks": chunk_ticks,
            "tbptt_steps_by_bucket": training_cfg.get("tbptt_steps_by_bucket", {}),
            "horizon_ticks_by_bucket": rollout_cfg.get("horizon_ticks_by_bucket", {}),
            "shuffle_train": bool(training_cfg.get("shuffle_train", True)),
            "sort_dataset_by_length": sort_dataset_by_length,
        }
        save_json(run_dir / "training_summary.json", training_summary)
        maybe_push_to_hub(config, export_dir)

    progress_bar.close()
    accelerator.end_training()


if __name__ == "__main__":
    main()