@10et/cli 1.0.0

package/scripts/train/v2/precompute.py

@@ -0,0 +1,150 @@
+"""
+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", "benchmark", "counterfactual", "synthetic"]:
+        path = os.path.join(data_dir, f"{split}.jsonl")
+        if not os.path.exists(path):
+            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 = {}
+    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):
+                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(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)}
+    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()

package/scripts/train/v2/transform_buffer.py

@@ -0,0 +1,227 @@
+"""
+Transform v1 training buffer (state, action, reward) to v2 format (current_state, goal, correct_tool).
+
+Reads .jfl/training-buffer.jsonl and .jfl/agents/*.toml to produce v2 training data
+with 70/15/15 train/val/test splits.
+"""
+
+import json
+import os
+import sys
+import random
+import argparse
+from pathlib import Path
+
+ACTION_TYPE_MAP = {
+    "fix": "fix_bug",
+    "refactor": "refactor_code",
+    "feature": "add_feature",
+    "test": "add_tests",
+    "config": "update_config",
+    "experiment": "run_experiment",
+}
+
+def load_domain(domain_path: str) -> dict:
+    with open(domain_path) as f:
+        return json.load(f)
+
+def load_agent_goals(agents_dir: str) -> dict[str, str]:
+    goals = {}
+    agents_path = Path(agents_dir)
+    if not agents_path.exists():
+        return goals
+
+    for toml_file in agents_path.glob("*.toml"):
+        name = None
+        metric = None
+        direction = None
+        scope = None
+
+        for line in toml_file.read_text().splitlines():
+            line = line.strip()
+            if line.startswith("name"):
+                name = line.split("=", 1)[1].strip().strip('"')
+            elif line.startswith("metric"):
+                metric = line.split("=", 1)[1].strip().strip('"')
+            elif line.startswith("direction"):
+                direction = line.split("=", 1)[1].strip().strip('"')
+            elif line.startswith("scope") and "=" in line and not line.startswith("scope_files"):
+                scope = line.split("=", 1)[1].strip().strip('"')
+
+        if name and metric and direction:
+            verb = "Improve" if direction == "maximize" else "Reduce"
+            scope_str = f" for {scope}" if scope else ""
+            goals[name] = f"{verb} {metric.replace('_', ' ')}{scope_str}"
+
+    return goals
+
+def format_state_text(state: dict) -> str:
+    dims = state.get("dimension_scores", {})
+    dims_str = ", ".join(f"{k}={v:.4f}" for k, v in dims.items()) if dims else "none"
+    deltas = state.get("recent_deltas", [])
+    deltas_str = ", ".join(f"{'+' if d >= 0 else ''}{d:.4f}" for d in deltas) if deltas else "none"
+    return "\n".join([
+        f"Agent: {state.get('agent', 'unknown')}",
+        f"Composite: {state.get('composite_score', 0):.4f}",
+        f"Tests: {state.get('tests_passing', 0)}/{state.get('tests_total', 0)}",
+        f"Trajectory: {state.get('trajectory_length', 0)}",
+        f"Dimensions: {dims_str}",
+        f"Recent deltas: {deltas_str}",
+    ])
+
+def map_action_type(action_type: str) -> str | None:
+    return ACTION_TYPE_MAP.get(action_type)
+
+def transform_entry(entry: dict, agent_goals: dict, valid_tools: set) -> dict | None:
+    state = entry.get("state", {})
+    action = entry.get("action", {})
+    reward = entry.get("reward", {})
+
+    action_type = action.get("type", "")
+    tool_name = map_action_type(action_type)
+    if not tool_name or tool_name not in valid_tools:
+        return None
+
+    agent_name = entry.get("agent", state.get("agent", ""))
+    goal = agent_goals.get(agent_name)
+    if not goal:
+        metric_hint = action.get("description", "")[:80]
+        goal = f"Improve codebase quality: {metric_hint}" if metric_hint else "Improve overall codebase quality"
+
+    current_state = format_state_text(state)
+
+    return {
+        "current_state": current_state,
+        "goal": goal,
+        "correct_tool": tool_name,
+        "source": "buffer",
+        "agent": agent_name,
+        "improved": reward.get("improved", False),
+        "reward_delta": reward.get("composite_delta", 0),
+    }
+
+def transform_buffer(
+    buffer_path: str,
+    agents_dir: str,
+    domain_path: str,
+    output_dir: str,
+    improved_only: bool = True,
+    seed: int = 42,
+) -> dict:
+    domain = load_domain(domain_path)
+    valid_tools = {t["name"] for t in domain["tools"]}
+    agent_goals = load_agent_goals(agents_dir)
+
+    entries = []
+    with open(buffer_path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entries.append(json.loads(line))
+            except json.JSONDecodeError:
+                continue
+
+    print(f"Loaded {len(entries)} raw training buffer entries")
+    if agent_goals:
+        print(f"Agent goals: {json.dumps(agent_goals, indent=2)}")
+
+    examples = []
+    skipped_type = 0
+    skipped_improved = 0
+
+    for entry in entries:
+        transformed = transform_entry(entry, agent_goals, valid_tools)
+        if transformed is None:
+            skipped_type += 1
+            continue
+
+        if improved_only and not transformed["improved"]:
+            skipped_improved += 1
+            continue
+
+        examples.append({
+            "current_state": transformed["current_state"],
+            "goal": transformed["goal"],
+            "correct_tool": transformed["correct_tool"],
+            "source": transformed["source"],
+        })
+
+    print(f"Transformed: {len(examples)} examples")
+    print(f"Skipped: {skipped_type} (unmapped type), {skipped_improved} (not improved)")
+
+    if not examples:
+        print("No examples to write. Check action types and improved filter.")
+        return {"train": 0, "val": 0, "test": 0}
+
+    random.seed(seed)
+    random.shuffle(examples)
+
+    n = len(examples)
+    train_end = int(n * 0.7)
+    val_end = int(n * 0.85)
+
+    splits = {
+        "train": examples[:train_end],
+        "val": examples[train_end:val_end],
+        "test": examples[val_end:],
+    }
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    counts = {}
+    for split_name, split_data in splits.items():
+        path = os.path.join(output_dir, f"{split_name}.jsonl")
+        with open(path, "w") as f:
+            for ex in split_data:
+                f.write(json.dumps(ex) + "\n")
+        counts[split_name] = len(split_data)
+        print(f"  {split_name}: {len(split_data)} examples -> {path}")
+
+    # Class balance report
+    from collections import Counter
+    tool_counts = Counter(ex["correct_tool"] for ex in examples)
+    print(f"\nClass distribution:")
+    for tool, count in tool_counts.most_common():
+        pct = count / len(examples) * 100
+        print(f"  {tool:25s} {count:4d} ({pct:5.1f}%)")
+
+    return counts
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Transform v1 training buffer to v2 format")
+    parser.add_argument("--buffer", default=".jfl/training-buffer.jsonl", help="Path to training buffer JSONL")
+    parser.add_argument("--agents", default=".jfl/agents", help="Path to agents directory")
+    parser.add_argument("--domain", default=None, help="Path to domain.json")
+    parser.add_argument("--output", default=".jfl/v2-data", help="Output directory for split data")
+    parser.add_argument("--all", action="store_true", help="Include entries where improved=false")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for split")
+    args = parser.parse_args()
+
+    # Default domain path: relative to this script
+    domain_path = args.domain
+    if domain_path is None:
+        domain_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "domain.json")
+
+    if not os.path.exists(args.buffer):
+        print(f"Training buffer not found: {args.buffer}")
+        sys.exit(1)
+
+    if not os.path.exists(domain_path):
+        print(f"Domain file not found: {domain_path}")
+        sys.exit(1)
+
+    transform_buffer(
+        buffer_path=args.buffer,
+        agents_dir=args.agents,
+        domain_path=domain_path,
+        output_dir=args.output,
+        improved_only=not args.all,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()