training/train_v2.py

"""
Memory Routing Agent - Training Pipeline v2

Fixes based on engineer feedback and Tinker docs alignment:

1. KL Divergence: Use proper Tinker estimators (kl_sample_train_v1/v2)
2. Reward Function: Use full composite reward from rl_env.py
3. Group Size: Increased to 32 (per Tinker rl/rl-hyperparams.mdx)
4. Batch Size: Increased to 64 groups per batch
5. Checkpointing: Non-blocking saves
6. Advantage Computation: Proper centering within groups (per Tinker rl/rl-loops.mdx)

Per Tinker docs:
- forward_backward_async returns Future, must await .result_async()
- importance_sampling loss requires: target_tokens, logprobs, advantages (all same length)
- save_state() for resumable checkpoints, save_weights_for_sampler() for sampling
- LR scaling: LR ∝ √batch_size
"""

import asyncio
import json
import os
import time
from datetime import datetime
from dataclasses import dataclass, field
from typing import List, Dict, Any, Tuple, Set, Optional
from dotenv import load_dotenv
import numpy as np

load_dotenv()

import tinker
from tinker import types
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.hyperparam_utils import get_lr

# =============================================================================
# CONFIGURATION - Aligned with Tinker recommendations
# =============================================================================

@dataclass
class TrainingConfig:
    # Model
    base_model: str = "meta-llama/Llama-3.1-8B"
    lora_rank: int = 32
    
    # SFT Config - Per Tinker sl-hyperparams.mdx
    sft_steps: int = 100  # Increased for better convergence (aim for 100+ steps)
    sft_batch_size: int = 32
    sft_eval_every: int = 10
    sft_early_stopping_patience: int = 5  # Stop if no improvement for N evals
    sft_min_steps: int = 30  # Minimum steps before early stopping
    sft_gradient_accumulation: int = 1  # Accumulate gradients over N batches
    
    # RL Config - Per Tinker rl/rl-hyperparams.mdx
    rl_iterations: int = 30  # More iterations for convergence
    rl_groups_per_batch: int = 64  # Increased from 32
    rl_group_size: int = 32  # Increased from 4 (per Tinker recommendation)
    rl_learning_rate: float = 2e-5
    rl_temperature: float = 0.7
    rl_max_tokens: int = 100
    rl_kl_threshold: float = 0.01  # Per Tinker: stable with KL < 0.01
    
    # Reward weights - Per PRD Section 4
    reward_f1_weight: float = 0.6
    reward_temp_weight: float = 0.2
    reward_parity_weight: float = 0.1
    reward_efficiency_weight: float = 0.1
    
    # Paths
    train_data: str = "training/processed_data/train_data.json"
    test_data: str = "training/processed_data/test_data.json"
    log_dir: str = field(default_factory=lambda: f"training/logs/run_{datetime.now().strftime('%Y%m%d_%H%M%S')}")


# =============================================================================
# MEMORY TAXONOMY
# =============================================================================

VALID_CATEGORIES = {
    "company.brand_core", "company.strategic_signatures", "company.knowledge_artifacts",
    "company.business_priorities", "company.tools_config", "company.performance_context",
    "user.communication_style", "user.strategic_approach", "user.role_context",
    "user.workflow_patterns", "user.session_history", "user.interaction_preferences",
    "none"
}

CATEGORY_PERSISTENCE = {
    "company.brand_core": "long", "company.strategic_signatures": "long",
    "company.knowledge_artifacts": "long", "company.business_priorities": "short",
    "company.tools_config": "medium", "company.performance_context": "rolling",
    "user.communication_style": "long", "user.strategic_approach": "long",
    "user.role_context": "medium", "user.workflow_patterns": "medium",
    "user.session_history": "short", "user.interaction_preferences": "evolving",
    "none": "short"
}

CATEGORY_SCOPE = {cat: cat.split(".")[0] if "." in cat else "none" for cat in VALID_CATEGORIES}

SYSTEM_PROMPT = """You route marketing conversations into structured memory categories.

Available categories:
- company.brand_core: Voice, values, positioning, identity anchors
- company.strategic_signatures: Decision frameworks, strategic heuristics
- company.knowledge_artifacts: Docs, style guides, playbooks
- company.business_priorities: Quarterly/seasonal goals, active campaigns
- company.tools_config: Integrations, API keys, workflow settings
- company.performance_context: Campaign metrics, retrospectives, learnings
- user.communication_style: Tone, verbosity, format expectations
- user.strategic_approach: Personal priorities, success definitions
- user.role_context: Title, scope, decision authority
- user.workflow_patterns: Review cadence, collaboration norms
- user.session_history: Immediate context, recent asks
- user.interaction_preferences: Coaching style, feedback expectations
- none: Irrelevant, vague, or transactional content

Respond with comma-separated categories. Use 'none' only if no other category applies."""


