jfl 0.4.4 → 0.6.0

package/scripts/train/v2/model.py

@@ -0,0 +1,112 @@
+"""
+v2 Policy Head — Transformer-based action selector.
+
+Architecture from Drew's Stratus tutorial:
+  (current_state_emb, goal_emb) -> state_proj + goal_proj -> fusion -> TransformerEncoder -> classifier -> action logits
+
+~8.7M params, ~17MB checkpoint. Replaces v1 MLP reward predictor.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PolicyHead(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int = 768,
+        hidden_dim: int = 512,
+        num_tools: int = 12,
+        num_layers: int = 4,
+        num_heads: int = 8,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.num_tools = num_tools
+
+        self.state_proj = nn.Linear(embedding_dim, hidden_dim)
+        self.goal_proj = nn.Linear(embedding_dim, hidden_dim)
+
+        self.fusion = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim * 4,
+            dropout=dropout,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=num_layers,
+        )
+
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_tools),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        current_state_emb: torch.Tensor,
+        goal_state_emb: torch.Tensor,
+    ) -> torch.Tensor:
+        state_h = self.state_proj(current_state_emb)
+        goal_h = self.goal_proj(goal_state_emb)
+
+        fused = self.fusion(torch.cat([state_h, goal_h], dim=-1))
+
+        x = fused.unsqueeze(1)
+        x = self.transformer(x)
+        x = x.squeeze(1)
+
+        x = self.norm(x)
+        logits = self.classifier(x)
+
+        return logits
+
+    def predict(
+        self,
+        current_state_emb: torch.Tensor,
+        goal_state_emb: torch.Tensor,
+        top_k: int = 3,
+    ) -> dict:
+        self.eval()
+        with torch.no_grad():
+            logits = self.forward(current_state_emb, goal_state_emb)
+            probs = F.softmax(logits, dim=-1)
+            top_probs, top_indices = torch.topk(probs, k=min(top_k, self.num_tools), dim=-1)
+
+        return {
+            "top_k_indices": top_indices,
+            "top_k_probs": top_probs,
+            "all_probs": probs,
+        }
+
+    @property
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)

package/scripts/train/v2/precompute.py

@@ -0,0 +1,132 @@
+"""
+Pre-compute Stratus embeddings for all unique texts in v2 training data.
+Caches embeddings as .npz files to avoid re-computation during training.
+"""
+
+import json
+import os
+import sys
+import argparse
+import numpy as np
+
+def get_stratus_embedder(api_url: str, api_key: str):
+    import requests
+
+    def embed_batch(texts: list[str]) -> list[list[float]]:
+        response = requests.post(
+            f"{api_url}/v1/embeddings",
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            },
+            json={
+                "model": "stratus-x1ac-base",
+                "input": texts,
+            },
+            timeout=30,
+        )
+        response.raise_for_status()
+        data = response.json()
+        return [d["embedding"] for d in data["data"]]
+
+    return embed_batch
+
+def collect_unique_texts(data_path: str) -> tuple[list[str], list[str]]:
+    states = set()
+    goals = set()
+
+    with open(data_path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            ex = json.loads(line)
+            states.add(ex["current_state"])
+            goals.add(ex["goal"])
+
+    return sorted(states), sorted(goals)
+
+def precompute_embeddings(
+    data_dir: str,
+    api_url: str,
+    api_key: str,
+    batch_size: int = 32,
+):
+    embedder = get_stratus_embedder(api_url, api_key)
+
+    all_states = set()
+    all_goals = set()
+
+    for split in ["train", "val", "test"]:
+        path = os.path.join(data_dir, f"{split}.jsonl")
+        if not os.path.exists(path):
+            print(f"  Skipping {split} (file not found)")
+            continue
+        states, goals = collect_unique_texts(path)
+        all_states.update(states)
+        all_goals.update(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)")
+
+    text_to_embedding = {}
+    for i in range(0, len(all_texts), batch_size):
+        batch = all_texts[i : i + batch_size]
+        try:
+            embeddings = embedder(batch)
+            for text, emb in zip(batch, embeddings):
+                text_to_embedding[text] = emb
+        except Exception as e:
+            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")
+
+    texts_list = sorted(text_to_embedding.keys())
+    text_to_idx = {t: i for i, t in enumerate(texts_list)}
+    embeddings_matrix = np.array([text_to_embedding[t] for t in texts_list], dtype=np.float32)
+
+    cache_path = os.path.join(data_dir, "embeddings_cache.npz")
+    np.savez(
+        cache_path,
+        embeddings=embeddings_matrix,
+        texts=np.array(texts_list, dtype=object),
+    )
+    print(f"Saved embedding cache: {cache_path} ({embeddings_matrix.shape})")
+
+    index_path = os.path.join(data_dir, "text_to_idx.json")
+    with open(index_path, "w") as f:
+        json.dump(text_to_idx, f)
+    print(f"Saved text index: {index_path} ({len(text_to_idx)} entries)")
+
+    return text_to_idx, embeddings_matrix
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Pre-compute Stratus embeddings for v2 training data")
+    parser.add_argument("--data-dir", default=".jfl/v2-data", help="Directory with train/val/test JSONL files")
+    parser.add_argument("--batch-size", type=int, default=32, help="Embedding batch size")
+    args = parser.parse_args()
+
+    api_url = os.environ.get("STRATUS_API_URL", "https://api.stratus.run")
+    api_key = os.environ.get("STRATUS_API_KEY", "")
+
+    if not api_key:
+        print("STRATUS_API_KEY not set")
+        sys.exit(1)
+
+    if not os.path.exists(args.data_dir):
+        print(f"Data directory not found: {args.data_dir}")
+        sys.exit(1)
+
+    precompute_embeddings(
+        data_dir=args.data_dir,
+        api_url=api_url,
+        api_key=api_key,
+        batch_size=args.batch_size,
+    )
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/train.py

@@ -0,0 +1,302 @@
+"""
+v2 Policy Head Training Loop.
+
+CrossEntropyLoss with label smoothing, cosine annealing with warmup,
+early stopping. Produces .pt checkpoint with model weights, config, and tool index.
+"""
+
+import json
+import os
+import sys
+import time
+import math
+import argparse
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+
+from model import PolicyHead
+from dataset import PolicyHeadDataset, load_embedding_cache
+
+
+def load_tool_index(domain_path: str) -> dict[str, int]:
+    with open(domain_path) as f:
+        domain = json.load(f)
+    return {tool["name"]: i for i, tool in enumerate(domain["tools"])}
+
+
+def get_lr_scheduler(optimizer, warmup_steps: int, total_steps: int):
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return float(step) / float(max(1, warmup_steps))
+        progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+        return max(0.0, 0.5 * (1.0 + math.cos(progress * math.pi)))
+
+    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def train_epoch(model, dataloader, criterion, optimizer, scheduler, device):
+    model.train()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+
+    for batch in dataloader:
+        state_emb = batch["state_emb"].to(device)
+        goal_emb = batch["goal_emb"].to(device)
+        labels = batch["label"].to(device)
+
+        optimizer.zero_grad()
+
+        logits = model(state_emb, goal_emb)
+        loss = criterion(logits, labels)
+
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+        optimizer.step()
+        scheduler.step()
+
+        total_loss += loss.item() * labels.size(0)
+        preds = logits.argmax(dim=-1)
+        correct += (preds == labels).sum().item()
+        total += labels.size(0)
+
+    return total_loss / max(total, 1), correct / max(total, 1)
+
+
+@torch.no_grad()
+def evaluate(model, dataloader, criterion, device):
+    model.eval()
+    total_loss = 0.0
+    correct = 0
+    total = 0
+
+    for batch in dataloader:
+        state_emb = batch["state_emb"].to(device)
+        goal_emb = batch["goal_emb"].to(device)
+        labels = batch["label"].to(device)
+
+        logits = model(state_emb, goal_emb)
+        loss = criterion(logits, labels)
+
+        total_loss += loss.item() * labels.size(0)
+        preds = logits.argmax(dim=-1)
+        correct += (preds == labels).sum().item()
+        total += labels.size(0)
+
+    return total_loss / max(total, 1), correct / max(total, 1)
+
+
+def train(args):
+    # 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}")
+
+    # Domain
+    domain_path = args.domain
+    tool_to_index = load_tool_index(domain_path)
+    index_to_tool = {v: k for k, v in tool_to_index.items()}
+    num_tools = len(tool_to_index)
+    print(f"Tools: {num_tools}")
+
+    # Embeddings cache
+    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")
+    else:
+        print("WARNING: No embedding cache found. Training with zero vectors.")
+        print("  Run: python precompute.py --data-dir", args.data_dir)
+
+    # Datasets
+    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)
+
+    train_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
+
+    num_workers = 0 if device == "mps" else min(4, os.cpu_count() or 1)
+    train_loader = DataLoader(train_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
+
+    print(f"Train: {len(train_ds)} examples")
+    if val_ds:
+        print(f"Val:   {len(val_ds)} examples")
+
+    # Model
+    embedding_dim = embeddings_matrix.shape[1] if embeddings_matrix is not None else 768
+    model = PolicyHead(
+        embedding_dim=embedding_dim,
+        hidden_dim=args.hidden_dim,
+        num_tools=num_tools,
+        num_layers=args.num_layers,
+        num_heads=args.num_heads,
+        dropout=args.dropout,
+    ).to(device)
+
+    print(f"Parameters: {model.num_parameters:,}")
+
+    # Warm start
+    if args.warm_start and os.path.exists(args.warm_start):
+        print(f"Warm-starting from: {args.warm_start}")
+        state_dict = torch.load(args.warm_start, map_location=device, weights_only=True)
+        if "model_state_dict" in state_dict:
+            state_dict = state_dict["model_state_dict"]
+        compatible = {}
+        for k, v in state_dict.items():
+            if k in model.state_dict() and v.shape == model.state_dict()[k].shape:
+                compatible[k] = v
+        model.load_state_dict(compatible, strict=False)
+        print(f"  Loaded {len(compatible)}/{len(state_dict)} layers")
+
+    # Loss, optimizer, scheduler
+    criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing)
+    optimizer = optim.AdamW(
+        model.parameters(),
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+    )
+    total_steps = len(train_loader) * args.epochs
+    scheduler = get_lr_scheduler(optimizer, args.warmup_steps, total_steps)
+
+    # Training loop
+    os.makedirs(args.output_dir, exist_ok=True)
+    best_val_acc = 0.0
+    best_val_loss = float("inf")
+    patience_counter = 0
+
+    print(f"\nStarting training for {args.epochs} epochs...")
+    print(f"{'Epoch':>5} {'Train Loss':>12} {'Train Acc':>10} {'Val Loss':>10} {'Val Acc':>9} {'LR':>10} {'Time':>8}")
+    print("-" * 75)
+
+    for epoch in range(1, args.epochs + 1):
+        t0 = time.time()
+
+        train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, scheduler, device)
+
+        val_loss, val_acc = (0.0, 0.0)
+        if val_loader:
+            val_loss, val_acc = evaluate(model, val_loader, criterion, device)
+
+        elapsed = time.time() - t0
+        lr = scheduler.get_last_lr()[0]
+
+        print(
+            f"{epoch:5d} {train_loss:12.4f} {train_acc:9.1%} {val_loss:10.4f} {val_acc:8.1%} {lr:10.2e} {elapsed:7.1f}s"
+        )
+
+        # Save best model
+        is_best = False
+        if val_loader:
+            if val_acc > best_val_acc:
+                best_val_acc = val_acc
+                best_val_loss = val_loss
+                is_best = True
+        elif train_loss < best_val_loss:
+            best_val_loss = train_loss
+            best_val_acc = train_acc
+            is_best = True
+
+        if is_best:
+            patience_counter = 0
+            checkpoint = {
+                "epoch": epoch,
+                "model_state_dict": model.state_dict(),
+                "val_accuracy": best_val_acc,
+                "val_loss": best_val_loss,
+                "num_tools": num_tools,
+                "tool_to_index": tool_to_index,
+                "index_to_tool": index_to_tool,
+                "config": {
+                    "embedding_dim": embedding_dim,
+                    "hidden_dim": args.hidden_dim,
+                    "num_layers": args.num_layers,
+                    "num_heads": args.num_heads,
+                    "dropout": args.dropout,
+                },
+            }
+            ckpt_path = os.path.join(args.output_dir, "best_policy_head.pt")
+            torch.save(checkpoint, ckpt_path)
+            print(f"  ✅ New best model (val_acc={best_val_acc:.1%})")
+        else:
+            patience_counter += 1
+
+        # Early stopping
+        if patience_counter >= args.patience:
+            print(f"\n  Early stopping at epoch {epoch} (no improvement for {args.patience} epochs)")
+            break
+
+    print(f"\nTraining complete. Best val accuracy: {best_val_acc:.1%}")
+
+    ckpt_path = os.path.join(args.output_dir, "best_policy_head.pt")
+    if os.path.exists(ckpt_path):
+        size_mb = os.path.getsize(ckpt_path) / 1024 / 1024
+        print(f"Checkpoint: {ckpt_path} ({size_mb:.1f} MB)")
+
+    # Write metadata for TypeScript bridge
+    meta = {
+        "version": 2,
+        "architecture": "transformer-4layer-512h",
+        "embedding_dim": embedding_dim,
+        "hidden_dim": args.hidden_dim,
+        "num_tools": num_tools,
+        "num_layers": args.num_layers,
+        "num_heads": args.num_heads,
+        "trained_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "trained_on": len(train_ds),
+        "val_accuracy": best_val_acc,
+        "val_loss": best_val_loss,
+        "device": device,
+        "parameters": model.num_parameters,
+        "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(args.output_dir, "policy-head-v2.json")
+    with open(meta_path, "w") as f:
+        json.dump(meta, f, indent=2)
+    print(f"Metadata: {meta_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train v2 policy head (transformer action selector)")
+    parser.add_argument("--data-dir", default=".jfl/v2-data", help="Directory with train/val/test JSONL + embeddings")
+    parser.add_argument("--domain", default=None, help="Path to domain.json")
+    parser.add_argument("--output-dir", default=".jfl/checkpoints", help="Output directory for checkpoints")
+    parser.add_argument("--warm-start", default=None, help="Path to pretrained checkpoint for warm start")
+    parser.add_argument("--epochs", type=int, default=50, help="Max training epochs")
+    parser.add_argument("--batch-size", type=int, default=64, help="Batch size")
+    parser.add_argument("--lr", type=float, default=3e-4, help="Learning rate")
+    parser.add_argument("--weight-decay", type=float, default=0.01, help="Weight decay")
+    parser.add_argument("--warmup-steps", type=int, default=100, help="Warmup steps for LR scheduler")
+    parser.add_argument("--patience", type=int, default=7, help="Early stopping patience")
+    parser.add_argument("--label-smoothing", type=float, default=0.1, help="Label smoothing for CrossEntropyLoss")
+    parser.add_argument("--hidden-dim", type=int, default=512, help="Hidden dimension")
+    parser.add_argument("--num-layers", type=int, default=4, help="Transformer encoder layers")
+    parser.add_argument("--num-heads", type=int, default=8, help="Attention heads")
+    parser.add_argument("--dropout", type=float, default=0.1, help="Dropout rate")
+    args = parser.parse_args()
+
+    if args.domain is None:
+        args.domain = os.path.join(os.path.dirname(os.path.abspath(__file__)), "domain.json")
+
+    if not os.path.exists(args.domain):
+        print(f"Domain file not found: {args.domain}")
+        sys.exit(1)
+
+    train(args)
+
+
+if __name__ == "__main__":
+    main()