Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,36 @@
+# Samantha Thought-Loop SFT
+
+- Base model: `google/t5gemma-l-l-prefixlm-it`
+- Cleaned dataset repo: `BRlkl/timed_natural_conversations_5min4`
+- Raw dataset repo: `BRlkl/timed_natural_conversations_5min4`
+- Latent slots: `256`
+- Fixed tick seconds: `1.0`
+- Chunk ticks: `300`
+- Max horizon ticks: `300`
+- Explicit time features: `False`
+- Gate loss mode: `bce`
+- Hazard soft target window: `16`
+- Hazard soft target tau ticks: `4.0`
+- Hazard censor weight: `0.25`
+- TBPTT by bucket: `{}`
+
+## Best Validation Metrics
+
+- `active_ticks`: `1888.000000`
+- `decoder_token_count`: `2494.000000`
+- `gate_accuracy`: `0.715042`
+- `gate_f1`: `0.154088`
+- `gate_positive_weight`: `16.000000`
+- `gate_precision`: `0.091932`
+- `gate_pred_positive_count`: `533.000000`
+- `gate_recall`: `0.475728`
+- `gate_target_positive_count`: `103.000000`
+- `hazard_censor_count`: `0.000000`
+- `hazard_event_count`: `0.000000`
+- `loss_decoder`: `3.564287`
+- `loss_gate`: `1.241260`
+- `loss_gate_ranking`: `0.000000`
+- `loss_state_delta`: `0.258536`
+- `loss_total`: `4.521086`
+- `soft_target_mass`: `0.000000`
+- `soft_target_positive_count`: `0.000000`

best_metrics.json

@@ -0,0 +1,20 @@
+{
+  "loss_total": 4.521085739135742,
+  "loss_gate": 1.2412604093551636,
+  "loss_gate_ranking": 0.0,
+  "loss_decoder": 3.564286708831787,
+  "loss_state_delta": 0.25853580236434937,
+  "gate_accuracy": 0.7150423526763916,
+  "gate_precision": 0.09193246066570282,
+  "gate_recall": 0.4757281541824341,
+  "gate_f1": 0.15408805012702942,
+  "gate_target_positive_count": 103.0,
+  "gate_pred_positive_count": 533.0,
+  "gate_positive_weight": 16.0,
+  "hazard_event_count": 0.0,
+  "hazard_censor_count": 0.0,
+  "soft_target_mass": 0.0,
+  "soft_target_positive_count": 0.0,
+  "decoder_token_count": 2494.0,
+  "active_ticks": 1888.0
+}

chat_template.jinja