# =============================================================================
# REWARD COMPUTATION - Full composite reward per PRD
# =============================================================================

@dataclass
class RewardComponents:
    r_f1: float = 0.0
    r_temp: float = 0.0
    r_parity: float = 0.0
    r_eff: float = 0.0
    r_total: float = 0.0
    format_valid: bool = True
    predicted: Set[str] = field(default_factory=set)
    gold: Set[str] = field(default_factory=set)


def parse_categories(text: str) -> Tuple[Set[str], bool]:
    """Parse comma-separated categories from model output."""
    if not text or not text.strip():
        return set(), False
    
    raw_cats = [c.strip().lower() for c in text.split(",")]
    valid_cats = {c for c in raw_cats if c in VALID_CATEGORIES}
    
    if not valid_cats:
        return set(), False
    
    # "none" must be exclusive
    if "none" in valid_cats and len(valid_cats) > 1:
        valid_cats.discard("none")
    
    return valid_cats, True


def compute_f1(predicted: Set[str], gold: Set[str]) -> float:
    """Compute F1 score between predicted and gold category sets."""
    if not predicted and not gold:
        return 1.0
    if not predicted or not gold:
        return 0.0
    
    tp = len(predicted & gold)
    precision = tp / len(predicted) if predicted else 0.0
    recall = tp / len(gold) if gold else 0.0
    
    if precision + recall == 0:
        return 0.0
    return 2 * (precision * recall) / (precision + recall)


def compute_temporal_reward(predicted: Set[str], gold: Set[str]) -> float:
    """Compute temporal alignment reward."""
    if not predicted or not gold:
        return 0.0
    
    from collections import Counter
    
    pred_pers = [CATEGORY_PERSISTENCE.get(c, "medium") for c in predicted]
    gold_pers = [CATEGORY_PERSISTENCE.get(c, "medium") for c in gold]
    
    def majority(items):
        if not items:
            return "medium"
        return Counter(items).most_common(1)[0][0]
    
    pred_p = majority(pred_pers)
    gold_p = majority(gold_pers)
    
    if pred_p == gold_p:
        return 1.0
    
    # Adjacent matches
    adjacent = {
        ("long", "medium"), ("medium", "long"),
        ("medium", "short"), ("short", "medium"),
        ("medium", "rolling"), ("rolling", "medium"),
        ("short", "rolling"), ("rolling", "short"),
    }
    
    if (pred_p, gold_p) in adjacent:
        return 0.5
    return 0.0


def compute_parity_reward(predicted: Set[str], gold: Set[str]) -> float:
    """Compute company/user scope alignment reward."""
    def get_scope(cats):
        scopes = {CATEGORY_SCOPE.get(c, "none") for c in cats}
        if "company" in scopes and "user" in scopes:
            return "mixed"
        elif "company" in scopes:
            return "company"
        elif "user" in scopes:
            return "user"
        return "none"
    
    return 1.0 if get_scope(predicted) == get_scope(gold) else 0.0


def compute_efficiency_reward(predicted: Set[str]) -> float:
    """Compute storage efficiency reward."""
    n = len(predicted)
    if n <= 3:
        return 1.0
    elif n == 4:
        return 0.7
    elif n == 5:
        return 0.4
    return 0.0


def compute_reward(predicted_text: str, gold_categories: List[str], config: TrainingConfig) -> RewardComponents:
    """
    Compute full composite reward.
    R_total = w1*R_F1 + w2*R_temp + w3*R_parity + w4*R_eff
    """
    result = RewardComponents()
    
    predicted, parse_success = parse_categories(predicted_text)
    gold = set(gold_categories)
    
    result.predicted = predicted
    result.gold = gold
    
    if not parse_success:
        result.format_valid = False
        result.r_total = -1.0
        return result
    
    result.r_f1 = compute_f1(predicted, gold)
    result.r_temp = compute_temporal_reward(predicted, gold)
    result.r_parity = compute_parity_reward(predicted, gold)
    result.r_eff = compute_efficiency_reward(predicted)
    
    result.r_total = (
        config.reward_f1_weight * result.r_f1 +
        config.reward_temp_weight * result.r_temp +
        config.reward_parity_weight * result.r_parity +
        config.reward_efficiency_weight * result.r_eff
    )
    
    return result


# =============================================================================
# DATA STRUCTURES - Per Tinker rl/rl-loops.mdx
# =============================================================================

@dataclass
class Rollout:
    """Single rollout from a problem."""
    prompt_tokens: List[int]
    gen_tokens: List[int]
    logprobs: List[float]
    reward: float
    reward_components: RewardComponents
    predicted: str
    gold: List[str]


