jfl 0.8.0 → 0.9.0

package/scripts/train/v2/online_train.py

@@ -0,0 +1,576 @@
+"""
+Online Learning Harness for v2 Policy Head.
+
+Implements Drew's recommended hybrid approach (section 8.3):
+  - Experience replay: 70% historical + 30% new data per batch
+  - Small learning rate (1e-5) to avoid catastrophic forgetting
+  - Validation monitoring with automatic rollback if degradation >10%
+  - Continuous checkpointing for recovery
+
+Usage:
+  # Fine-tune on new data with experience replay
+  python online_train.py --new-data .jfl/v2-data/new.jsonl --checkpoint .jfl/checkpoints/best_policy_head.pt
+
+  # Continuous mode: watch for new data and retrain automatically
+  python online_train.py --watch --checkpoint .jfl/checkpoints/best_policy_head.pt
+
+Drew's architecture decision:
+  Pre-training: offline, 3e-4 LR, full dataset
+  Online: continuous, 1e-5 LR, experience replay, validation gating
+  Batch retraining: weekly, full offline, reset to best checkpoint
+"""
+
+import json
+import os
+import sys
+import time
+import math
+import random
+import shutil
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset, ConcatDataset, Subset
+
+from model import PolicyHead
+from dataset import PolicyHeadDataset, load_embedding_cache
+
+
+# ============================================================================
+# Experience Replay Buffer
+# ============================================================================
+
+class ExperienceReplayBuffer:
+    """
+    Maintains a pool of historical training examples for replay.
+
+    Drew's recommendation (section 8.3):
+      Mix new transitions (30%) + sampled historical data (70%) in each batch.
+    """
+
+    def __init__(self, max_size: int = 10000, new_ratio: float = 0.3, seed: int = 42):
+        self.max_size = max_size
+        self.new_ratio = new_ratio
+        self.rng = random.Random(seed)
+        self.historical: list[dict] = []
+        self.new_examples: list[dict] = []
+
+    def add_historical(self, examples: list[dict]):
+        """Add examples to the historical pool."""
+        self.historical.extend(examples)
+        # Reservoir sampling if over max size
+        if len(self.historical) > self.max_size:
+            self.historical = self.rng.sample(self.historical, self.max_size)
+
+    def add_new(self, examples: list[dict]):
+        """Add new examples that will be weighted higher in sampling."""
+        self.new_examples.extend(examples)
+
+    def sample_batch(self, batch_size: int) -> list[dict]:
+        """
+        Sample a mixed batch: new_ratio % new + (1-new_ratio) % historical.
+        If not enough new examples, fill with historical.
+        """
+        n_new = min(
+            int(batch_size * self.new_ratio),
+            len(self.new_examples)
+        )
+        n_historical = batch_size - n_new
+
+        batch = []
+
+        if n_new > 0 and self.new_examples:
+            batch.extend(self.rng.sample(
+                self.new_examples,
+                min(n_new, len(self.new_examples))
+            ))
+
+        if n_historical > 0 and self.historical:
+            batch.extend(self.rng.sample(
+                self.historical,
+                min(n_historical, len(self.historical))
+            ))
+
+        self.rng.shuffle(batch)
+        return batch
+
+    def get_mixed_dataset_indices(
+        self, n_historical: int, n_new: int
+    ) -> tuple[list[int], list[int]]:
+        """Return indices for creating a mixed dataset split."""
+        hist_indices = self.rng.sample(
+            range(len(self.historical)),
+            min(n_historical, len(self.historical))
+        ) if self.historical else []
+
+        new_indices = self.rng.sample(
+            range(len(self.new_examples)),
+            min(n_new, len(self.new_examples))
+        ) if self.new_examples else []
+
+        return hist_indices, new_indices
+
+    @property
+    def total_size(self) -> int:
+        return len(self.historical) + len(self.new_examples)
+
+    def stats(self) -> dict:
+        return {
+            "historical": len(self.historical),
+            "new": len(self.new_examples),
+            "total": self.total_size,
+            "new_ratio": self.new_ratio,
+            "max_size": self.max_size,
+        }
+
+
+# ============================================================================
+# Validation Monitor
+# ============================================================================
+
+class ValidationMonitor:
+    """
+    Tracks validation metrics and triggers rollback on degradation.
+
+    Drew's recommendation:
+      Track performance on held-out test set, rollback if degradation detected.
+      Rollback plan: Automatic rollback if L3 metrics degrade >10%.
+    """
+
+    def __init__(
+        self,
+        degradation_threshold: float = 0.10,
+        patience: int = 3,
+        checkpoint_dir: str = ".jfl/checkpoints",
+    ):
+        self.degradation_threshold = degradation_threshold
+        self.patience = patience
+        self.checkpoint_dir = checkpoint_dir
+        self.baseline_accuracy: float | None = None
+        self.best_accuracy: float = 0.0
+        self.history: list[dict] = []
+        self.degradation_count: int = 0
+
+    def set_baseline(self, accuracy: float):
+        """Set the baseline accuracy from pre-trained model."""
+        self.baseline_accuracy = accuracy
+        self.best_accuracy = accuracy
+        print(f"  Baseline accuracy: {accuracy:.1%}")
+
+    def check(self, epoch: int, val_accuracy: float, val_loss: float) -> dict:
+        """
+        Check if model has degraded beyond threshold.
+
+        Returns:
+          {
+            "action": "continue" | "rollback" | "save_best",
+            "reason": str,
+            "degradation": float,
+          }
+        """
+        self.history.append({
+            "epoch": epoch,
+            "val_accuracy": val_accuracy,
+            "val_loss": val_loss,
+            "timestamp": time.time(),
+        })
+
+        result = {
+            "action": "continue",
+            "reason": "",
+            "degradation": 0.0,
+        }
+
+        if self.baseline_accuracy is None:
+            self.set_baseline(val_accuracy)
+            return result
+
+        # Check for improvement
+        if val_accuracy > self.best_accuracy:
+            self.best_accuracy = val_accuracy
+            self.degradation_count = 0
+            result["action"] = "save_best"
+            result["reason"] = f"New best accuracy: {val_accuracy:.1%} (was {self.best_accuracy:.1%})"
+            return result
+
+        # Check for degradation
+        degradation = (self.baseline_accuracy - val_accuracy) / self.baseline_accuracy
+        result["degradation"] = degradation
+
+        if degradation > self.degradation_threshold:
+            self.degradation_count += 1
+            if self.degradation_count >= self.patience:
+                result["action"] = "rollback"
+                result["reason"] = (
+                    f"Accuracy degraded {degradation:.1%} from baseline "
+                    f"({val_accuracy:.1%} vs {self.baseline_accuracy:.1%}) "
+                    f"for {self.degradation_count} consecutive checks"
+                )
+            else:
+                result["reason"] = (
+                    f"Degradation {degradation:.1%} detected "
+                    f"({self.degradation_count}/{self.patience} until rollback)"
+                )
+        else:
+            self.degradation_count = 0
+
+        return result
+
+    def save_rollback_checkpoint(self, model, optimizer, epoch: int, path: str):
+        """Save a checkpoint that can be rolled back to."""
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        torch.save({
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "baseline_accuracy": self.baseline_accuracy,
+            "best_accuracy": self.best_accuracy,
+        }, path)
+
+
+# ============================================================================
+# Online Training Loop
+# ============================================================================
+
+def online_train(args):
+    """
+    Fine-tune policy head on new data with experience replay.
+
+    Key differences from offline train.py:
+      - Lower learning rate (1e-5 vs 3e-4)
+      - Experience replay (70% historical + 30% new)
+      - Validation monitoring with rollback
+      - Warm-starts from existing checkpoint (required)
+    """
+    # Device
+    if torch.cuda.is_available():
+        device = "cuda"
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = "mps"
+    else:
+        device = "cpu"
+    print(f"Device: {device}")
+
+    # Load existing checkpoint (required for online learning)
+    if not os.path.exists(args.checkpoint):
+        print(f"ERROR: Checkpoint not found: {args.checkpoint}")
+        print("Online learning requires a pre-trained checkpoint.")
+        print("Run offline training first: python train.py")
+        sys.exit(1)
+
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    config = checkpoint.get("config", {})
+    tool_to_index = checkpoint["tool_to_index"]
+    index_to_tool = checkpoint.get("index_to_tool", {str(v): k for k, v in tool_to_index.items()})
+    num_tools = checkpoint.get("num_tools", len(tool_to_index))
+    baseline_accuracy = checkpoint.get("val_accuracy", 0.0)
+
+    print(f"Loaded checkpoint: {args.checkpoint}")
+    print(f"  Baseline val accuracy: {baseline_accuracy:.1%}")
+    print(f"  Tools: {num_tools}")
+
+    # Embeddings
+    embeddings_matrix, text_to_idx = load_embedding_cache(args.data_dir)
+    if embeddings_matrix is not None:
+        print(f"Embedding cache: {embeddings_matrix.shape[0]} texts, {embeddings_matrix.shape[1]}-dim")
+
+    embedding_dim = config.get("embedding_dim", 768)
+
+    # Model
+    model = PolicyHead(
+        embedding_dim=embedding_dim,
+        hidden_dim=config.get("hidden_dim", 512),
+        num_tools=num_tools,
+        num_layers=config.get("num_layers", 4),
+        num_heads=config.get("num_heads", 8),
+        dropout=config.get("dropout", 0.1),
+    ).to(device)
+
+    # Load pre-trained weights
+    model.load_state_dict(checkpoint["model_state_dict"])
+    print(f"  Loaded {model.num_parameters:,} parameters")
+
+    # Load datasets
+    # Historical = existing train data
+    train_path = os.path.join(args.data_dir, "train.jsonl")
+    val_path = os.path.join(args.data_dir, "val.jsonl")
+
+    if not os.path.exists(train_path):
+        print(f"Training data not found: {train_path}")
+        sys.exit(1)
+
+    historical_ds = PolicyHeadDataset(train_path, tool_to_index, embeddings_matrix, text_to_idx)
+    val_ds = PolicyHeadDataset(val_path, tool_to_index, embeddings_matrix, text_to_idx) if os.path.exists(val_path) else None
+
+    # New data (counterfactual + recent real)
+    new_paths = []
+    if args.new_data and os.path.exists(args.new_data):
+        new_paths.append(args.new_data)
+
+    # Also check for counterfactual data
+    cf_path = os.path.join(args.data_dir, "counterfactual.jsonl")
+    if os.path.exists(cf_path) and cf_path not in new_paths:
+        new_paths.append(cf_path)
+
+    new_datasets = []
+    for p in new_paths:
+        ds = PolicyHeadDataset(p, tool_to_index, embeddings_matrix, text_to_idx)
+        if len(ds) > 0:
+            new_datasets.append(ds)
+            print(f"  New data: {p} ({len(ds)} examples)")
+
+    if not new_datasets:
+        print("No new data to train on. Nothing to do.")
+        return
+
+    new_ds = ConcatDataset(new_datasets) if len(new_datasets) > 1 else new_datasets[0]
+
+    # Experience replay: mix historical (70%) + new (30%)
+    replay_ratio = args.replay_ratio
+    n_new = len(new_ds)
+    n_historical = int(n_new * (1 - replay_ratio) / replay_ratio)
+    n_historical = min(n_historical, len(historical_ds))
+
+    print(f"\n  Experience replay:")
+    print(f"    Historical pool: {len(historical_ds)} examples")
+    print(f"    New data: {n_new} examples")
+    print(f"    Sampling: {n_historical} historical + {n_new} new = {n_historical + n_new} total")
+    print(f"    Ratio: {n_new/(n_historical+n_new):.0%} new / {n_historical/(n_historical+n_new):.0%} historical")
+
+    # Create mixed dataset via random sampling
+    rng = random.Random(args.seed)
+    historical_indices = rng.sample(range(len(historical_ds)), n_historical)
+    historical_subset = Subset(historical_ds, historical_indices)
+
+    mixed_ds = ConcatDataset([historical_subset, new_ds])
+
+    num_workers = 0 if device == "mps" else min(4, os.cpu_count() or 1)
+    train_loader = DataLoader(mixed_ds, batch_size=args.batch_size, shuffle=True, num_workers=num_workers)
+    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=num_workers) if val_ds else None
+
+    # Optimizer with low LR (Drew: 1e-5 for online vs 1e-3 for pre-training)
+    optimizer = optim.AdamW(
+        model.parameters(),
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+    )
+
+    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)  # Less smoothing for fine-tuning
+
+    # Validation monitor
+    monitor = ValidationMonitor(
+        degradation_threshold=args.degradation_threshold,
+        patience=args.rollback_patience,
+        checkpoint_dir=args.output_dir,
+    )
+    monitor.set_baseline(baseline_accuracy)
+
+    # Save rollback checkpoint
+    rollback_path = os.path.join(args.output_dir, "rollback_checkpoint.pt")
+    shutil.copy2(args.checkpoint, rollback_path)
+    print(f"  Rollback checkpoint: {rollback_path}")
+
+    # Training loop
+    print(f"\n  Online fine-tuning for {args.epochs} epochs (lr={args.lr})...")
+    print(f"  {'Epoch':>5} {'Train Loss':>12} {'Train Acc':>10} {'Val Loss':>10} {'Val Acc':>9} {'Status':>12}")
+    print("  " + "-" * 70)
+
+    from train import train_epoch, evaluate
+
+    for epoch in range(1, args.epochs + 1):
+        train_loss, train_acc = train_epoch(
+            model, train_loader, criterion, optimizer,
+            # Use constant LR (no scheduler for online)
+            type("FakeScheduler", (), {"step": lambda self: None, "get_last_lr": lambda self: [args.lr]})(),
+            device,
+        )
+
+        val_loss, val_acc = (0.0, 0.0)
+        if val_loader:
+            val_loss, val_acc = evaluate(model, val_loader, criterion, device)
+
+        # Check validation monitor
+        check_result = monitor.check(epoch, val_acc, val_loss)
+        status = check_result["action"]
+
+        status_str = {
+            "continue": "✓",
+            "save_best": "★ best",
+            "rollback": "⚠ ROLLBACK",
+        }.get(status, status)
+
+        print(
+            f"  {epoch:5d} {train_loss:12.4f} {train_acc:9.1%} {val_loss:10.4f} {val_acc:8.1%} {status_str:>12}"
+        )
+
+        if status == "save_best":
+            # Save new best
+            save_checkpoint(model, optimizer, epoch, val_acc, val_loss, config,
+                          tool_to_index, index_to_tool, num_tools,
+                          len(mixed_ds), device, args.output_dir)
+
+        elif status == "rollback":
+            print(f"\n  ⚠ ROLLING BACK: {check_result['reason']}")
+            print(f"  Restoring: {rollback_path}")
+
+            # Restore from rollback checkpoint
+            rollback_ckpt = torch.load(rollback_path, map_location=device, weights_only=False)
+            model.load_state_dict(rollback_ckpt["model_state_dict"])
+            print(f"  Rolled back to baseline accuracy: {baseline_accuracy:.1%}")
+
+            # Write rollback event
+            write_rollback_event(args.output_dir, epoch, check_result)
+            break
+
+    # Final save if we completed without rollback
+    if check_result.get("action") != "rollback":
+        save_checkpoint(model, optimizer, args.epochs, val_acc, val_loss, config,
+                       tool_to_index, index_to_tool, num_tools,
+                       len(mixed_ds), device, args.output_dir)
+        print(f"\n  Online training complete. Final val accuracy: {val_acc:.1%}")
+
+    # Write training event
+    write_training_event(args.output_dir, {
+        "type": "online_training",
+        "epochs": epoch,
+        "final_val_accuracy": val_acc,
+        "baseline_accuracy": baseline_accuracy,
+        "new_examples": n_new,
+        "historical_examples": n_historical,
+        "replay_ratio": replay_ratio,
+        "rollback": check_result.get("action") == "rollback",
+    })
+
+
+def save_checkpoint(model, optimizer, epoch, val_acc, val_loss, config,
+                   tool_to_index, index_to_tool, num_tools, trained_on,
+                   device, output_dir):
+    """Save checkpoint + metadata."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    ckpt_path = os.path.join(output_dir, "best_policy_head.pt")
+    torch.save({
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "val_accuracy": val_acc,
+        "val_loss": val_loss,
+        "num_tools": num_tools,
+        "tool_to_index": tool_to_index,
+        "index_to_tool": index_to_tool,
+        "config": config,
+        "training_mode": "online",
+    }, ckpt_path)
+
+    meta = {
+        "version": 2,
+        "architecture": f"transformer-{config.get('num_layers', 4)}layer-{config.get('hidden_dim', 512)}h",
+        "embedding_dim": config.get("embedding_dim", 768),
+        "hidden_dim": config.get("hidden_dim", 512),
+        "num_tools": num_tools,
+        "trained_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "trained_on": trained_on,
+        "val_accuracy": val_acc,
+        "training_mode": "online",
+        "tool_to_index": tool_to_index,
+        "index_to_tool": {str(k): v for k, v in index_to_tool.items()},
+        "checkpoint_path": os.path.abspath(ckpt_path),
+    }
+    meta_path = os.path.join(output_dir, "policy-head-v2.json")
+    with open(meta_path, "w") as f:
+        json.dump(meta, f, indent=2)
+
+
+def write_rollback_event(output_dir, epoch, check_result):
+    """Record rollback event for monitoring."""
+    event_path = os.path.join(output_dir, "training-events.jsonl")
+    event = {
+        "type": "rollback",
+        "epoch": epoch,
+        "reason": check_result["reason"],
+        "degradation": check_result["degradation"],
+        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+    }
+    with open(event_path, "a") as f:
+        f.write(json.dumps(event) + "\n")
+
+
+def write_training_event(output_dir, data):
+    """Record training event for monitoring."""
+    event_path = os.path.join(output_dir, "training-events.jsonl")
+    data["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+    os.makedirs(output_dir, exist_ok=True)
+    with open(event_path, "a") as f:
+        f.write(json.dumps(data) + "\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Online learning for v2 Policy Head with experience replay"
+    )
+    parser.add_argument(
+        "--checkpoint", required=True,
+        help="Path to pre-trained checkpoint (.pt)"
+    )
+    parser.add_argument(
+        "--data-dir", default=".jfl/v2-data",
+        help="Directory with train/val JSONL + embeddings"
+    )
+    parser.add_argument(
+        "--new-data", default=None,
+        help="Path to new training data JSONL"
+    )
+    parser.add_argument(
+        "--output-dir", default=".jfl/checkpoints",
+        help="Output directory for checkpoints"
+    )
+    parser.add_argument(
+        "--domain", default=None,
+        help="Path to domain.json (uses default if not specified)"
+    )
+    parser.add_argument(
+        "--epochs", type=int, default=10,
+        help="Max fine-tuning epochs (default: 10, much less than offline)"
+    )
+    parser.add_argument(
+        "--batch-size", type=int, default=32,
+        help="Batch size"
+    )
+    parser.add_argument(
+        "--lr", type=float, default=1e-5,
+        help="Learning rate (default: 1e-5, much lower than offline 3e-4)"
+    )
+    parser.add_argument(
+        "--weight-decay", type=float, default=0.01,
+        help="Weight decay"
+    )
+    parser.add_argument(
+        "--replay-ratio", type=float, default=0.3,
+        help="Fraction of batch that is new data (default: 0.3 = 30%% new)"
+    )
+    parser.add_argument(
+        "--degradation-threshold", type=float, default=0.10,
+        help="Rollback if val accuracy drops more than this fraction (default: 0.10 = 10%%)"
+    )
+    parser.add_argument(
+        "--rollback-patience", type=int, default=3,
+        help="Number of consecutive degraded epochs before rollback"
+    )
+    parser.add_argument(
+        "--seed", type=int, default=42,
+        help="Random seed"
+    )
+    args = parser.parse_args()
+
+    if args.domain is None:
+        args.domain = os.path.join(os.path.dirname(os.path.abspath(__file__)), "domain.json")
+
+    online_train(args)
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/precompute.py

@@ -57,21 +57,39 @@ def precompute_embeddings(
     all_states = set()
     all_goals = set()
 
-    for split in ["train", "val", "test"]:
+    for split in ["train", "val", "test", "benchmark", "counterfactual", "synthetic"]:
         path = os.path.join(data_dir, f"{split}.jsonl")
         if not os.path.exists(path):
-            print(f"  Skipping {split} (file not found)")
+            if split in ["train", "val", "test"]:
+                print(f"  Skipping {split} (file not found)")
             continue
         states, goals = collect_unique_texts(path)
         all_states.update(states)
         all_goals.update(goals)
+        if split not in ["train", "val", "test"]:
+            print(f"  Added {split}: {len(states)} states, {len(goals)} goals")
 
     all_texts = sorted(all_states | all_goals)
     print(f"Unique texts to embed: {len(all_texts)} ({len(all_states)} states, {len(all_goals)} goals)")
 
+    # Load existing cache to avoid re-embedding
     text_to_embedding = {}
-    for i in range(0, len(all_texts), batch_size):
-        batch = all_texts[i : i + batch_size]
+    cache_path = os.path.join(data_dir, "embeddings_cache.npz")
+    index_path = os.path.join(data_dir, "text_to_idx.json")
+    if os.path.exists(cache_path) and os.path.exists(index_path):
+        existing_idx = json.load(open(index_path))
+        existing_emb = np.load(cache_path, allow_pickle=True)["embeddings"]
+        for text, idx in existing_idx.items():
+            if idx < len(existing_emb):
+                text_to_embedding[text] = existing_emb[idx].tolist()
+        print(f"  Loaded {len(text_to_embedding)} cached embeddings")
+
+    # Only embed new texts
+    new_texts = [t for t in all_texts if t not in text_to_embedding]
+    print(f"  New texts to embed: {len(new_texts)} (cached: {len(text_to_embedding)})")
+
+    for i in range(0, len(new_texts), batch_size):
+        batch = new_texts[i : i + batch_size]
         try:
             embeddings = embedder(batch)
             for text, emb in zip(batch, embeddings):
@@ -80,8 +98,8 @@ def precompute_embeddings(
             print(f"  Error embedding batch {i}-{i + len(batch)}: {e}")
             continue
 
-        done = min(i + batch_size, len(all_texts))
-        print(f"  Embedded {done}/{len(all_texts)} texts")
+        done = min(i + batch_size, len(new_texts))
+        print(f"  Embedded {done}/{len(new_texts)} new texts")
 
     texts_list = sorted(text_to_embedding.keys())
     text_to_idx = {t: i for i, t in enumerate(texts_list)}