@@ -0,0 +1,4 @@
+{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '
+' + message['content'] | trim + '<end_of_turn>
+' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model
+'}}{% endif %}

config.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+def load_config(path: str | Path) -> dict[str, Any]:
+    with Path(path).open("r", encoding="utf-8") as handle:
+        data = yaml.safe_load(handle)
+    if not isinstance(data, dict):
+        raise ValueError(f"Config at {path} must deserialize into a mapping.")
+    return data
+
+
+def deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
+    merged = deepcopy(base)
+    for key, value in override.items():
+        if isinstance(value, dict) and isinstance(merged.get(key), dict):
+            merged[key] = deep_merge(merged[key], value)
+        else:
+            merged[key] = value
+    return merged
+
+
+def ensure_dir(path: str | Path) -> Path:
+    resolved = Path(path)
+    resolved.mkdir(parents=True, exist_ok=True)
+    return resolved
+
+
+def save_json(path: str | Path, payload: dict[str, Any]) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    with target.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+
+
+def flatten_for_wandb(payload: dict[str, Any], prefix: str = "") -> dict[str, Any]:
+    flat: dict[str, Any] = {}
+    for key, value in payload.items():
+        full_key = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            flat.update(flatten_for_wandb(value, prefix=full_key))
+        else:
+            flat[full_key] = value
+    return flat
+
+
+def fingerprint_payload(payload: dict[str, Any]) -> str:
+    serialized = json.dumps(payload, sort_keys=True, ensure_ascii=True, separators=(",", ":"))
+    return hashlib.sha256(serialized.encode("utf-8")).hexdigest()[:16]

data.py

@@ -0,0 +1,521 @@
+from __future__ import annotations
+
+import hashlib
+import math
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from datasets import Dataset as HFDataset
+from datasets import load_dataset
+from torch.utils.data import BatchSampler, Dataset
+
+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, fingerprint_payload
+
+
+OBS_ROLE_NONE = 0
+OBS_ROLE_USER = 1
+OBS_ROLE_AGENT_FEEDBACK = 2
+SFT_CACHE_SCHEMA_VERSION = 2
+
+DURATION_BUCKET_ORDER = {
+    "short": 0,
+    "medium": 1,
+    "long": 2,
+    "hour_scale": 3,
+}
+
+
+def load_clean_split(dataset_cfg: dict[str, Any], split: str) -> HFDataset:
+    local_parquet_dir = dataset_cfg.get("cleaned_local_parquet_dir")
+    if local_parquet_dir and Path(local_parquet_dir).exists():
+        parquet_path = Path(local_parquet_dir) / f"{split}.parquet"
+        if parquet_path.exists():
+            return load_dataset("parquet", data_files={split: str(parquet_path)}, split=split)
+
+    cleaned_repo_id = dataset_cfg.get("cleaned_repo_id")
+    if not cleaned_repo_id:
+        raise ValueError("Dataset config must provide either cleaned_local_parquet_dir or cleaned_repo_id.")
+
+    source_split = dataset_cfg.get("source_split")
+    if source_split:
+        dataset = load_dataset(cleaned_repo_id, split=str(source_split))
+        if split not in {"train", "validation"}:
+            return dataset
+
+        validation_ratio = float(dataset_cfg.get("validation_ratio", 0.02))
+        split_seed = int(dataset_cfg.get("split_seed", 17))
+        selected_indices = [
+            index
+            for index, row in enumerate(dataset)
+            if assign_deterministic_split(str(row["id"]), validation_ratio, split_seed) == split
+        ]
+        return dataset.select(selected_indices)
+
+    return load_dataset(cleaned_repo_id, split=split)
+
+
+def tokenize_text(tokenizer: Any, text: str | None, max_length: int) -> tuple[list[int], list[int]]:
+    if not text:
+        return [tokenizer.pad_token_id], [0]
+    encoded = tokenizer(
+        text,
+        add_special_tokens=True,
+        truncation=True,
+        max_length=max_length,
+    )
+    return encoded["input_ids"], encoded["attention_mask"]
+
+
+def tokenizer_supports_chat_template(tokenizer: Any) -> bool:
+    return bool(getattr(tokenizer, "chat_template", None))
+
+
+def should_use_base_chat_template(config: dict[str, Any], tokenizer: Any) -> bool:
+    dataset_cfg = config.get("dataset", {})
+    inference_cfg = config.get("inference", {})
+    if "use_base_chat_template" in inference_cfg:
+        return bool(inference_cfg["use_base_chat_template"]) and tokenizer_supports_chat_template(tokenizer)
+    return bool(dataset_cfg.get("use_base_chat_template", False)) and tokenizer_supports_chat_template(tokenizer)
+
+
+def build_user_generation_observation(tokenizer: Any, text: str, use_base_chat_template: bool) -> str:
+    if use_base_chat_template:
+        return tokenizer.apply_chat_template(
+            [{"role": "user", "content": text}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+    return text
+
+
+def build_assistant_feedback_observation(tokenizer: Any, text: str, use_base_chat_template: bool) -> str:
+    del tokenizer
+    if use_base_chat_template:
+        return f"<start_of_turn>model\n{text}<end_of_turn>\n"
+    return text
+
+
+def build_assistant_decoder_target(tokenizer: Any, text: str, use_base_chat_template: bool) -> str:
+    if not use_base_chat_template:
+        return text
+
+    prompt_text = tokenizer.apply_chat_template(
+        [{"role": "user", "content": ""}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+    full_text = tokenizer.apply_chat_template(
+        [
+            {"role": "user", "content": ""},
+            {"role": "assistant", "content": text},
+        ],
+        tokenize=False,
+        add_generation_prompt=False,
+    )
+    if not full_text.startswith(prompt_text):
+        raise ValueError("Chat template full text did not start with the generation prompt prefix.")
+    return full_text[len(prompt_text) :]
+
+
+def duration_bucket_rank(bucket: str | None) -> int:
+    return DURATION_BUCKET_ORDER.get(str(bucket), len(DURATION_BUCKET_ORDER))
+
+
+def assign_deterministic_split(row_id: str, validation_ratio: float, seed: int) -> str:
+    digest = hashlib.sha1(f"{seed}:{row_id}".encode("utf-8")).hexdigest()
+    score = int(digest[:8], 16) / 0xFFFFFFFF
+    return "validation" if score < validation_ratio else "train"
+
+
+def stable_tick_index(time_seconds: float, tick_seconds: float) -> int:
+    ratio = max(0.0, float(time_seconds)) / tick_seconds
+    return int(math.ceil(ratio - 1e-9))
+
+
+def quantize_message_sequence(
+    messages: list[dict[str, Any]],
+    *,
+    tick_seconds: float,
+    feedback_delay_seconds: float,
+) -> list[dict[str, Any]]:
+    raw_events: list[dict[str, Any]] = []
+    event_order = 0
+    for message in messages:
+        speaker = str(message["speaker"])
+        text = str(message["text"])
+        time_seconds = float(message["t"])
+
+        if speaker == "user":
+            raw_events.append(
+                {
+                    "time_seconds": time_seconds,
+                    "order": event_order,
+                    "kind": "user",
+                    "text": text,
+                }
+            )
+            event_order += 1
+            continue
+
+        raw_events.append(
+            {
+                "time_seconds": time_seconds,
+                "order": event_order,
+                "kind": "agent_speak",
+                "text": text,
+            }
+        )
+        event_order += 1
+        raw_events.append(
+            {
+                "time_seconds": time_seconds + feedback_delay_seconds,
+                "order": event_order,
+                "kind": "agent_feedback",
+                "text": text,
+            }
+        )
+        event_order += 1
+
+    raw_events.sort(key=lambda item: (float(item["time_seconds"]), int(item["order"])))
+
+    quantized_events: list[dict[str, Any]] = []
+    previous_tick = -1
+    for event in raw_events:
+        desired_tick = stable_tick_index(float(event["time_seconds"]), tick_seconds)
+        assigned_tick = max(desired_tick, previous_tick + 1)
+        quantized_events.append(
+            {
+                "tick": assigned_tick,
+                "kind": event["kind"],
+                "text": event["text"],
+            }
+        )
+        previous_tick = assigned_tick
+
+    return quantized_events
+
+
+def resolve_bucket_horizon_ticks(duration_bucket: str, rollout_cfg: dict[str, Any]) -> int:
+    bucket_map = rollout_cfg.get("horizon_ticks_by_bucket", {})
+    default_horizon = int(rollout_cfg.get("max_horizon_ticks", 36000))
+    return int(bucket_map.get(duration_bucket, default_horizon))
+
+
+def build_fixed_tick_conversation(
+    *,
+    row: dict[str, Any],
+    tokenizer: Any,
+    config: dict[str, Any],
+    rollout_cfg: dict[str, Any],
+    max_observation_tokens: int,
+    max_decoder_tokens: int,
+) -> dict[str, Any] | None:
+    tick_seconds = float(rollout_cfg["tick_seconds"])
+    chunk_ticks = int(rollout_cfg["chunk_ticks"])
+    feedback_delay_seconds = float(rollout_cfg["post_speech_feedback_delay_seconds"])
+    window_strategy = str(rollout_cfg.get("long_window_strategy", "tail"))
+    duration_bucket = str(row.get("meta", {}).get("duration_bucket", "short"))
+    use_base_chat_template = should_use_base_chat_template(config, tokenizer)
+
+    quantized_events = quantize_message_sequence(
+        messages=row["messages"],
+        tick_seconds=tick_seconds,
+        feedback_delay_seconds=feedback_delay_seconds,
+    )
+    total_ticks = max((event["tick"] for event in quantized_events), default=0) + 1
+
+    horizon_ticks = resolve_bucket_horizon_ticks(duration_bucket, rollout_cfg)
+    if bool(rollout_cfg.get("drop_overlong_examples", False)) and total_ticks > horizon_ticks:
+        return None
+
+    if total_ticks <= horizon_ticks:
+        effective_start_tick = 0
+        effective_total_ticks = total_ticks
+    elif window_strategy == "tail":
+        effective_start_tick = total_ticks - horizon_ticks
+        effective_total_ticks = horizon_ticks
+    else:
+        effective_start_tick = 0
+        effective_total_ticks = horizon_ticks
+
+    chunk_count = max(1, math.ceil(effective_total_ticks / chunk_ticks))
+    chunk_lengths = [
+        min(chunk_ticks, max(0, effective_total_ticks - chunk_index * chunk_ticks))
+        for chunk_index in range(chunk_count)
+    ]
+    chunk_events: list[list[dict[str, Any]]] = [[] for _ in range(chunk_count)]
+
+    event_count = 0
+    for event in quantized_events:
+        absolute_tick = int(event["tick"])
+        if absolute_tick < effective_start_tick or absolute_tick >= effective_start_tick + effective_total_ticks:
+            continue
+
+        relative_tick = absolute_tick - effective_start_tick
+        chunk_index = relative_tick // chunk_ticks
+        offset = relative_tick % chunk_ticks
+
+        if event["kind"] == "user":
+            observation_role = OBS_ROLE_USER
+            observation_text = build_user_generation_observation(
+                tokenizer,
+                str(event["text"]),
+                use_base_chat_template,
+            )
+            gate_target = 0
+            decoder_text = None
+        elif event["kind"] == "agent_feedback":
+            observation_role = OBS_ROLE_AGENT_FEEDBACK
+            observation_text = build_assistant_feedback_observation(
+                tokenizer,
+                str(event["text"]),
+                use_base_chat_template,
+            )
+            gate_target = 0
+            decoder_text = None
+        else:
+            observation_role = OBS_ROLE_NONE
+            observation_text = None
+            gate_target = 1
+            decoder_text = build_assistant_decoder_target(
+                tokenizer,
+                str(event["text"]),
+                use_base_chat_template,
+            )
+
+        observation_input_ids, observation_attention_mask = tokenize_text(
+            tokenizer=tokenizer,
+            text=observation_text,
+            max_length=max_observation_tokens,
+        )
+        decoder_labels, _ = tokenize_text(
+            tokenizer=tokenizer,
+            text=decoder_text,
+            max_length=max_decoder_tokens,
+        )
+        if gate_target == 0:
+            decoder_labels = []
+
+        chunk_events[chunk_index].append(
+            {
+                "offset": int(offset),
+                "absolute_tick": absolute_tick,
+                "observation_role": int(observation_role),
+                "observation_input_ids": observation_input_ids,
+                "observation_attention_mask": observation_attention_mask,
+                "gate_target": int(gate_target),
+                "decoder_labels": decoder_labels,
+            }
+        )
+        event_count += 1
+
+    return {
+        "row_id": str(row["id"]),
+        "scenario_id": row.get("meta", {}).get("scenario_id"),
+        "duration_bucket": duration_bucket,
+        "bucket_rank": duration_bucket_rank(duration_bucket),
+        "total_ticks": int(total_ticks),
+        "effective_start_tick": int(effective_start_tick),
+        "effective_total_ticks": int(effective_total_ticks),
+        "chunk_count": int(chunk_count),
+        "chunk_lengths": chunk_lengths,
+        "chunk_events": chunk_events,
+        "event_count": int(event_count),
+    }
+
+
+def build_chunk_batch(
+    *,
+    conversations: list[dict[str, Any]],
+    chunk_index: int,
+    pad_token_id: int,
+    chunk_ticks: int,
+    tick_seconds: float,
+) -> dict[str, torch.Tensor] | None:
+    active_lengths: list[int] = []
+    max_observation_tokens = 1
+    max_decoder_tokens = 1
+
+    for conversation in conversations:
+        if chunk_index >= int(conversation["chunk_count"]):
+            active_lengths.append(0)
+            continue
+        step_count = int(conversation["chunk_lengths"][chunk_index])
+        active_lengths.append(step_count)
+        for event in conversation["chunk_events"][chunk_index]:
+            max_observation_tokens = max(max_observation_tokens, len(event["observation_input_ids"]))
+            max_decoder_tokens = max(max_decoder_tokens, len(event["decoder_labels"]) or 1)
+
+    max_steps = max(active_lengths, default=0)
+    if max_steps <= 0:
+        return None
+
+    batch_size = len(conversations)
+    observation_input_ids = torch.full(
+        (batch_size, max_steps, max_observation_tokens),
+        fill_value=pad_token_id,
+        dtype=torch.long,
+    )
+    observation_attention_mask = torch.zeros((batch_size, max_steps, max_observation_tokens), dtype=torch.long)
+    decoder_labels = torch.full((batch_size, max_steps, max_decoder_tokens), fill_value=-100, dtype=torch.long)
+    tick_mask = torch.zeros((batch_size, max_steps), dtype=torch.bool)
+    gate_target = torch.zeros((batch_size, max_steps), dtype=torch.float32)
+    observation_role = torch.zeros((batch_size, max_steps), dtype=torch.long)
+    delta_seconds = torch.zeros((batch_size, max_steps), dtype=torch.float32)
+    elapsed_seconds = torch.zeros((batch_size, max_steps), dtype=torch.float32)
+
+    for batch_index, conversation in enumerate(conversations):
+        step_count = active_lengths[batch_index]
+        if step_count <= 0:
+            continue
+
+        tick_mask[batch_index, :step_count] = True
+        for local_tick in range(step_count):
+            global_tick = chunk_index * chunk_ticks + local_tick
+            delta_seconds[batch_index, local_tick] = 0.0 if global_tick == 0 else tick_seconds
+            elapsed_seconds[batch_index, local_tick] = global_tick * tick_seconds
+
+        for event in conversation["chunk_events"][chunk_index]:
+            offset = int(event["offset"])
+            observation_role[batch_index, offset] = int(event["observation_role"])
+            gate_target[batch_index, offset] = float(event["gate_target"])
+
+            observation_ids = event["observation_input_ids"]
+            observation_mask = event["observation_attention_mask"]
+            observation_input_ids[batch_index, offset, : len(observation_ids)] = torch.tensor(
+                observation_ids,
+                dtype=torch.long,
+            )
+            observation_attention_mask[batch_index, offset, : len(observation_mask)] = torch.tensor(
+                observation_mask,
+                dtype=torch.long,
+            )
+
+            labels = event["decoder_labels"]
+            if labels:
+                decoder_labels[batch_index, offset, : len(labels)] = torch.tensor(labels, dtype=torch.long)
+
+    return {
+        "tick_mask": tick_mask,
+        "gate_target": gate_target,
+        "observation_role": observation_role,
+        "observation_input_ids": observation_input_ids,
+        "observation_attention_mask": observation_attention_mask,
+        "decoder_labels": decoder_labels,
+        "delta_seconds": delta_seconds,
+        "elapsed_seconds": elapsed_seconds,
+    }
+
+
+class ThoughtLoopConversationDataset(Dataset):
+    def __init__(self, config: dict[str, Any], tokenizer: Any, split: str) -> None:
+        self.config = config
+        self.tokenizer = tokenizer
+        self.split = split
+        self.examples = self._load_or_build()
+
+    def _cache_path(self) -> Path:
+        dataset_cfg = self.config["dataset"]
+        model_cfg = self.config["model"]
+        rollout_cfg = self.config["rollout"]
+        cache_cfg = self.config["cache"]
+        cache_root = ensure_dir(cache_cfg["preprocessed_root"])
+        payload = {
+            "cache_schema_version": SFT_CACHE_SCHEMA_VERSION,
+            "split": self.split,
+            "dataset": dataset_cfg,
+            "model_name": model_cfg["base_model_name"],
+            "rollout": rollout_cfg,
+            "max_observation_tokens": model_cfg["max_observation_tokens"],
+            "max_decoder_tokens": model_cfg["max_decoder_tokens"],
+            "tokenizer_vocab_size": getattr(self.tokenizer, "vocab_size", None),
+        }
+        return cache_root / f"{self.split}_{fingerprint_payload(payload)}.pt"
+
+    def _load_or_build(self) -> list[dict[str, Any]]:
+        cache_path = self._cache_path()
+        if cache_path.exists():
+            return torch.load(cache_path, map_location="cpu")
+
+        dataset = load_clean_split(self.config["dataset"], self.split)
+        rollout_cfg = self.config["rollout"]
+        model_cfg = self.config["model"]
+
+        examples: list[dict[str, Any]] = []
+        for row in dataset:
+            example = build_fixed_tick_conversation(
+                row=row,
+                tokenizer=self.tokenizer,
+                config=self.config,
+                rollout_cfg=rollout_cfg,
+                max_observation_tokens=int(model_cfg["max_observation_tokens"]),
+                max_decoder_tokens=int(model_cfg["max_decoder_tokens"]),
+            )
+            if example is not None:
+                examples.append(example)
+
+        if bool(rollout_cfg.get("sort_by_duration_bucket", True)):
+            examples.sort(key=lambda item: (int(item["bucket_rank"]), int(item["total_ticks"]), str(item["row_id"])))
+
+        torch.save(examples, cache_path)
+        return examples
+
+    def __len__(self) -> int:
+        return len(self.examples)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        return self.examples[index]
+
+
+class DurationBucketBatchSampler(BatchSampler):
+    def __init__(
+        self,
+        dataset: ThoughtLoopConversationDataset,
+        batch_size: int,
+    ) -> None:
+        self.dataset = dataset
+        self.batch_size = batch_size
+        self.batches: list[list[int]] = []
+
+        current_bucket: str | None = None
+        current_batch: list[int] = []
+        for index, example in enumerate(self.dataset.examples):
+            bucket = str(example["duration_bucket"])
+            if current_batch and (bucket != current_bucket or len(current_batch) >= self.batch_size):
+                self.batches.append(current_batch)
+                current_batch = []
+            current_bucket = bucket
+            current_batch.append(index)
+            if len(current_batch) >= self.batch_size:
+                self.batches.append(current_batch)
+                current_batch = []
+
+        if current_batch:
+            self.batches.append(current_batch)
+
+    def __iter__(self) -> Any:
+        return iter(self.batches)
+
+    def __len__(self) -> int:
+        return len(self.batches)
+
+
+def identity_collate(batch: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    return batch
+
+
+def estimate_total_chunk_microsteps(
+    *,
+    dataset: ThoughtLoopConversationDataset,
+    batch_sampler: DurationBucketBatchSampler,
+) -> int:
+    total_microsteps = 0
+    for batch_indices in batch_sampler.batches:
+        total_microsteps += max(int(dataset.examples[index]["chunk_count"]) for index in batch_indices)
+    return total_microsteps

inference.py

@@ -0,0 +1,546 @@
+from __future__ import annotations
+
+import argparse
+from collections import deque
+import json
+import math
+import sys
+import threading
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+from data import (
+    OBS_ROLE_AGENT_FEEDBACK,
+    OBS_ROLE_NONE,
+    OBS_ROLE_USER,
+    build_assistant_feedback_observation,
+    build_user_generation_observation,
+    quantize_message_sequence,
+    should_use_base_chat_template,
+    tokenize_text,
+)
+from model import ThoughtLoopT5Gemma
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Interactive inference for Samantha thought-loop checkpoints, with a direct HF baseline fallback."
+    )
+    parser.add_argument("--model-path", required=True, help="Local export directory or Hugging Face repo id.")
+    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    parser.add_argument("--history-json", default=None, help="Optional JSON array of {t, speaker, text} turns to replay.")
+    parser.add_argument("--user", default=None, help="Optional single user message to answer.")
+    parser.add_argument("--interactive", action="store_true", help="Run a live terminal session.")
+    parser.add_argument("--tick-seconds", type=float, default=None, help="Runtime tick size. Defaults to the trained value.")
+    parser.add_argument("--max-thought-seconds", type=float, default=300.0)
+    parser.add_argument("--settle-thought-seconds", type=float, default=1.0)
+    parser.add_argument("--max-autonomous-utterances", type=int, default=4, help="One-shot --user safety cap.")
+    parser.add_argument("--gate-threshold", type=float, default=0.5, help="Trained binary gate decision threshold.")
+    parser.add_argument("--max-new-tokens", type=int, default=160)
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument(
+        "--debug-gate",
+        "--show-gate",
+        dest="debug_gate",
+        action="store_true",
+        help="Print the gate head probability on every runtime tick.",
+    )
+    return parser.parse_args()
+
+
+def build_generation_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    kwargs: dict[str, Any] = {
+        "max_new_tokens": int(args.max_new_tokens),
+        "do_sample": float(args.temperature) > 0.0,
+    }
+    if kwargs["do_sample"]:
+        kwargs["temperature"] = float(args.temperature)
+        kwargs["top_p"] = float(args.top_p)
+    return kwargs
+
+
+def normalize_history_speaker(speaker: str | None) -> str | None:
+    normalized = str(speaker or "").strip().lower()
+    if normalized in {"user", "human", "prompter"}:
+        return "user"
+    if normalized in {"assistant", "agent", "model", "gpt"}:
+        return "assistant"
+    return None
+
+
+def load_base_seq2seq_model(
+    model_path: str,
+    *,
+    device: str | torch.device,
+) -> tuple[torch.nn.Module, Any]:
+    resolved_device = torch.device(device)
+    dtype = torch.bfloat16 if resolved_device.type == "cuda" else torch.float32
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_path, torch_dtype=dtype)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    model.to(resolved_device)
+    model.eval()
+    return model, tokenizer
+
+
+def load_inference_backend(
+    model_path: str,
+    *,
+    device: str | torch.device,
+) -> tuple[str, Any, Any | None]:
+    try:
+        model = ThoughtLoopT5Gemma.from_pretrained(model_path, device=device, map_location=device)
+        model.eval()
+        return "thought_loop", model, None
+    except FileNotFoundError as exc:
+        if Path(str(exc.filename)).name != "sft_config.json":
+            raise
+        print(
+            "[inference] No sft_config.json found at model path; loading as a raw Hugging Face seq2seq baseline.",
+            file=sys.stderr,
+        )
+        base_model, tokenizer = load_base_seq2seq_model(model_path, device=device)
+        return "base_seq2seq", base_model, tokenizer
+
+
+class BaseSeq2SeqSession:
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        tokenizer: Any,
+        device: str | torch.device,
+        generation_kwargs: dict[str, Any],
+    ) -> None:
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device(device)
+        self.generation_kwargs = generation_kwargs
+        self.history: list[dict[str, str]] = []
+
+    def replay_history(self, messages: list[dict[str, Any]]) -> None:
+        ordered_messages = sorted(messages, key=lambda item: float(item.get("t", 0.0)))
+        for message in ordered_messages:
+            role = normalize_history_speaker(message.get("speaker"))
+            text = str(message.get("text", "")).strip()
+            if role is None or not text:
+                continue
+            self.history.append({"role": role, "content": text})
+
+    def _format_prompt(self, user_text: str) -> str:
+        messages = [*self.history, {"role": "user", "content": user_text}]
+        if getattr(self.tokenizer, "chat_template", None):
+            return self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+
+        lines: list[str] = []
+        for message in messages:
+            role = "User" if message["role"] == "user" else "Assistant"
+            lines.append(f"{role}: {message['content']}")
+        lines.append("Assistant:")
+        return "\n".join(lines)
+
+    def respond(
+        self,
+        user_text: str,
+        *,
+        max_thought_seconds: float,
+        settle_thought_seconds: float,
+        max_autonomous_utterances: int,
+    ) -> tuple[list[str], float]:
+        del max_thought_seconds, settle_thought_seconds, max_autonomous_utterances
+        prompt = self._format_prompt(user_text)
+        encoded = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            generated = self.model.generate(**encoded, **self.generation_kwargs)
+        response = self.tokenizer.decode(generated[0], skip_special_tokens=True).strip()
+        self.history.append({"role": "user", "content": user_text})
+        self.history.append({"role": "assistant", "content": response})
+        return [response], 0.0
+
+
+class TerminalInputBuffer:
+    def __init__(self) -> None:
+        self._lines: deque[str] = deque()
+        self._lock = threading.Lock()
+
+    def push(self, line: str) -> None:
+        with self._lock:
+            self._lines.append(line)
+
+    def drain(self) -> list[str]:
+        with self._lock:
+            lines = list(self._lines)
+            self._lines.clear()
+        return lines
+
+    def prepend_many(self, lines: list[str]) -> None:
+        if not lines:
+            return
+        with self._lock:
+            for line in reversed(lines):
+                self._lines.appendleft(line)
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._lines)
+
+
+def start_terminal_reader(line_buffer: TerminalInputBuffer, stop_event: threading.Event) -> threading.Thread:
+    def read_loop() -> None:
+        while not stop_event.is_set():
+            line = sys.stdin.readline()
+            if line == "":
+                stop_event.set()
+                break
+            line_buffer.push(line.rstrip("\n"))
+
+    thread = threading.Thread(target=read_loop, name="terminal-input-reader", daemon=True)
+    thread.start()
+    return thread
+
+
+def print_gate_debug(tick_index: int, event_name: str, gate_probability: float) -> None:
+    print(f"[tick={tick_index} event={event_name} gate={gate_probability:.6f}]", flush=True)
+
+
+class ThoughtLoopSession:
+    def __init__(
+        self,
+        model: ThoughtLoopT5Gemma,
+        tick_seconds: float,
+        gate_threshold: float,
+        generation_kwargs: dict[str, Any],
+    ) -> None:
+        self.model = model
+        self.tokenizer = model.tokenizer
+        self.tick_seconds = tick_seconds
+        self.gate_threshold = gate_threshold
+        self.generation_kwargs = generation_kwargs
+        self.z_state = model.initial_state(batch_size=1, device=model.device)
+        self.z_mask = model.initial_state_mask(batch_size=1, device=model.device)
+        self.use_base_chat_template = should_use_base_chat_template(model.config, self.tokenizer)
+        self.elapsed_seconds = 0.0
+        self.next_tick_index = 0
+        self.feedback_delay_seconds = float(model.config.get("rollout", {}).get("post_speech_feedback_delay_seconds", tick_seconds))
+
+    def _format_observation(self, observation_role: int, observation_text: str | None) -> str | None:
+        if observation_text is None:
+            return None
+        if observation_role == OBS_ROLE_USER:
+            return build_user_generation_observation(
+                self.tokenizer,
+                observation_text,
+                self.use_base_chat_template,
+            )
+        if observation_role == OBS_ROLE_AGENT_FEEDBACK:
+            return build_assistant_feedback_observation(
+                self.tokenizer,
+                observation_text,
+                self.use_base_chat_template,
+            )
+        return observation_text
+
+    def _advance_raw(self, delta_seconds: float, observation_role: int, observation_text: str | None) -> float:
+        formatted_observation_text = self._format_observation(observation_role, observation_text)
+        input_ids, attention_mask = tokenize_text(
+            tokenizer=self.tokenizer,
+            text=formatted_observation_text,
+            max_length=int(self.model.config["model"]["max_observation_tokens"]),
+        )
+        input_ids_tensor = torch.tensor([input_ids], dtype=torch.long, device=self.model.device)
+        attention_mask_tensor = torch.tensor([attention_mask], dtype=torch.long, device=self.model.device)
+        role_tensor = torch.tensor([observation_role], dtype=torch.long, device=self.model.device)
+
+        self.elapsed_seconds += delta_seconds
+
+        with torch.no_grad():
+            self.z_state, gate_logits, self.z_mask = self.model.rollout_step_with_mask(
+                z_state=self.z_state,
+                state_mask=self.z_mask,
+                observation_input_ids=input_ids_tensor,
+                observation_attention_mask=attention_mask_tensor,
+                observation_role_ids=role_tensor,
+                delta_seconds=torch.tensor([delta_seconds], dtype=torch.float32, device=self.model.device),
+                elapsed_seconds=torch.tensor([self.elapsed_seconds], dtype=torch.float32, device=self.model.device),
+                since_last_user_seconds=torch.zeros(1, dtype=torch.float32, device=self.model.device),
+                since_last_agent_seconds=torch.zeros(1, dtype=torch.float32, device=self.model.device),
+            )
+        return torch.sigmoid(gate_logits).item()
+
+    def _advance_tick(self, observation_role: int, observation_text: str | None) -> float:
+        delta_seconds = 0.0 if self.next_tick_index == 0 else self.tick_seconds
+        gate_probability = self._advance_raw(delta_seconds, observation_role, observation_text)
+        self.next_tick_index += 1
+        return gate_probability
+
+    def advance_tick(self, observation_role: int, observation_text: str | None) -> float:
+        return self._advance_tick(observation_role, observation_text)
+
+    def replay_history(self, messages: list[dict[str, Any]]) -> None:
+        quantized_events = quantize_message_sequence(
+            messages=messages,
+            tick_seconds=self.tick_seconds,
+            feedback_delay_seconds=self.feedback_delay_seconds,
+        )
+        events_by_tick = {int(event["tick"]): event for event in quantized_events}
+        if not events_by_tick:
+            return
+
+        last_tick = max(events_by_tick)
+        while self.next_tick_index <= last_tick:
+            event = events_by_tick.get(self.next_tick_index)
+            if event is None:
+                self._advance_tick(OBS_ROLE_NONE, None)
+                continue
+
+            if event["kind"] == "user":
+                self._advance_tick(OBS_ROLE_USER, str(event["text"]))
+            elif event["kind"] == "agent_feedback":
+                self._advance_tick(OBS_ROLE_AGENT_FEEDBACK, str(event["text"]))
+            else:
+                self._advance_tick(OBS_ROLE_NONE, None)
+
+    def _generate_response(self) -> str:
+        generated = self.model.generate_from_state(
+            self.z_state,
+            encoder_attention_mask=self.z_mask,
+            **self.generation_kwargs,
+        )
+        return self.tokenizer.decode(generated[0], skip_special_tokens=True).strip()
+
+    def generate_response(self) -> str:
+        return self._generate_response()
+
+    def respond(
+        self,
+        user_text: str,
+        *,
+        max_thought_seconds: float,
+        settle_thought_seconds: float,
+        max_autonomous_utterances: int,
+        debug_gate: bool = False,
+    ) -> tuple[list[str], float]:
+        gate_probability = self._advance_tick(OBS_ROLE_USER, user_text)
+        if debug_gate:
+            print_gate_debug(self.next_tick_index - 1, "user", gate_probability)
+        utterances: list[str] = []
+        thought_ticks = 0
+        settle_ticks = max(1, int(math.ceil(settle_thought_seconds / self.tick_seconds)))
+        max_thought_ticks = max(1, int(math.ceil(max_thought_seconds / self.tick_seconds)))
+        idle_ticks_after_last_speech = 0
+        pending_assistant_feedback: str | None = None
+
+        while thought_ticks < max_thought_ticks:
+            if pending_assistant_feedback is not None:
+                event_name = "assistant_feedback"
+                gate_probability = self._advance_tick(OBS_ROLE_AGENT_FEEDBACK, pending_assistant_feedback)
+                pending_assistant_feedback = None
+            else:
+                event_name = "idle"
+                gate_probability = self._advance_tick(OBS_ROLE_NONE, None)
+            if debug_gate:
+                print_gate_debug(self.next_tick_index - 1, event_name, gate_probability)
+            thought_ticks += 1
+            idle_ticks_after_last_speech += 1
+
+            if gate_probability >= self.gate_threshold:
+                response = self._generate_response()
+                utterances.append(response)
+                idle_ticks_after_last_speech = 0
+                if len(utterances) >= max_autonomous_utterances:
+                    break
+                pending_assistant_feedback = response
+                continue
+
+            if utterances and idle_ticks_after_last_speech >= settle_ticks:
+                break
+
+        return utterances, thought_ticks * self.tick_seconds
+
+
+def consume_next_live_user_line(
+    line_buffer: TerminalInputBuffer,
+    *,
+    stop_event: threading.Event,
+    status_text: str,
+) -> str | None:
+    lines = line_buffer.drain()
+    if not lines:
+        return None
+
+    selected_line: str | None = None
+    deferred_lines: list[str] = []
+    for raw_line in lines:
+        line = raw_line.strip()
+        if not line:
+            continue
+        command = line.lower()
+        if command in {"/exit", "/quit", "exit", "quit"}:
+            stop_event.set()
+            continue
+        if command == "/status":
+            print(status_text, flush=True)
+            continue
+        if selected_line is None:
+            selected_line = line
+        else:
+            deferred_lines.append(line)
+
+    line_buffer.prepend_many(deferred_lines)
+    return selected_line
+
+
+def run_live_thought_loop_terminal(
+    session: ThoughtLoopSession,
+    *,
+    gate_threshold: float,
+    debug_gate: bool,
+) -> None:
+    line_buffer = TerminalInputBuffer()
+    stop_event = threading.Event()
+    start_terminal_reader(line_buffer, stop_event)
+    pending_assistant_feedback: str | None = None
+
+    print("Live Samantha session.", flush=True)
+    print("Type anytime and press Enter. Each line is one user event on the next 1 Hz tick.", flush=True)
+    print("Commands: /status, /exit", flush=True)
+
+    while not stop_event.is_set():
+        tick_started_at = time.monotonic()
+        status_text = (
+            f"[status] tick={session.next_tick_index} elapsed={session.elapsed_seconds:.1f}s "
+            f"queued_user_lines={len(line_buffer)} pending_assistant_feedback={pending_assistant_feedback is not None}"
+        )
+
+        if pending_assistant_feedback is not None:
+            observation_role = OBS_ROLE_AGENT_FEEDBACK
+            observation_text = pending_assistant_feedback
+            event_name = "assistant_feedback"
+            pending_assistant_feedback = None
+        else:
+            user_line = consume_next_live_user_line(
+                line_buffer,
+                stop_event=stop_event,
+                status_text=status_text,
+            )
+            if stop_event.is_set():
+                break
+
+            if user_line is not None:
+                observation_role = OBS_ROLE_USER
+                observation_text = user_line
+                event_name = "user"
+                print(f"user> {user_line}", flush=True)
+            else:
+                observation_role = OBS_ROLE_NONE
+                observation_text = None
+                event_name = "idle"
+
+        gate_probability = session.advance_tick(observation_role, observation_text)
+        if debug_gate:
+            print_gate_debug(session.next_tick_index - 1, event_name, gate_probability)
+
+        if observation_role != OBS_ROLE_USER and gate_probability >= gate_threshold:
+            response = session.generate_response()
+            print(f"assistant> {response}", flush=True)
+            pending_assistant_feedback = response
+
+        sleep_seconds = session.tick_seconds - (time.monotonic() - tick_started_at)
+        if sleep_seconds > 0.0:
+            time.sleep(sleep_seconds)
+
+
+def run_blocking_base_terminal(session: BaseSeq2SeqSession) -> None:
+    print("Raw HF baseline terminal mode. This model has no Samantha Z state or speech gate.", flush=True)
+    print("Type a message and press Enter. Commands: /exit", flush=True)
+    while True:
+        user_text = input("user> ").strip()
+        if user_text.lower() in {"/exit", "/quit", "exit", "quit"}:
+            break
+        if not user_text:
+            continue
+        responses, _ = session.respond(
+            user_text,
+            max_thought_seconds=0.0,
+            settle_thought_seconds=0.0,
+            max_autonomous_utterances=1,
+        )
+        for response in responses:
+            print(f"assistant> {response}", flush=True)
+
+
+def main() -> None:
+    args = parse_args()
+    backend_kind, model, tokenizer = load_inference_backend(args.model_path, device=args.device)
+    generation_kwargs = build_generation_kwargs(args)
+
+    if backend_kind == "thought_loop":
+        tick_seconds = (
+            float(args.tick_seconds)
+            if args.tick_seconds is not None
+            else float(model.config.get("rollout", {}).get("tick_seconds", 1.0))
+        )
+        session: Any = ThoughtLoopSession(
+            model=model,
+            tick_seconds=tick_seconds,
+            gate_threshold=args.gate_threshold,
+            generation_kwargs=generation_kwargs,
+        )
+    else:
+        assert tokenizer is not None
+        session = BaseSeq2SeqSession(
+            model=model,
+            tokenizer=tokenizer,
+            device=args.device,
+            generation_kwargs=generation_kwargs,
+        )
+
+    if args.history_json:
+        history = json.loads(Path(args.history_json).read_text(encoding="utf-8"))
+        if not isinstance(history, list):
+            raise ValueError("History JSON must be a list of turn objects.")
+        session.replay_history(history)
+
+    if args.user:
+        responses, thought_time = session.respond(
+            args.user,
+            max_thought_seconds=args.max_thought_seconds,
+            settle_thought_seconds=args.settle_thought_seconds,
+            max_autonomous_utterances=args.max_autonomous_utterances,
+            debug_gate=bool(args.debug_gate),
+        )
+        print(
+            json.dumps(
+                {
+                    "responses": responses,
+                    "thought_seconds": thought_time,
+                },
+                indent=2,
+                ensure_ascii=False,
+            )
+        )
+
+    if args.interactive or args.user is None:
+        if backend_kind == "thought_loop":
+            run_live_thought_loop_terminal(
+                session,
+                gate_threshold=float(args.gate_threshold),
+                debug_gate=bool(args.debug_gate),
+            )
+        else:
+            run_blocking_base_terminal(session)
+
+
+if __name__ == "__main__":
+    main()

initial_latent_z.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c8a7ba9282608a91981002cf45998c493d47ba4ccb5538fae1f183e9659d344
+size 1048656

model.py

@@ -0,0 +1,777 @@
+from __future__ import annotations
+
+import json
+import math
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_model, save_file, save_model
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.modeling_outputs import BaseModelOutput
+
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+
+class ParameterlessRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.eps = eps
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        original_dtype = inputs.dtype
+        normalized = inputs.float()
+        normalized = normalized * torch.rsqrt(normalized.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return normalized.to(dtype=original_dtype)
+
+
+class TemporalFeatureProjector(nn.Module):
+    def __init__(self, hidden_size: int, num_time_tokens: int, scalar_count: int = 4) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_time_tokens = num_time_tokens
+        self.mlp = nn.Sequential(
+            nn.Linear(scalar_count, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size * num_time_tokens),
+        )
+        self.norm = nn.LayerNorm(hidden_size)
+
+    def forward(self, temporal_features: torch.Tensor) -> torch.Tensor:
+        projected = self.mlp(temporal_features)
+        projected = projected.view(temporal_features.size(0), self.num_time_tokens, self.hidden_size)
+        return self.norm(projected)
+
+
+class ThoughtLoopT5Gemma(nn.Module):
+    def __init__(self, config: dict[str, Any]) -> None:
+        super().__init__()
+        self.config = config
+        model_cfg = config["model"]
+        training_cfg = config.get("training", {})
+
+        backbone_dtype = getattr(torch, model_cfg.get("dtype", "bfloat16"))
+        backbone_kwargs: dict[str, Any] = {
+            # transformers currently warns that torch_dtype is deprecated in favor of dtype,
+            # but torch_dtype remains compatible across more installed versions.
+            "torch_dtype": backbone_dtype,
+        }
+        attn_implementation = model_cfg.get("attn_implementation")
+        if attn_implementation:
+            backbone_kwargs["attn_implementation"] = attn_implementation
+
+        self.backbone = AutoModelForSeq2SeqLM.from_pretrained(model_cfg["base_model_name"], **backbone_kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_cfg["base_model_name"])
+
+        if training_cfg.get("gradient_checkpointing", True):
+            self.backbone.gradient_checkpointing_enable()
+            if getattr(self.backbone.config, "use_cache", None) is not None:
+                self.backbone.config.use_cache = False
+
+        self.encoder = self.backbone.get_encoder()
+        self.decoder = self.backbone.get_decoder()
+        self.input_embeddings = self.backbone.get_input_embeddings()
+        hidden_size = int(self.backbone.config.decoder.hidden_size)
+        self.hidden_size = hidden_size
+        self.z_slots = int(model_cfg["z_slots"])
+        self.thought_loop_proposal_mode = str(
+            model_cfg.get("thought_loop_proposal_mode", "latent_prefix")
+        ).strip().lower()
+        if self.thought_loop_proposal_mode not in {"latent_prefix", "observation_hidden_compression"}:
+            raise ValueError(f"Unsupported thought_loop_proposal_mode: {self.thought_loop_proposal_mode}")
+        self.preserve_observation_encoder_manifold = bool(
+            model_cfg.get(
+                "preserve_observation_encoder_manifold",
+                self.thought_loop_proposal_mode == "observation_hidden_compression",
+            )
+        )
+        self.observation_encoder_use_state_context = bool(
+            model_cfg.get("observation_encoder_use_state_context", False)
+        )
+        self.latent_attention_mask_mode = str(
+            model_cfg.get("latent_attention_mask_mode", "observed")
+        ).strip().lower()
+        if self.latent_attention_mask_mode not in {"observed", "full"}:
+            raise ValueError(f"Unsupported latent_attention_mask_mode: {self.latent_attention_mask_mode}")
+        self.use_explicit_time_features = bool(model_cfg.get("use_explicit_time_features", False))
+        self.num_time_tokens = int(model_cfg["num_time_tokens"]) if self.use_explicit_time_features else 0
+        self.observation_role_count = 3
+        magicnorm_eps = float(model_cfg.get("magicnorm_eps", 1e-6))
+
+        # Keep the newly initialized recurrent/gating modules in fp32. We cast only the
+        # tensors handed into the bf16 backbone. This avoids dtype crashes while preserving
+        # stable optimizer state for the custom modules.
+        self.z_init = nn.Parameter(torch.randn(self.z_slots, hidden_size) * 0.02)
+        self.segment_embeddings = nn.Embedding(3, hidden_size)
+        self.observation_role_embeddings = nn.Embedding(self.observation_role_count, hidden_size)
+        self.temporal_projector = (
+            TemporalFeatureProjector(hidden_size, self.num_time_tokens) if self.use_explicit_time_features else None
+        )
+
+        self.state_gate = nn.Sequential(
+            nn.Linear(hidden_size * 2, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+        self._initialize_update_gate_bias(model_cfg)
+        self.state_gate_input_norm = ParameterlessRMSNorm(hidden_size * 2, eps=magicnorm_eps)
+        self.proposed_state_norm = ParameterlessRMSNorm(hidden_size, eps=magicnorm_eps)
+        self.recurrent_state_norm = ParameterlessRMSNorm(hidden_size, eps=magicnorm_eps)
+        self.gate_context_norm = ParameterlessRMSNorm(hidden_size, eps=magicnorm_eps)
+
+        self.gate_query = nn.Parameter(torch.randn(1, 1, hidden_size) * 0.02)
+        self.gate_pool = nn.MultiheadAttention(
+            embed_dim=hidden_size,
+            num_heads=int(model_cfg.get("gate_attention_heads", 4)),
+            batch_first=True,
+        )
+        self.gate_head = nn.Sequential(
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, hidden_size),
+            nn.GELU(),
+            nn.Linear(hidden_size, 1),
+        )
+
+        self._drop_unused_vision_modules()
+
+    def _initialize_update_gate_bias(self, model_cfg: dict[str, Any]) -> None:
+        if "initial_update_gate_bias" not in model_cfg:
+            return
+        final_gate_layer = self.state_gate[-1]
+        if not isinstance(final_gate_layer, nn.Linear) or final_gate_layer.bias is None:
+            raise RuntimeError("state_gate must end in a Linear layer with bias to use initial_update_gate_bias.")
+        nn.init.constant_(final_gate_layer.bias, float(model_cfg["initial_update_gate_bias"]))
+
+    def _drop_unused_vision_modules(self) -> None:
+        for module_path in (
+            "encoder.vision_tower",
+            "encoder.vision_model",
+            "model.encoder.vision_tower",
+            "model.encoder.vision_model",
+        ):
+            parent: Any = self.backbone
+            path_parts = module_path.split(".")
+            try:
+                for part in path_parts[:-1]:
+                    parent = getattr(parent, part)
+            except AttributeError:
+                continue
+
+            child_name = path_parts[-1]
+            child = getattr(parent, child_name, None)
+            if isinstance(child, nn.Module):
+                setattr(parent, child_name, None)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def backbone_dtype(self) -> torch.dtype:
+        return self.input_embeddings.weight.dtype
+
+    @property
+    def recurrent_dtype(self) -> torch.dtype:
+        return self.z_init.dtype
+
+    def trainable_parameter_count(self) -> int:
+        return sum(parameter.numel() for parameter in self.parameters() if parameter.requires_grad)
+
+    def _decoder_blocks(self) -> Any:
+        for attribute_name in ("block", "layers", "h"):
+            blocks = getattr(self.decoder, attribute_name, None)
+            if blocks is not None:
+                return blocks
+        return None
+
+    def decoder_layer_count(self) -> int:
+        blocks = self._decoder_blocks()
+        return len(blocks) if blocks is not None else 0
+
+    def set_gate_trainable(self, trainable: bool) -> None:
+        self.gate_query.requires_grad_(trainable)
+        self.gate_pool.requires_grad_(trainable)
+        self.gate_head.requires_grad_(trainable)
+        self.gate_context_norm.requires_grad_(trainable)
+
+    def set_decoder_trainable_fraction(self, trainable_fraction: float) -> int:
+        blocks = self._decoder_blocks()
+        if blocks is None:
+            return 0
+
+        clamped_fraction = min(max(trainable_fraction, 0.0), 1.0)
+        total_blocks = len(blocks)
+        trainable_blocks = min(total_blocks, math.ceil(total_blocks * clamped_fraction)) if clamped_fraction > 0 else 0
+        first_trainable_index = total_blocks - trainable_blocks
+
+        for block_index, block in enumerate(blocks):
+            block.requires_grad_(block_index >= first_trainable_index)
+
+        for attribute_name in ("final_layer_norm", "norm", "layer_norm"):
+            maybe_module = getattr(self.decoder, attribute_name, None)
+            if isinstance(maybe_module, nn.Module):
+                maybe_module.requires_grad_(trainable_blocks > 0)
+
+        decoder_embeddings = getattr(self.decoder, "embed_tokens", None)
+        if isinstance(decoder_embeddings, nn.Module):
+            decoder_embeddings.requires_grad_(trainable_blocks > 0)
+
+        for module_path in (
+            "lm_head",
+            "model.lm_head",
+            "backbone.lm_head",
+            "shared",
+            "model.shared",
+        ):
+            parent: Any = self
+            path_parts = module_path.split(".")
+            try:
+                for part in path_parts[:-1]:
+                    parent = getattr(parent, part)
+            except AttributeError:
+                continue
+
+            child_name = path_parts[-1]
+            child = getattr(parent, child_name, None)
+            if isinstance(child, nn.Module):
+                child.requires_grad_(trainable_blocks > 0)
+
+        return trainable_blocks
+
+    def initial_state(self, batch_size: int, device: torch.device | None = None) -> torch.Tensor:
+        target_device = device or self.device
+        initial_state = self.z_init.unsqueeze(0).expand(batch_size, -1, -1).to(device=target_device)
+        return self.recurrent_state_norm(initial_state)
+
+    def initial_state_mask(self, batch_size: int, device: torch.device | None = None) -> torch.Tensor:
+        target_device = device or self.device
+        if self.latent_attention_mask_mode == "full":
+            return torch.ones(batch_size, self.z_slots, dtype=torch.long, device=target_device)
+        if self.thought_loop_proposal_mode == "observation_hidden_compression":
+            return torch.zeros(batch_size, self.z_slots, dtype=torch.long, device=target_device)
+        return torch.ones(batch_size, self.z_slots, dtype=torch.long, device=target_device)
+
+    def _build_temporal_tensor(
+        self,
+        delta_seconds: torch.Tensor,
+        elapsed_seconds: torch.Tensor,
+        since_last_user_seconds: torch.Tensor,
+        since_last_agent_seconds: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.use_explicit_time_features:
+            raise RuntimeError("Temporal features were requested, but this model is configured to disable them.")
+        stacked = torch.stack(
+            [
+                torch.log1p(delta_seconds),
+                torch.log1p(elapsed_seconds),
+                torch.log1p(since_last_user_seconds),
+                torch.log1p(since_last_agent_seconds),
+            ],
+            dim=-1,
+        )
+        return stacked
+
+    def compress_hidden_states_to_slots(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        target_slots: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        slots = int(target_slots if target_slots is not None else self.z_slots)
+        batch_size, _, hidden_size = hidden_states.shape
+        compressed_hidden = hidden_states.new_zeros((batch_size, slots, hidden_size))
+        compressed_mask = torch.zeros(batch_size, slots, dtype=torch.long, device=hidden_states.device)
+        attention_mask = attention_mask.to(device=hidden_states.device, dtype=torch.bool)
+
+        for batch_index in range(batch_size):
+            valid_states = hidden_states[batch_index, attention_mask[batch_index]]
+            valid_length = int(valid_states.shape[0])
+            if valid_length <= 0:
+                continue
+
+            if valid_length <= slots:
+                compressed_hidden[batch_index, :valid_length] = valid_states
+                compressed_mask[batch_index, :valid_length] = 1
+                continue
+
+            for slot_index in range(slots):
+                start = (slot_index * valid_length) // slots
+                end = ((slot_index + 1) * valid_length) // slots
+                if end <= start:
+                    end = min(valid_length, start + 1)
+                compressed_hidden[batch_index, slot_index] = valid_states[start:end].mean(dim=0)
+                compressed_mask[batch_index, slot_index] = 1
+
+        return compressed_hidden, compressed_mask
+
+    def _rollout_step_impl(
+        self,
+        z_state: torch.Tensor,
+        observation_input_ids: torch.Tensor,
+        observation_attention_mask: torch.Tensor,
+        observation_role_ids: torch.Tensor,
+        delta_seconds: torch.Tensor,
+        elapsed_seconds: torch.Tensor,
+        since_last_user_seconds: torch.Tensor,
+        since_last_agent_seconds: torch.Tensor,
+        previous_state_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch_size = z_state.size(0)
+        recurrent_dtype = z_state.dtype
+        backbone_dtype = self.backbone_dtype
+
+        if self.thought_loop_proposal_mode == "latent_prefix":
+            aux_tokens: list[torch.Tensor] = []
+            aux_masks: list[torch.Tensor] = []
+            if self.use_explicit_time_features:
+                temporal_features = self._build_temporal_tensor(
+                    delta_seconds=delta_seconds,
+                    elapsed_seconds=elapsed_seconds,
+                    since_last_user_seconds=since_last_user_seconds,
+                    since_last_agent_seconds=since_last_agent_seconds,
+                ).to(dtype=recurrent_dtype)
+                assert self.temporal_projector is not None
+                time_tokens = self.temporal_projector(temporal_features)
+                time_tokens = time_tokens + self.segment_embeddings.weight[1].view(1, 1, -1).to(dtype=recurrent_dtype)
+                aux_tokens.append(time_tokens)
+                aux_masks.append(torch.ones(batch_size, self.num_time_tokens, device=z_state.device, dtype=torch.long))
+
+            role_tokens = self.observation_role_embeddings(observation_role_ids).unsqueeze(1).to(dtype=recurrent_dtype)
+            role_tokens = role_tokens + self.segment_embeddings.weight[1].view(1, 1, -1).to(dtype=recurrent_dtype)
+            aux_tokens.append(role_tokens)
+            aux_masks.append(torch.ones(batch_size, 1, device=z_state.device, dtype=torch.long))
+
+            observation_embeds = self.input_embeddings(observation_input_ids).to(dtype=recurrent_dtype)
+            observation_embeds = observation_embeds + self.segment_embeddings.weight[2].view(1, 1, -1).to(
+                dtype=recurrent_dtype
+            )
+            observation_embeds = observation_embeds + self.observation_role_embeddings(observation_role_ids).unsqueeze(1).to(
+                dtype=recurrent_dtype
+            )
+
+            z_tokens = z_state + self.segment_embeddings.weight[0].view(1, 1, -1).to(dtype=recurrent_dtype)
+            encoder_inputs = torch.cat([z_tokens, *aux_tokens, observation_embeds], dim=1)
+
+            z_mask = torch.ones(batch_size, self.z_slots, device=z_state.device, dtype=torch.long)
+            encoder_attention_mask = torch.cat([z_mask, *aux_masks, observation_attention_mask.long()], dim=1)
+
+            # The pretrained T5Gemma backbone was loaded in bf16. Passing fp32 inputs_embeds
+            # into its bf16 Linear layers causes: expected mat1 and mat2 to have same dtype.
+            encoder_outputs = self.encoder(
+                inputs_embeds=encoder_inputs.to(dtype=backbone_dtype),
+                attention_mask=encoder_attention_mask,
+                return_dict=True,
+            )
+            latent_prefix_state = encoder_outputs.last_hidden_state[:, : self.z_slots, :].to(dtype=recurrent_dtype)
+            proposed_state = latent_prefix_state
+            proposal_mask = torch.ones(batch_size, self.z_slots, dtype=torch.long, device=z_state.device)
+        else:
+            use_state_context = (
+                self.observation_encoder_use_state_context
+                and previous_state_mask is not None
+                and torch.any(previous_state_mask > 0)
+            )
+            if use_state_context:
+                observation_embeds = self.input_embeddings(observation_input_ids).to(dtype=recurrent_dtype)
+                encoder_inputs = torch.cat([z_state, observation_embeds], dim=1)
+                encoder_attention_mask = torch.cat(
+                    [
+                        previous_state_mask.to(device=z_state.device, dtype=torch.long),
+                        observation_attention_mask.long(),
+                    ],
+                    dim=1,
+                )
+                observation_encoder_outputs = self.encoder(
+                    inputs_embeds=encoder_inputs.to(dtype=backbone_dtype),
+                    attention_mask=encoder_attention_mask,
+                    return_dict=True,
+                )
+                if self.latent_attention_mask_mode == "full":
+                    proposed_state = observation_encoder_outputs.last_hidden_state[:, : self.z_slots, :].to(
+                        dtype=recurrent_dtype
+                    )
+                    proposal_mask = torch.ones(batch_size, self.z_slots, dtype=torch.long, device=z_state.device)
+                    observation_outputs = None
+                else:
+                    observation_outputs = observation_encoder_outputs.last_hidden_state[:, self.z_slots :, :].to(
+                        dtype=recurrent_dtype
+                    )
+            else:
+                observation_encoder_outputs = self.encoder(
+                    input_ids=observation_input_ids,
+                    attention_mask=observation_attention_mask.long(),
+                    return_dict=True,
+                )
+                observation_outputs = observation_encoder_outputs.last_hidden_state.to(dtype=recurrent_dtype)
+            if observation_outputs is not None:
+                proposed_state, proposal_mask = self.compress_hidden_states_to_slots(
+                    observation_outputs,
+                    observation_attention_mask,
+                    target_slots=self.z_slots,
+                )
+                proposed_state = proposed_state.to(dtype=recurrent_dtype)
+                if self.latent_attention_mask_mode == "full":
+                    proposed_state = proposed_state.clone()
+                    inactive_slot_mask = ~proposal_mask.bool()
+                    proposed_state[inactive_slot_mask] = z_state[inactive_slot_mask]
+                    proposal_mask = torch.ones(batch_size, self.z_slots, dtype=torch.long, device=z_state.device)
+
+            no_observation_mask = proposal_mask.sum(dim=1) <= 0
+            if torch.any(no_observation_mask):
+                proposed_state = proposed_state.clone()
+                proposal_mask = proposal_mask.clone()
+                proposed_state[no_observation_mask] = z_state[no_observation_mask]
+                if previous_state_mask is None:
+                    proposal_mask[no_observation_mask] = 1
+                else:
+                    proposal_mask[no_observation_mask] = previous_state_mask[no_observation_mask].to(
+                        device=z_state.device,
+                        dtype=torch.long,
+                    )
+
+        if not self.preserve_observation_encoder_manifold:
+            proposed_state = self.proposed_state_norm(proposed_state)
+
+        gate_inputs = self.state_gate_input_norm(torch.cat([z_state, proposed_state], dim=-1))
+        update_gate = torch.sigmoid(self.state_gate(gate_inputs))
+        raw_next_state = update_gate * proposed_state + (1.0 - update_gate) * z_state
+        active_slot_mask = proposal_mask.bool().unsqueeze(-1)
+        next_state = torch.where(active_slot_mask, raw_next_state, z_state)
+        if not self.preserve_observation_encoder_manifold:
+            next_state = self.recurrent_state_norm(next_state)
+        if previous_state_mask is not None:
+            next_state_mask = torch.maximum(
+                previous_state_mask.to(device=z_state.device, dtype=torch.long),
+                proposal_mask,
+            )
+        else:
+            next_state_mask = proposal_mask
+
+        pooled_query = self.gate_query.to(dtype=next_state.dtype).expand(batch_size, -1, -1)
+        gate_context = self.gate_context_norm(next_state)
+        pooled_state, _ = self.gate_pool(pooled_query, gate_context, gate_context, need_weights=False)
+        gate_logits = self.gate_head(pooled_state.squeeze(1)).squeeze(-1)
+        return next_state, gate_logits, next_state_mask
+
+    def rollout_step(
+        self,
+        z_state: torch.Tensor,
+        observation_input_ids: torch.Tensor,
+        observation_attention_mask: torch.Tensor,
+        observation_role_ids: torch.Tensor,
+        delta_seconds: torch.Tensor,
+        elapsed_seconds: torch.Tensor,
+        since_last_user_seconds: torch.Tensor,
+        since_last_agent_seconds: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        next_state, gate_logits, _ = self._rollout_step_impl(
+            z_state=z_state,
+            observation_input_ids=observation_input_ids,
+            observation_attention_mask=observation_attention_mask,
+            observation_role_ids=observation_role_ids,
+            delta_seconds=delta_seconds,
+            elapsed_seconds=elapsed_seconds,
+            since_last_user_seconds=since_last_user_seconds,
+            since_last_agent_seconds=since_last_agent_seconds,
+            previous_state_mask=None,
+        )
+        return next_state, gate_logits
+
+    def rollout_step_with_mask(
+        self,
+        z_state: torch.Tensor,
+        state_mask: torch.Tensor,
+        observation_input_ids: torch.Tensor,
+        observation_attention_mask: torch.Tensor,
+        observation_role_ids: torch.Tensor,
+        delta_seconds: torch.Tensor,
+        elapsed_seconds: torch.Tensor,
+        since_last_user_seconds: torch.Tensor,
+        since_last_agent_seconds: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self._rollout_step_impl(
+            z_state=z_state,
+            observation_input_ids=observation_input_ids,
+            observation_attention_mask=observation_attention_mask,
+            observation_role_ids=observation_role_ids,
+            delta_seconds=delta_seconds,
+            elapsed_seconds=elapsed_seconds,
+            since_last_user_seconds=since_last_user_seconds,
+            since_last_agent_seconds=since_last_agent_seconds,
+            previous_state_mask=state_mask,
+        )
+
+    def _resolve_generation_eos_token_id(self) -> int | list[int] | None:
+        eos_token_id = getattr(self.backbone.generation_config, "eos_token_id", None)
+        if eos_token_id is None:
+            eos_token_id = getattr(self.backbone.config, "eos_token_id", None)
+        return eos_token_id
+
+    def _resolve_fill_token_id(self) -> int:
+        eos_token_id = self._resolve_generation_eos_token_id()
+        if isinstance(eos_token_id, list) and eos_token_id:
+            return int(eos_token_id[0])
+        if isinstance(eos_token_id, int):
+            return eos_token_id
+        if self.tokenizer.eos_token_id is not None:
+            return int(self.tokenizer.eos_token_id)
+        if self.tokenizer.pad_token_id is not None:
+            return int(self.tokenizer.pad_token_id)
+        return 0
+
+    def _teacher_forced_decoder_inputs(self, labels: torch.Tensor) -> torch.Tensor:
+        if labels.ndim != 2 or labels.size(0) <= 0:
+            raise ValueError("labels must be a rank-2 tensor with non-zero batch size.")
+        safe_labels = labels.clone()
+        safe_labels[safe_labels == -100] = self._resolve_fill_token_id()
+        if hasattr(self.backbone, "prepare_decoder_input_ids_from_labels"):
+            return self.backbone.prepare_decoder_input_ids_from_labels(labels=safe_labels)
+        return self.backbone._shift_right(safe_labels)
+
+    def _build_self_generated_decoder_inputs(
+        self,
+        labels: torch.Tensor,
+        generated: torch.Tensor,
+        *,
+        self_generated_prefix_tokens: int,
+    ) -> torch.Tensor:
+        batch_size, target_length = labels.shape
+        teacher_inputs = self._teacher_forced_decoder_inputs(labels).to(device=labels.device)
+        if target_length <= 0:
+            return teacher_inputs
+
+        prefix_token_count = min(max(self_generated_prefix_tokens, 0), max(target_length - 1, 0))
+        if prefix_token_count <= 0:
+            return teacher_inputs
+
+        generated_tokens = generated.to(device=labels.device, dtype=torch.long)
+        if generated_tokens.ndim != 2:
+            raise ValueError("generated tokens must be rank-2.")
+        if generated_tokens.size(0) != batch_size:
+            raise ValueError("generated batch size must match labels batch size.")
+        if generated_tokens.size(1) > 0 and torch.equal(generated_tokens[:, :1], teacher_inputs[:, :1]):
+            generated_tokens = generated_tokens[:, 1:]
+
+        fill_token_id = self._resolve_fill_token_id()
+        if generated_tokens.size(1) < prefix_token_count:
+            padding = labels.new_full(
+                (batch_size, prefix_token_count - generated_tokens.size(1)),
+                fill_value=fill_token_id,
+            )
+            generated_tokens = torch.cat([generated_tokens, padding], dim=1)
+        else:
+            generated_tokens = generated_tokens[:, :prefix_token_count]
+
+        hybrid_inputs = teacher_inputs.clone()
+        hybrid_inputs[:, 1 : 1 + prefix_token_count] = generated_tokens
+        return hybrid_inputs
+
+    def decoder_loss(
+        self,
+        z_state: torch.Tensor,
+        labels: torch.Tensor,
+        *,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if z_state.numel() == 0:
+            zero = torch.zeros((), device=self.device)
+            return zero, zero
+        encoder_mask = (
+            encoder_attention_mask.to(device=z_state.device, dtype=torch.long)
+            if encoder_attention_mask is not None
+            else torch.ones(z_state.shape[:2], dtype=torch.long, device=z_state.device)
+        )
+        token_count = (labels != -100).sum()
+        if token_count.item() == 0:
+            zero = torch.zeros((), device=z_state.device)
+            return zero, zero
+        outputs = self.backbone(
+            encoder_outputs=BaseModelOutput(last_hidden_state=z_state.to(dtype=self.backbone_dtype)),
+            attention_mask=encoder_mask,
+            labels=labels,
+            return_dict=True,
+        )
+        loss_sum = outputs.loss * token_count
+        return loss_sum, token_count
+
+    def decoder_self_generated_loss(
+        self,
+        z_state: torch.Tensor,
+        labels: torch.Tensor,
+        *,
+        generation_kwargs: dict[str, Any] | None = None,
+        self_generated_prefix_tokens: int | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if z_state.numel() == 0:
+            zero = torch.zeros((), device=self.device)
+            return zero, zero
+
+        encoder_mask = (
+            encoder_attention_mask.to(device=z_state.device, dtype=torch.long)
+            if encoder_attention_mask is not None
+            else torch.ones(z_state.shape[:2], dtype=torch.long, device=z_state.device)
+        )
+        token_count = (labels != -100).sum()
+        if token_count.item() == 0:
+            zero = torch.zeros((), device=z_state.device)
+            return zero, zero
+
+        prefix_token_count = max(0, int(self_generated_prefix_tokens if self_generated_prefix_tokens is not None else labels.shape[1]))
+        if prefix_token_count <= 0:
+            return self.decoder_loss(z_state, labels, encoder_attention_mask=encoder_attention_mask)
+
+        effective_generation_kwargs = {
+            "max_new_tokens": prefix_token_count,
+            "do_sample": False,
+            "num_beams": 1,
+            "return_dict_in_generate": False,
+        }
+        if generation_kwargs:
+            effective_generation_kwargs.update(generation_kwargs)
+
+        was_training = self.backbone.training
+        if was_training:
+            self.backbone.eval()
+        try:
+            with torch.no_grad():
+                generated = self.backbone.generate(
+                    encoder_outputs=BaseModelOutput(last_hidden_state=z_state.to(dtype=self.backbone_dtype)),
+                    attention_mask=encoder_mask,
+                    **effective_generation_kwargs,
+                )
+        finally:
+            if was_training:
+                self.backbone.train()
+
+        decoder_input_ids = self._build_self_generated_decoder_inputs(
+            labels,
+            generated,
+            self_generated_prefix_tokens=prefix_token_count,
+        )
+        outputs = self.backbone(
+            encoder_outputs=BaseModelOutput(last_hidden_state=z_state.to(dtype=self.backbone_dtype)),
+            attention_mask=encoder_mask,
+            decoder_input_ids=decoder_input_ids,
+            return_dict=True,
+        )
+        logits = outputs.logits.float()
+        loss_sum = F.cross_entropy(
+            logits.reshape(-1, logits.size(-1)),
+            labels.reshape(-1),
+            ignore_index=-100,
+            reduction="sum",
+        )
+        return loss_sum, token_count
+
+    @torch.no_grad()
+    def first_token_exact_match_stats(
+        self,
+        z_state: torch.Tensor,
+        labels: torch.Tensor,
+        *,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if z_state.numel() == 0:
+            zero = torch.zeros((), device=self.device)
+            return zero, zero
+
+        first_logits, valid_mask = self.first_token_logits(
+            z_state,
+            labels,
+            encoder_attention_mask=encoder_attention_mask,
+        )
+        valid_count = valid_mask.sum()
+        if valid_count.item() == 0:
+            zero = torch.zeros((), device=self.device)
+            return zero, zero
+
+        first_targets = labels[:, 0]
+        predicted_first_tokens = first_logits.argmax(dim=-1)
+        correct = ((predicted_first_tokens == first_targets) & valid_mask).sum()
+        return correct, valid_count
+
+    def first_token_logits(
+        self,
+        z_state: torch.Tensor,
+        labels: torch.Tensor,
+        *,
+        encoder_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if z_state.numel() == 0:
+            empty_logits = torch.zeros((0, 0), device=self.device)
+            empty_mask = torch.zeros((0,), dtype=torch.bool, device=self.device)
+            return empty_logits, empty_mask
+
+        first_targets = labels[:, 0]
+        valid_mask = first_targets != -100
+        encoder_mask = (
+            encoder_attention_mask.to(device=z_state.device, dtype=torch.long)
+            if encoder_attention_mask is not None
+            else torch.ones(z_state.shape[:2], dtype=torch.long, device=z_state.device)
+        )
+        decoder_input_ids = self._teacher_forced_decoder_inputs(labels)[:, :1].to(device=z_state.device)
+        outputs = self.backbone(
+            encoder_outputs=BaseModelOutput(last_hidden_state=z_state.to(dtype=self.backbone_dtype)),
+            attention_mask=encoder_mask,
+            decoder_input_ids=decoder_input_ids,
+            return_dict=True,
+        )
+        return outputs.logits[:, 0, :], valid_mask
+
+    @torch.no_grad()
+    def generate_from_state(
+        self,
+        z_state: torch.Tensor,
+        *,
+        encoder_attention_mask: torch.Tensor | None = None,
+        **generation_kwargs: Any,
+    ) -> torch.Tensor:
+        encoder_mask = (
+            encoder_attention_mask.to(device=z_state.device, dtype=torch.long)
+            if encoder_attention_mask is not None
+            else torch.ones(z_state.shape[:2], dtype=torch.long, device=z_state.device)
+        )
+        return self.backbone.generate(
+            encoder_outputs=BaseModelOutput(last_hidden_state=z_state.to(dtype=self.backbone_dtype)),
+            attention_mask=encoder_mask,
+            **generation_kwargs,
+        )
+
+    def save_pretrained(self, output_dir: str | Path, tokenizer: Any | None = None) -> None:
+        output_path = Path(output_dir)
+        output_path.mkdir(parents=True, exist_ok=True)
+        (output_path / "sft_config.json").write_text(
+            json.dumps(self.config, indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+        # The T5Gemma backbone exposes tied/shared weights, which raw save_file refuses.
+        save_model(self, str(output_path / "model.safetensors"))
+        save_file({"z_init": self.z_init.detach().cpu()}, str(output_path / "initial_latent_z.safetensors"))
+        active_tokenizer = tokenizer or self.tokenizer
+        active_tokenizer.save_pretrained(output_path)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path_or_repo_id: str,
+        device: str | torch.device = "cpu",
+        map_location: str | torch.device = "cpu",
+    ) -> "ThoughtLoopT5Gemma":
+        local_path = Path(model_path_or_repo_id)
+        if not local_path.exists():
+            local_path = Path(snapshot_download(repo_id=model_path_or_repo_id))
+
+        config = json.loads((local_path / "sft_config.json").read_text(encoding="utf-8"))
+        model = cls(config)
+        load_model(model, str(local_path / "model.safetensors"), device=str(map_location))
+        model.to(device)
+        return model

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f7b8bf028336d8f0695dc3fd48cbe296efc8501953bca95d3a8727f3a3e86a1
+size 2518360980

sft_b200.yaml

@@ -0,0 +1,112 @@
+dataset:
+  raw_repo_id: BRlkl/timed_natural_conversations_5min4
+  cleaned_repo_id: BRlkl/timed_natural_conversations_5min4
+  cleaned_local_parquet_dir: artifacts/timed_natural_conversations_5min4/parquet
+  source_split: train
+  validation_ratio: 0.02
+  split_seed: 17
+  use_base_chat_template: true
+
+model:
+  base_model_name: google/t5gemma-l-l-prefixlm-it
+  initial_model_path: BRlkl/test_multiturn2
+  dtype: bfloat16
+  attn_implementation: sdpa
+  magicnorm_eps: 1.0e-6
+  z_slots: 256
+  num_time_tokens: 0
+  use_explicit_time_features: false
+  gate_attention_heads: 4
+  max_observation_tokens: 256
+  max_decoder_tokens: 256
+  thought_loop_proposal_mode: observation_hidden_compression
+  preserve_observation_encoder_manifold: true
+  observation_encoder_use_state_context: true
+  latent_attention_mask_mode: full
+  initial_update_gate_bias: 6.0
+
+rollout:
+  tick_seconds: 1.0
+  chunk_ticks: 300
+  max_horizon_ticks: 300
+  sort_by_duration_bucket: false
+  drop_overlong_examples: true
+  long_window_strategy: tail
+  post_speech_feedback_delay_seconds: 1.0
+
+training:
+  seed: 17
+  micro_batch_size: 12
+  eval_batch_size: 1
+  gradient_accumulation_steps: 3
+  num_workers: 0
+  num_train_epochs: 1
+  learning_rate: 5.0e-6
+  weight_decay: 0.05
+  adam_beta1: 0.9
+  adam_beta2: 0.95
+  adam_epsilon: 1.0e-8
+  warmup_ratio: 0.03
+  gradient_checkpointing: true
+  fused_adamw: true
+  mixed_precision: bf16
+  max_grad_norm: 0.1
+  tbptt_start_steps: 300
+  tbptt_end_steps: 300
+  gate_loss_mode: bce
+  gate_soft_target_early_window_ticks: 5
+  gate_soft_target_late_window_ticks: 5
+  gate_soft_target_tau_ticks: 1.5
+  gate_soft_target_shape: gaussian
+  gate_soft_target_peak: 1.0
+  gate_soft_target_floor: 0.0
+  gate_soft_target_block_user_ticks: true
+  hazard_soft_target_window_ticks: 16
+  hazard_soft_target_tau_ticks: 4.0
+  hazard_censor_weight: 0.25
+  use_ranked_gate_objective: false
+  gate_bce_weight: 1.0
+  gate_ranking_weight: 1.0
+  gate_ranking_margin: 0.5
+  gate_ranking_lookback: 64
+  gradual_unfreeze_decoder: false
+  decoder_initial_trainable_fraction: 0.0
+  decoder_unfreeze_start_fraction: 0.0
+  decoder_unfreeze_end_fraction: 0.5
+  freeze_gate_head_first_half: false
+  gate_head_freeze_fraction: 0.5
+  gate_teacher_logit_scale: 20.0
+  gate_positive_weight: 16.0
+  use_dynamic_gate_positive_weight: false
+  dynamic_gate_positive_weight_min: 8.0
+  dynamic_gate_positive_weight_max: 16.0
+  gate_loss_weight: 0.75
+  decoder_loss_weight: 1.0
+  state_delta_weight: 0.1
+  shuffle_train: true
+  sort_dataset_by_length: false
+  logging_steps: 1
+  eval_steps: 51
+  checkpoint_steps: 50
+  eval_max_batches: 12
+
+cache:
+  preprocessed_root: cache/preprocessed_5min
+
+paths:
+  run_root: runs
+  export_root: exports
+
+inference:
+  use_base_chat_template: true
+
+wandb:
+  enabled: true
+  project: samantha-sft-5min
+  run_name: t5gemma2-thoughtloop-5min-from-multiturn3
+
+hub:
+  model_repo_id: BRlkl/test-true2
+  private: false
+  push_checkpoints: false
+  checkpoint_path_prefix: checkpoints

sft_config.json

@@ -0,0 +1,115 @@
+{
+  "dataset": {
+    "raw_repo_id": "BRlkl/timed_natural_conversations_5min4",
+    "cleaned_repo_id": "BRlkl/timed_natural_conversations_5min4",
+    "cleaned_local_parquet_dir": "artifacts/timed_natural_conversations_5min4/parquet",
+    "source_split": "train",
+    "validation_ratio": 0.02,
+    "split_seed": 17,
+    "use_base_chat_template": true
+  },
+  "model": {
+    "base_model_name": "google/t5gemma-l-l-prefixlm-it",
+    "initial_model_path": "BRlkl/test_multiturn2",
+    "dtype": "bfloat16",
+    "attn_implementation": "sdpa",
+    "magicnorm_eps": 1e-06,
+    "z_slots": 256,
+    "num_time_tokens": 0,
+    "use_explicit_time_features": false,
+    "gate_attention_heads": 4,
+    "max_observation_tokens": 256,
+    "max_decoder_tokens": 256,
+    "thought_loop_proposal_mode": "observation_hidden_compression",
+    "preserve_observation_encoder_manifold": true,
+    "observation_encoder_use_state_context": true,
+    "latent_attention_mask_mode": "full",
+    "initial_update_gate_bias": 6.0
+  },
+  "rollout": {
+    "tick_seconds": 1.0,
+    "chunk_ticks": 300,
+    "max_horizon_ticks": 300,
+    "sort_by_duration_bucket": false,
+    "drop_overlong_examples": true,
+    "long_window_strategy": "tail",
+    "post_speech_feedback_delay_seconds": 1.0
+  },
+  "training": {
+    "seed": 17,
+    "micro_batch_size": 12,
+    "eval_batch_size": 1,
+    "gradient_accumulation_steps": 3,
+    "num_workers": 0,
+    "num_train_epochs": 1,
+    "learning_rate": 5e-06,
+    "weight_decay": 0.05,
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.95,
+    "adam_epsilon": 1e-08,
+    "warmup_ratio": 0.03,
+    "gradient_checkpointing": true,
+    "fused_adamw": true,
+    "mixed_precision": "bf16",
+    "max_grad_norm": 0.1,
+    "tbptt_start_steps": 300,
+    "tbptt_end_steps": 300,
+    "gate_loss_mode": "bce",
+    "gate_soft_target_early_window_ticks": 5,
+    "gate_soft_target_late_window_ticks": 5,
+    "gate_soft_target_tau_ticks": 1.5,
+    "gate_soft_target_shape": "gaussian",
+    "gate_soft_target_peak": 1.0,
+    "gate_soft_target_floor": 0.0,
+    "gate_soft_target_block_user_ticks": true,
+    "hazard_soft_target_window_ticks": 16,
+    "hazard_soft_target_tau_ticks": 4.0,
+    "hazard_censor_weight": 0.25,
+    "use_ranked_gate_objective": false,
+    "gate_bce_weight": 1.0,
+    "gate_ranking_weight": 1.0,
+    "gate_ranking_margin": 0.5,
+    "gate_ranking_lookback": 64,
+    "gradual_unfreeze_decoder": false,
+    "decoder_initial_trainable_fraction": 0.0,
+    "decoder_unfreeze_start_fraction": 0.0,
+    "decoder_unfreeze_end_fraction": 0.5,
+    "freeze_gate_head_first_half": false,
+    "gate_head_freeze_fraction": 0.5,
+    "gate_teacher_logit_scale": 20.0,
+    "gate_positive_weight": 16.0,
+    "use_dynamic_gate_positive_weight": false,
+    "dynamic_gate_positive_weight_min": 8.0,
+    "dynamic_gate_positive_weight_max": 16.0,
+    "gate_loss_weight": 0.75,
+    "decoder_loss_weight": 1.0,
+    "state_delta_weight": 0.1,
+    "shuffle_train": true,
+    "sort_dataset_by_length": false,
+    "logging_steps": 1,
+    "eval_steps": 51,
+    "checkpoint_steps": 50,
+    "eval_max_batches": 12
+  },
+  "cache": {
+    "preprocessed_root": "cache/preprocessed_5min"
+  },
+  "paths": {
+    "run_root": "runs",
+    "export_root": "exports"
+  },
+  "inference": {
+    "use_base_chat_template": true
+  },
+  "wandb": {
+    "enabled": true,
+    "project": "samantha-sft-5min",
+    "run_name": "t5gemma2-thoughtloop-5min-from-multiturn3"
+  },
+  "hub": {
+    "model_repo_id": "BRlkl/test-true2",
+    "private": false,
+    "push_checkpoints": false,
+    "checkpoint_path_prefix": "checkpoints"
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:923de64f4705efd22f6880dc5e6a9b2993de2706a03d0d48773e0173e32a074f
+size 34362428

tokenizer_config.json

@@ -0,0 +1,17 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<eos>",
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<mask>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "padding_side": "right",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

train.py

@@ -0,0 +1,1402 @@
+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()