@dataclass
class RolloutGroup:
    """Group of rollouts for the same problem - per Tinker EnvGroupBuilder pattern."""
    problem_id: int
    rollouts: List[Rollout]
    
    def get_rewards(self) -> List[float]:
        return [r.reward for r in self.rollouts]
    
    def is_constant_reward(self) -> bool:
        rewards = self.get_rewards()
        return len(set(round(r, 4) for r in rewards)) == 1


# =============================================================================
# LOGGING
# =============================================================================

class TrainingLogger:
    def __init__(self, log_dir: str):
        os.makedirs(log_dir, exist_ok=True)
        self.log_dir = log_dir
        self.sft_log = open(os.path.join(log_dir, "sft_metrics.jsonl"), "w")
        self.rl_log = open(os.path.join(log_dir, "rl_metrics.jsonl"), "w")
        self.start_time = time.time()
        
    def log_sft(self, step: int, metrics: Dict):
        metrics["step"] = step
        metrics["elapsed_time"] = time.time() - self.start_time
        self.sft_log.write(json.dumps(metrics) + "\n")
        self.sft_log.flush()
        
        test_str = f"{metrics.get('test_loss', 0):.4f}" if isinstance(metrics.get('test_loss'), (int, float)) else "N/A"
        print(f"[SFT {step:3d}] Loss: {metrics.get('train_loss', 0):.4f} | Test: {test_str} | Time: {metrics.get('step_time', 0):.1f}s", flush=True)
    
    def log_rl(self, iteration: int, metrics: Dict):
        metrics["iteration"] = iteration
        metrics["elapsed_time"] = time.time() - self.start_time
        self.rl_log.write(json.dumps(metrics) + "\n")
        self.rl_log.flush()
        
        print(f"[RL {iteration:3d}] "
              f"Reward: {metrics.get('mean_reward', 0):.3f} (±{metrics.get('std_reward', 0):.3f}) | "
              f"Acc: {metrics.get('accuracy', 0):.1%} | "
              f"KL_v1: {metrics.get('kl_v1', 0):.4f} | "
              f"KL_v2: {metrics.get('kl_v2', 0):.4f} | "
              f"Active: {metrics.get('active_groups', 0)}/{metrics.get('total_groups', 0)} | "
              f"Time: {metrics.get('iter_time', 0):.1f}s", flush=True)
    
    def close(self):
        self.sft_log.close()
        self.rl_log.close()


# =============================================================================
# ADVANTAGE COMPUTATION - Per Tinker rl/rl-loops.mdx
# =============================================================================

def compute_group_advantages(groups: List[RolloutGroup]) -> List[List[float]]:
    """
    Compute advantages by centering rewards within each group.
    
    Per Tinker docs (rl/rl-loops.mdx):
    "We compute advantages by centering rewards within each problem group."
    
    This is the GRPO-style advantage: A(x,y) = r(y) - mean(r(y') for y' in group)
    """
    all_advantages = []
    
    for group in groups:
        rewards = np.array(group.get_rewards())
        mean_reward = rewards.mean()
        
        # Normalize by std for stability (optional but recommended)
        std_reward = rewards.std()
        if std_reward > 1e-8:
            advantages = (rewards - mean_reward) / std_reward
        else:
            advantages = rewards - mean_reward
        
        all_advantages.append(advantages.tolist())
    
    return all_advantages


# =============================================================================
# KL DIVERGENCE ESTIMATORS - Per Tinker rl/rl-hyperparams.mdx
# =============================================================================

def compute_kl_estimators(old_logprobs: List[float], new_logprobs: List[float]) -> Tuple[float, float]:
    """
    Compute KL divergence estimators per Tinker docs.
    
    Per rl/rl-hyperparams.mdx:
    - kl_sample_train_v1: E[log(p/q)] approximation
    - kl_sample_train_v2: Alternative estimator
    
    Both should be >= 0 in expectation. We compute:
    - v1: mean(new_lp - old_lp) -- this is E[log(p_new/p_old)]
    - v2: mean(exp(new_lp - old_lp) - 1 - (new_lp - old_lp)) -- Taylor expansion
    
    Note: These estimate KL(p_new || p_old), not the reverse.
    """
    if not old_logprobs or not new_logprobs:
        return 0.0, 0.0
    
    log_ratios = []
    for old_lp, new_lp in zip(old_logprobs, new_logprobs):
        if old_lp != 0.0:  # Skip prompt tokens
            log_ratios.append(new_lp - old_lp)
    
    if not log_ratios:
        return 0.0, 0.0
    
    log_ratios = np.array(log_ratios)
    
    # v1: Simple mean of log ratios
    # This estimates E_p[log(p/q)] but with samples from q
    # For on-policy, this should be close to 0
    kl_v1 = float(np.mean(log_ratios))
    
    # v2: Unbiased estimator using importance weights
    # E[exp(log_ratio) - 1 - log_ratio] = KL(p||q)
    # This is always >= 0 by Jensen's inequality
    ratios = np.exp(np.clip(log_ratios, -20, 20))  # Clip for numerical stability
    kl_v2 = float(np.mean(ratios - 1 - log_ratios))
    
    return kl_v1, kl_v2


# =============================================================================
# DATUM CONSTRUCTION - Per Tinker losses.mdx
# =============================================================================

def build_rl_datum(rollout: Rollout, advantage: float) -> types.Datum:
    """
    Build a Datum for importance_sampling loss.
    
    Per Tinker losses.mdx, importance_sampling requires:
    - target_tokens: array[(N,), int] - Target token IDs from sampler
    - logprobs: array[(N,), float] - Reference log probabilities from sampler
    - advantages: array[(N,), float] - Advantage values
    
    All must have length N = model_input.length
    """
    prompt_tokens = rollout.prompt_tokens
    gen_tokens = rollout.gen_tokens
    sampler_logprobs = rollout.logprobs
    
    n_prompt = len(prompt_tokens)
    n_gen = len(gen_tokens)
    
    # Full sequence: prompt + generated
    full_tokens = prompt_tokens + gen_tokens
    
    # Model input: all except last token
    input_tokens = full_tokens[:-1]
    
    # Target: all except first token (next-token prediction)
    target_tokens = full_tokens[1:]
    
    n_input = len(input_tokens)
    
    # Logprobs: 0 for prompt positions, actual for generation
    # Sampler logprobs correspond to gen_tokens
    full_logprobs = [0.0] * (n_prompt - 1) + sampler_logprobs
    
    # Advantages: 0 for prompt, actual for generation
    full_advantages = [0.0] * (n_prompt - 1) + [advantage] * n_gen
    
    # Verify lengths match
    assert len(target_tokens) == n_input, f"target_tokens: {len(target_tokens)} vs input: {n_input}"
    assert len(full_logprobs) == n_input, f"logprobs: {len(full_logprobs)} vs input: {n_input}"
    assert len(full_advantages) == n_input, f"advantages: {len(full_advantages)} vs input: {n_input}"
    
    return types.Datum(
        model_input=types.ModelInput.from_ints(input_tokens),
        loss_fn_inputs=dict(
            target_tokens=target_tokens,
            logprobs=full_logprobs,
            advantages=full_advantages
        )
    )


# =============================================================================
# ROLLOUT COLLECTION
# =============================================================================

async def collect_rollouts(
    sampling_client: tinker.SamplingClient,
    renderer: renderers.Renderer,
    train_data: List[Dict],
    config: TrainingConfig
) -> List[RolloutGroup]:
    """
    Collect rollouts organized by problem groups.
    
    Per Tinker rl/rl-loops.mdx:
    - Generate group_size rollouts per unique problem
    - This enables variance reduction via advantage centering
    """
    stop_sequences = renderer.get_stop_sequences()
    params = types.SamplingParams(
        max_tokens=config.rl_max_tokens,
        temperature=config.rl_temperature,
        stop=stop_sequences
    )
    
    # Sample random problems
    n_problems = min(config.rl_groups_per_batch, len(train_data))
    problem_indices = np.random.choice(len(train_data), size=n_problems, replace=False)
    
    rollout_groups = []
    
    for problem_idx in problem_indices:
        example = train_data[problem_idx]
        gold = example.get("categories", [])
        messages = example.get("messages", [])
        
        # Build prompt (exclude assistant response)
        prompt_messages = messages[:-1] if messages else []
        if not prompt_messages:
            continue
        
        prompt = renderer.build_generation_prompt(prompt_messages)
        prompt_tokens = prompt.to_ints()
        
        # Generate group_size rollouts for this problem
        result = sampling_client.sample(
            prompt=prompt,
            sampling_params=params,
            num_samples=config.rl_group_size
        ).result()
        
        rollouts = []
        for seq in result.sequences:
            response, success = renderer.parse_response(seq.tokens)
            predicted = response["content"] if success else ""
            reward_comp = compute_reward(predicted, gold, config)
            
            # Only include if we have logprobs
            if seq.logprobs and len(seq.logprobs) == len(seq.tokens):
                rollouts.append(Rollout(
                    prompt_tokens=prompt_tokens,
                    gen_tokens=seq.tokens,
                    logprobs=seq.logprobs,
                    reward=reward_comp.r_total,
                    reward_components=reward_comp,
                    predicted=predicted,
                    gold=gold
                ))
        
        if rollouts:
            rollout_groups.append(RolloutGroup(
                problem_id=int(problem_idx),
                rollouts=rollouts
            ))
    
    return rollout_groups


def filter_constant_reward_groups(groups: List[RolloutGroup]) -> List[RolloutGroup]:
    """
    Remove groups where all rollouts have the same reward.
    These provide no learning signal (gradient is zero).
    
    Per Tinker rl/rl-loops.mdx:
    "We can optionally filter out groups with all successes or all failures
    as these have policy gradients of zero."
    """
    return [g for g in groups if not g.is_constant_reward()]


# =============================================================================
# SFT PHASE - Per Tinker supervised-learning docs
# =============================================================================

async def run_sft(
    service_client: tinker.ServiceClient,
    training_client: tinker.TrainingClient,
    renderer: renderers.Renderer,
    train_data: List[Dict],
    test_data: List[Dict],
    config: TrainingConfig,
    logger: TrainingLogger
) -> Tuple[str, str]:
    """
    Run SFT phase per Tinker supervised-learning docs.
    
    Improvements based on Tinker docs:
    1. Data shuffling each epoch (sl-basic.mdx)
    2. Early stopping to prevent overfitting (lora-primer.mdx)
    3. Proper LR from get_lr() which accounts for LoRA scaling (sl-hyperparams.mdx)
    4. Overlapping requests for better throughput (async.mdx)
    5. Gradient accumulation support for larger effective batch sizes
    """
    print("\n" + "=" * 70, flush=True)
    print("PHASE 1: SUPERVISED FINE-TUNING", flush=True)
    print("=" * 70, flush=True)
    
    # Get LoRA-adjusted learning rate per sl-hyperparams.mdx
    # LR(m) = lr_base * M_LoRA * (2000/H_m)^P_m
    lr = get_lr(config.base_model)
    
    # Effective batch size for LR scaling
    effective_batch = config.sft_batch_size * config.sft_gradient_accumulation
    
    print(f"Learning rate: {lr:.2e} (LoRA-adjusted)", flush=True)
    print(f"Steps: {config.sft_steps}, Batch size: {config.sft_batch_size}", flush=True)
    print(f"Gradient accumulation: {config.sft_gradient_accumulation}", flush=True)
    print(f"Effective batch size: {effective_batch}", flush=True)
    print(f"Early stopping patience: {config.sft_early_stopping_patience} evals", flush=True)
    print()
    
    # Convert to Datum - per rendering.mdx, use build_supervised_example
    def to_datum(item):
        messages = item.get("messages", [])
        tokens, weights = renderer.build_supervised_example(messages)
        if hasattr(tokens, 'tolist'):
            tokens = tokens.tolist()
        if hasattr(weights, 'tolist'):
            weights = weights.tolist()
        return types.Datum(
            model_input=types.ModelInput.from_ints(tokens[:-1]),
            loss_fn_inputs=dict(target_tokens=tokens[1:], weights=weights[1:])
        )
    
    train_datums = [to_datum(item) for item in train_data]
    test_datums = [to_datum(item) for item in test_data[:50]]
    
    # Count completion tokens for LoRA capacity check (per lora-primer.mdx)
    total_completion_tokens = sum(
        sum(d.loss_fn_inputs['weights'].tolist()) for d in train_datums
    )
    print(f"Total completion tokens: {total_completion_tokens:,}", flush=True)
    print(f"(LoRA works well when completion tokens < LoRA params)", flush=True)
    print()
    
    # Early stopping state
    best_test_loss = float('inf')
    best_checkpoint = None
    patience_counter = 0
    
    # Shuffle indices for each epoch
    indices = list(range(len(train_datums)))
    epoch = 0
    idx_ptr = 0
    
    for step in range(config.sft_steps):
        step_start = time.time()
        
        # Shuffle at epoch boundary
        if idx_ptr + config.sft_batch_size > len(indices):
            np.random.shuffle(indices)
            idx_ptr = 0
            epoch += 1
        
        # Get batch with shuffled indices
        batch_indices = indices[idx_ptr:idx_ptr + config.sft_batch_size]
        idx_ptr += config.sft_batch_size
        batch = [train_datums[i] for i in batch_indices]
        
        # Forward-backward per Tinker async pattern
        # Per async.mdx: "submit your next request while the current one is running"
        fwd_future = await training_client.forward_backward_async(batch, loss_fn="cross_entropy")
        
        # Submit optim step immediately (overlapping with forward-backward)
        optim_future = await training_client.optim_step_async(
            types.AdamParams(learning_rate=lr, beta1=0.9, beta2=0.95, eps=1e-8)
        )
        
        # Now wait for results
        fwd_result = await fwd_future.result_async()
        await optim_future.result_async()
        
        # Compute loss
        logprobs = np.concatenate([o['logprobs'].tolist() for o in fwd_result.loss_fn_outputs])
        weights_arr = np.concatenate([d.loss_fn_inputs['weights'].tolist() for d in batch])
        train_loss = -np.dot(logprobs, weights_arr) / max(weights_arr.sum(), 1)
        
        step_time = time.time() - step_start
        metrics = {
            "train_loss": float(train_loss),
            "step_time": step_time,
            "epoch": epoch,
            "learning_rate": lr
        }
        
        # Evaluate periodically
        if step % config.sft_eval_every == 0 or step == config.sft_steps - 1:
            # Use forward (not forward_backward) for eval to avoid gradient accumulation
            eval_future = await training_client.forward_backward_async(test_datums, loss_fn="cross_entropy")
            eval_result = await eval_future.result_async()
            test_logprobs = np.concatenate([o['logprobs'].tolist() for o in eval_result.loss_fn_outputs])
            test_weights = np.concatenate([d.loss_fn_inputs['weights'].tolist() for d in test_datums])
            test_loss = -np.dot(test_logprobs, test_weights) / max(test_weights.sum(), 1)
            metrics["test_loss"] = float(test_loss)
            
            # Save checkpoint
            save_future = await training_client.save_weights_for_sampler_async(name=f"sft_step_{step:04d}")
            save_result = await save_future.result_async()
            metrics["checkpoint"] = save_result.path
            
            # Early stopping check
            if test_loss < best_test_loss:
                best_test_loss = test_loss
                best_checkpoint = save_result.path
                patience_counter = 0
                metrics["is_best"] = True
            else:
                patience_counter += 1
                metrics["is_best"] = False
            
            metrics["patience_counter"] = patience_counter
            
            # Check early stopping (after minimum steps)
            if step >= config.sft_min_steps and patience_counter >= config.sft_early_stopping_patience:
                print(f"\nEarly stopping at step {step} (no improvement for {patience_counter} evals)", flush=True)
                logger.log_sft(step, metrics)
                break
        
        logger.log_sft(step, metrics)
    
    # Save final state for RL continuation
    # Per save-load.mdx: use save_state() for resumable checkpoints
    state_future = await training_client.save_state_async(name="sft_final")
    state_result = await state_future.result_async()
    
    # Also save sampler weights
    sampler_future = await training_client.save_weights_for_sampler_async(name="sft_final_sampler")
    sampler_result = await sampler_future.result_async()
    
    print(f"\nSFT Complete.", flush=True)
    print(f"  Final checkpoint: {sampler_result.path}", flush=True)
    print(f"  Best checkpoint (loss={best_test_loss:.4f}): {best_checkpoint}", flush=True)
    print(f"  State for RL: {state_result.path}", flush=True)
    
    return state_result.path, sampler_result.path


# =============================================================================
# RL PHASE
# =============================================================================

async def run_rl(
    service_client: tinker.ServiceClient,
    training_client: tinker.TrainingClient,
    sft_state_path: str,
    renderer: renderers.Renderer,
    train_data: List[Dict],
    test_data: List[Dict],
    config: TrainingConfig,
    logger: TrainingLogger
) -> str:
    """
    Run RL phase with proper advantage computation.
    
    Per Tinker rl/rl-loops.mdx:
    1. Create policy with current weights
    2. Generate rollouts (group_size per problem)
    3. Process trajectory data into training examples
    4. Update model parameters
    """
    print("\n" + "=" * 70, flush=True)
    print("PHASE 2: REINFORCEMENT LEARNING", flush=True)
    print("=" * 70, flush=True)
    
    # Load SFT weights
    print(f"Loading SFT state: {sft_state_path}", flush=True)
    await training_client.load_state_async(sft_state_path)
    
    print(f"Iterations: {config.rl_iterations}", flush=True)
    print(f"Groups per batch: {config.rl_groups_per_batch}", flush=True)
    print(f"Group size: {config.rl_group_size}", flush=True)
    print(f"Total rollouts per iteration: {config.rl_groups_per_batch * config.rl_group_size}", flush=True)
    print(f"Learning rate: {config.rl_learning_rate:.2e}", flush=True)
    print(f"KL threshold: {config.rl_kl_threshold}", flush=True)
    print()
    
    best_reward = -float('inf')
    best_checkpoint = None
    
    for iteration in range(config.rl_iterations):
        iter_start = time.time()
        
        # 1. Save current weights for sampling
        save_future = await training_client.save_weights_for_sampler_async(name=f"rl_iter_{iteration:03d}")
        save_result = await save_future.result_async()
        sampling_client = service_client.create_sampling_client(model_path=save_result.path)
        
        # 2. Collect rollouts organized by problem groups
        rollout_groups = await collect_rollouts(
            sampling_client, renderer, train_data, config
        )
        
        # 3. Filter constant-reward groups
        active_groups = filter_constant_reward_groups(rollout_groups)
        
        # Collect all rewards for metrics
        all_rewards = []
        all_reward_components = {"f1": [], "temp": [], "parity": [], "eff": []}
        for group in rollout_groups:
            for rollout in group.rollouts:
                all_rewards.append(rollout.reward)
                if rollout.reward_components.format_valid:
                    all_reward_components["f1"].append(rollout.reward_components.r_f1)
                    all_reward_components["temp"].append(rollout.reward_components.r_temp)
                    all_reward_components["parity"].append(rollout.reward_components.r_parity)
                    all_reward_components["eff"].append(rollout.reward_components.r_eff)
        
        # 4. Compute advantages (centered within groups)
        group_advantages = compute_group_advantages(active_groups)
        
        # 5. Build training data
        training_data = []
        for group, advantages in zip(active_groups, group_advantages):
            for rollout, advantage in zip(group.rollouts, advantages):
                try:
                    datum = build_rl_datum(rollout, advantage)
                    training_data.append((datum, rollout))
                except AssertionError as e:
                    print(f"Warning: Skipping datum: {e}", flush=True)
        
        # 6. Update model and compute KL
        kl_v1_samples = []
        kl_v2_samples = []
        
        if training_data:
            datums = [d[0] for d in training_data]
            
            fwd_future = await training_client.forward_backward_async(
                datums, loss_fn="importance_sampling"
            )
            optim_future = await training_client.optim_step_async(
                types.AdamParams(learning_rate=config.rl_learning_rate, beta1=0.9, beta2=0.95, eps=1e-8)
            )
            
            fwd_result = await fwd_future.result_async()
            await optim_future.result_async()
            
            # Compute KL divergence estimators
            for i, output in enumerate(fwd_result.loss_fn_outputs):
                new_logprobs = output['logprobs'].tolist()
                old_logprobs = datums[i].loss_fn_inputs['logprobs'].tolist()
                v1, v2 = compute_kl_estimators(old_logprobs, new_logprobs)
                kl_v1_samples.append(v1)
                kl_v2_samples.append(v2)
        
        iter_time = time.time() - iter_start
        
        # Compute metrics
        mean_reward = np.mean(all_rewards) if all_rewards else 0
        std_reward = np.std(all_rewards) if all_rewards else 0
        accuracy = sum(1 for r in all_rewards if r > 0) / len(all_rewards) if all_rewards else 0
        kl_v1 = np.mean(kl_v1_samples) if kl_v1_samples else 0
        kl_v2 = np.mean(kl_v2_samples) if kl_v2_samples else 0
        
        metrics = {
            "mean_reward": float(mean_reward),
            "std_reward": float(std_reward),
            "accuracy": accuracy,
            "kl_v1": float(kl_v1),  # Log ratio mean (can be negative)
            "kl_v2": float(kl_v2),  # Unbiased estimator (always >= 0)
            "total_groups": len(rollout_groups),
            "active_groups": len(active_groups),
            "num_training_examples": len(training_data),
            "iter_time": iter_time,
            "checkpoint": save_result.path,
            # Reward component breakdown
            "mean_r_f1": float(np.mean(all_reward_components["f1"])) if all_reward_components["f1"] else 0,
            "mean_r_temp": float(np.mean(all_reward_components["temp"])) if all_reward_components["temp"] else 0,
            "mean_r_parity": float(np.mean(all_reward_components["parity"])) if all_reward_components["parity"] else 0,
            "mean_r_eff": float(np.mean(all_reward_components["eff"])) if all_reward_components["eff"] else 0,
        }
        
        logger.log_rl(iteration, metrics)
        
        # Track best checkpoint
        if mean_reward > best_reward:
            best_reward = mean_reward
            best_checkpoint = save_result.path
        
        # KL threshold warning (using v2 which is always >= 0)
        if kl_v2 > config.rl_kl_threshold:
            print(f"WARNING: KL_v2 {kl_v2:.4f} exceeds threshold {config.rl_kl_threshold}", flush=True)
    
    # Save final checkpoint
    final_future = await training_client.save_weights_for_sampler_async(name="rl_final")
    final_result = await final_future.result_async()
    
    print(f"\nRL Complete. Final: {final_result.path}", flush=True)
    print(f"Best checkpoint (reward={best_reward:.3f}): {best_checkpoint}", flush=True)
    
    return final_result.path


# =============================================================================
# EVALUATION
# =============================================================================

async def evaluate_model(
    service_client: tinker.ServiceClient,
    model_path: str,
    renderer: renderers.Renderer,
    test_data: List[Dict],
    config: TrainingConfig,
    n_samples: int = 100
) -> Dict[str, float]:
    """Evaluate model on test data."""
    print(f"\nEvaluating: {model_path}", flush=True)
    
    sampling_client = service_client.create_sampling_client(model_path=model_path)
    stop_sequences = renderer.get_stop_sequences()
    params = types.SamplingParams(max_tokens=100, temperature=0.1, stop=stop_sequences)
    
    correct_any = 0
    correct_exact = 0
    total_f1 = 0
    total_reward = 0
    
    for item in test_data[:n_samples]:
        messages = item.get("messages", [])
        gold = item.get("categories", [])
        
        prompt_messages = messages[:-1] if messages else []
        if not prompt_messages:
            continue
        
        prompt = renderer.build_generation_prompt(prompt_messages)
        result = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=1).result()
        response, _ = renderer.parse_response(result.sequences[0].tokens)
        pred = response["content"]
        
        reward_comp = compute_reward(pred, gold, config)
        total_reward += reward_comp.r_total
        
        if reward_comp.predicted & reward_comp.gold:
            correct_any += 1
        if reward_comp.predicted == reward_comp.gold:
            correct_exact += 1
        total_f1 += reward_comp.r_f1
    
    n = min(n_samples, len(test_data))
    return {
        "any_match": correct_any / n,
        "exact_match": correct_exact / n,
        "f1": total_f1 / n,
        "mean_reward": total_reward / n
    }


# =============================================================================
# MAIN
# =============================================================================

async def main():
    config = TrainingConfig()
    
    print("=" * 70, flush=True)
    print("MEMORY ROUTING AGENT - TRAINING PIPELINE v2", flush=True)
    print("=" * 70, flush=True)
    print(f"Log directory: {config.log_dir}", flush=True)
    print(f"Model: {config.base_model}", flush=True)
    print(f"RL Groups: {config.rl_groups_per_batch}, Group Size: {config.rl_group_size}", flush=True)
    print()
    
    # Initialize
    service_client = tinker.ServiceClient()
    tokenizer = get_tokenizer(config.base_model)
    renderer = renderers.get_renderer(name="llama3", tokenizer=tokenizer)
    
    # Load data
    with open(config.train_data, "r") as f:
        train_data = json.load(f)
    with open(config.test_data, "r") as f:
        test_data = json.load(f)
    
    print(f"Train: {len(train_data)}, Test: {len(test_data)}", flush=True)
    
    # Create logger
    logger = TrainingLogger(config.log_dir)
    
    # Create training client
    training_client = await service_client.create_lora_training_client_async(
        base_model=config.base_model, rank=config.lora_rank
    )
    
    # Run SFT
    sft_state, sft_sampler = await run_sft(
        service_client, training_client, renderer,
        train_data, test_data, config, logger
    )
    
    # Evaluate SFT
    print("\n" + "-" * 70, flush=True)
    sft_results = await evaluate_model(service_client, sft_sampler, renderer, test_data, config)
    print(f"SFT: Any={sft_results['any_match']:.1%}, Exact={sft_results['exact_match']:.1%}, "
          f"F1={sft_results['f1']:.1%}, Reward={sft_results['mean_reward']:.3f}", flush=True)
    
    # Run RL
    rl_final = await run_rl(
        service_client, training_client, sft_state,
        renderer, train_data, test_data, config, logger
    )
    
    # Evaluate RL
    print("\n" + "-" * 70, flush=True)
    rl_results = await evaluate_model(service_client, rl_final, renderer, test_data, config)
    print(f"RL: Any={rl_results['any_match']:.1%}, Exact={rl_results['exact_match']:.1%}, "
          f"F1={rl_results['f1']:.1%}, Reward={rl_results['mean_reward']:.3f}", flush=True)
    
    logger.close()
    
    # Summary
    print("\n" + "=" * 70, flush=True)
    print("TRAINING COMPLETE", flush=True)
    print("=" * 70, flush=True)
    print(f"Logs: {config.log_dir}", flush=True)
    print(f"SFT: {sft_sampler}", flush=True)
    print(f"RL: {rl_final}", flush=True)
    print()
    print("Performance Comparison:", flush=True)
    print(f"{'Metric':<15} {'SFT':>10} {'RL':>10} {'Delta':>10}", flush=True)
    print("-" * 45, flush=True)
    for metric in ['any_match', 'exact_match', 'f1', 'mean_reward']:
        sft_val = sft_results[metric]
        rl_val = rl_results[metric]
        delta = rl_val - sft_val
        if metric == 'mean_reward':
            print(f"{metric:<15} {sft_val:>10.3f} {rl_val:>10.3f} {delta:>+10.3f}", flush=True)
        else:
            print(f"{metric:<15} {sft_val:>10.1%} {rl_val:>10.1%} {delta:>+10.1%}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())