@elizaos/training 2.0.0-alpha.11

package/python/pyproject.toml

@@ -0,0 +1,41 @@
+[project]
+name = "elizaos-training"
+version = "0.2.0"
+description = "RL training pipeline for ElizaOS agents using Atropos"
+authors = [
+    {name = "ElizaOS Contributors"}
+]
+requires-python = ">=3.10"
+dependencies = [
+    "atroposlib>=0.1.0",
+    "asyncpg>=0.29.0",
+    "python-dotenv>=1.0.0",
+    "pydantic>=2.0.0",
+    "numpy>=1.24.0",
+    "openai>=1.0.0",
+    "litellm>=1.0.0",
+    "tabulate>=0.9.0",
+    "vllm>=0.4.0",
+    "transformers>=4.40.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0.0",
+    "pytest-asyncio>=0.21.0",
+    "black>=23.0.0",
+    "ruff>=0.1.0",
+]
+
+[project.scripts]
+train-mmo = "scripts.train_mmo:main"
+check-windows = "scripts.check_windows:main"
+run-migrations = "scripts.run_migrations:main"
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.black]
+line-length = 100
+target-version = ["py310"]

package/python/requirements-ci.txt

@@ -0,0 +1,31 @@
+# Lightweight dependencies for CI test runs (avoids GPU / vLLM installs)
+atroposlib>=0.3.0
+
+# Database (integration tests)
+asyncpg>=0.29.0
+psycopg2-binary>=2.9.9
+
+# HTTP/API
+httpx>=0.26.0
+aiohttp>=3.9.0
+requests>=2.31.0
+openai>=1.0.0
+
+# Config / typing
+python-dotenv>=1.0.0
+pydantic>=2.5.0
+pyyaml>=6.0.1
+
+# Testing
+pytest>=7.4.0
+pytest-asyncio>=0.21.0
+
+# Utilities used by the training pipeline
+wandb>=0.16.0
+tqdm>=4.66.0
+psutil>=5.9.0
+numpy>=1.24.0
+tenacity>=8.2.0
+rich>=13.0.0
+jsonlines>=4.0.0
+

package/python/requirements.txt

@@ -0,0 +1,87 @@
+# Babylon RL Training - Atropos + Tinker Framework
+# Supports: Tinker Cloud (RECOMMENDED), CUDA (GPU), MLX (Apple Silicon), CPU
+
+# ===========================================
+# Tinker API (RECOMMENDED - Cloud Training)
+# ===========================================
+# Tinker provides cloud-based training without local GPU requirements
+# Get API key from: https://tinker-docs.thinkingmachines.ai/
+tinker>=0.1.0
+
+# ===========================================
+# Core Atropos Framework (for environments)
+# ===========================================
+atroposlib>=0.3.0
+
+# ===========================================
+# Database
+# ===========================================
+asyncpg>=0.29.0
+psycopg2-binary>=2.9.9
+
+# ===========================================
+# HTTP/API
+# ===========================================
+httpx>=0.26.0
+aiohttp>=3.9.0
+requests>=2.31.0
+
+# ===========================================
+# OpenAI-compatible client (for RLAIF judge)
+# ===========================================
+openai>=1.0.0
+
+# ===========================================
+# Configuration
+# ===========================================
+pyyaml>=6.0.1
+python-dotenv>=1.0.0
+pydantic>=2.5.0
+pydantic-cli>=3.0.0
+
+# ===========================================
+# Testing
+# ===========================================
+pytest>=7.4.0
+pytest-asyncio>=0.21.0
+
+# ===========================================
+# Experiment Tracking (Optional)
+# ===========================================
+wandb>=0.16.0  # W&B integration, falls back to offline mode if no API key
+
+# ===========================================
+# Utilities
+# ===========================================
+tqdm>=4.66.0
+psutil>=5.9.0
+numpy>=1.24.0
+tenacity>=8.2.0
+rich>=13.0.0
+jsonlines>=4.0.0
+
+# ============================================
+# OPTIONAL: Local Training (GPU/CPU)
+# ============================================
+# Only needed if NOT using Tinker for training
+# Uncomment if you need local fallback:
+#
+# torch>=2.1.0
+# transformers>=4.36.0
+# peft>=0.8.0
+# vllm>=0.3.0
+# accelerate>=1.12.0
+
+# ============================================
+# OPTIONAL: MLX Backend (Apple Silicon only)
+# ============================================
+# Install on Mac with Apple Silicon:
+#   pip install mlx mlx-lm
+#
+# For fine-tuning support:
+#   pip install mlx-lm[finetuning]
+#
+# Recommended models for MLX:
+#   - mlx-community/Qwen2.5-3B-Instruct-4bit
+#   - mlx-community/Qwen2.5-7B-Instruct-4bit
+#   - mlx-community/Qwen3-4B-4bit

package/python/scripts/__init__.py

@@ -0,0 +1,4 @@
+"""Training scripts for Babylon RL system"""
+
+
+

package/python/scripts/import_json_trajectories.py

@@ -0,0 +1,412 @@
+#!/usr/bin/env python3
+"""
+Import JSON Trajectories to PostgreSQL
+
+Bridges JSON-mode generated trajectories to the database for integration testing
+and production seeding. This enables testing the database pipeline with data
+generated by the TypeScript simulation engine.
+
+Usage:
+    # Import from default training-data-output directory
+    python scripts/import_json_trajectories.py
+    
+    # Import from custom directory
+    python scripts/import_json_trajectories.py --source ./my-trajectories
+    
+    # Dry run (validate without inserting)
+    python scripts/import_json_trajectories.py --dry-run
+    
+    # Verbose output
+    python scripts/import_json_trajectories.py --verbose
+
+Environment:
+    DATABASE_URL: PostgreSQL connection string (required)
+
+Requirements:
+    - psycopg2
+    - JSON trajectory files generated by generate-training-data.ts
+"""
+
+import argparse
+import json
+import logging
+import os
+import sys
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional
+import hashlib
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.data_bridge.reader import JsonTrajectoryReader, validate_llm_calls
+from src.training.rubric_loader import normalize_archetype, has_custom_rubric
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ImportStats:
+    """Statistics for import operation."""
+    total_files: int = 0
+    valid_trajectories: int = 0
+    invalid_trajectories: int = 0
+    inserted: int = 0
+    skipped_existing: int = 0
+    failed: int = 0
+    archetypes_seen: Dict[str, int] = None
+    
+    def __post_init__(self):
+        if self.archetypes_seen is None:
+            self.archetypes_seen = {}
+    
+    def record_archetype(self, archetype: str):
+        normalized = normalize_archetype(archetype)
+        self.archetypes_seen[normalized] = self.archetypes_seen.get(normalized, 0) + 1
+
+
+def get_db_connection():
+    """Get database connection from environment."""
+    database_url = os.environ.get("DATABASE_URL")
+    if not database_url:
+        raise ValueError(
+            "DATABASE_URL environment variable not set. "
+            "Example: postgresql://babylon:password@localhost:5433/babylon"
+        )
+    
+    try:
+        import psycopg2
+    except ImportError:
+        raise ImportError(
+            "psycopg2 is required for database import. "
+            "Install with: pip install psycopg2-binary"
+        )
+    
+    return psycopg2.connect(database_url)
+
+
+def generate_snowflake_id() -> str:
+    """Generate a unique ID similar to TypeScript snowflake."""
+    timestamp = int(datetime.now().timestamp() * 1000)
+    random_part = hashlib.sha256(os.urandom(8)).hexdigest()[:8]
+    return f"{timestamp}{random_part}"
+
+
+def extract_archetype_from_trajectory(traj_data: Dict) -> str:
+    """
+    Extract archetype from trajectory data.
+    
+    Priority:
+    1. trajectory.archetype field
+    2. First step's action.parameters.archetype
+    3. Default to 'trader'
+    """
+    # Check trajectory-level archetype
+    archetype = traj_data.get("archetype")
+    if archetype and archetype != "default":
+        return normalize_archetype(archetype)
+    
+    # Check steps for archetype in action parameters
+    steps_json = traj_data.get("stepsJson", "[]")
+    steps = json.loads(steps_json) if isinstance(steps_json, str) else steps_json
+    
+    for step in steps:
+        action = step.get("action", {})
+        params = action.get("parameters", {})
+        step_archetype = params.get("archetype")
+        if step_archetype:
+            return normalize_archetype(step_archetype)
+    
+    return "trader"
+
+
+def validate_trajectory(traj_data: Dict) -> tuple[bool, List[str]]:
+    """
+    Validate a trajectory for import.
+    
+    Returns:
+        (is_valid, list_of_issues)
+    """
+    issues = []
+    
+    # Required fields
+    required = ["trajectoryId", "agentId", "windowId"]
+    for field in required:
+        if not traj_data.get(field):
+            issues.append(f"Missing required field: {field}")
+    
+    # Steps validation
+    steps_json = traj_data.get("stepsJson", "[]")
+    try:
+        steps = json.loads(steps_json) if isinstance(steps_json, str) else steps_json
+        if len(steps) == 0:
+            issues.append("No steps in trajectory")
+        else:
+            is_valid_llm, llm_issues = validate_llm_calls(steps)
+            if not is_valid_llm:
+                issues.extend(llm_issues)
+    except json.JSONDecodeError as e:
+        issues.append(f"Invalid stepsJson: {e}")
+    
+    # Archetype validation
+    archetype = extract_archetype_from_trajectory(traj_data)
+    if archetype != "default" and not has_custom_rubric(archetype):
+        issues.append(f"Unknown archetype: {archetype}")
+    
+    return len(issues) == 0, issues
+
+
+def check_trajectory_exists(conn, trajectory_id: str) -> bool:
+    """Check if trajectory already exists in database."""
+    cur = conn.cursor()
+    cur.execute(
+        'SELECT 1 FROM trajectories WHERE "trajectoryId" = %s LIMIT 1',
+        (trajectory_id,)
+    )
+    exists = cur.fetchone() is not None
+    cur.close()
+    return exists
+
+
+def insert_trajectory(conn, traj_data: Dict) -> bool:
+    """
+    Insert a trajectory into the database.
+    
+    Returns:
+        True if inserted, False if failed
+    """
+    trajectory_id = traj_data.get("trajectoryId")
+    archetype = extract_archetype_from_trajectory(traj_data)
+    
+    # Parse timestamps
+    start_time = traj_data.get("startTime")
+    if isinstance(start_time, str):
+        start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
+    elif start_time is None:
+        start_time = datetime.now()
+    
+    end_time = traj_data.get("endTime")
+    if isinstance(end_time, str):
+        end_time = datetime.fromisoformat(end_time.replace("Z", "+00:00"))
+    elif end_time is None:
+        end_time = datetime.now()
+    
+    # Prepare values
+    values = (
+        traj_data.get("id", generate_snowflake_id()),
+        trajectory_id,
+        traj_data.get("agentId"),
+        archetype,
+        traj_data.get("windowId"),
+        traj_data.get("scenarioId"),
+        traj_data.get("stepsJson", "[]"),
+        json.dumps(traj_data.get("rewardComponents", {})),
+        json.dumps(traj_data.get("metrics", {})),
+        json.dumps(traj_data.get("metadata", {})),
+        float(traj_data.get("finalPnL", 0)),
+        int(traj_data.get("episodeLength", 0)),
+        float(traj_data.get("totalReward", 0)),
+        traj_data.get("finalStatus", "completed"),
+        float(traj_data.get("finalBalance")) if traj_data.get("finalBalance") else None,
+        int(traj_data.get("tradesExecuted", 0)),
+        int(traj_data.get("postsCreated", 0)),
+        float(traj_data.get("aiJudgeReward")) if traj_data.get("aiJudgeReward") else None,
+        traj_data.get("isTrainingData", True),
+        traj_data.get("isEvaluation", False),
+        traj_data.get("usedInTraining", False),
+        start_time,
+        end_time,
+        int(traj_data.get("durationMs", 0)),
+        1,  # windowHours
+        datetime.now(),
+        datetime.now(),
+    )
+    
+    cur = conn.cursor()
+    cur.execute(
+        '''
+        INSERT INTO trajectories (
+            "id", "trajectoryId", "agentId", "archetype", "windowId",
+            "scenarioId", "stepsJson", "rewardComponentsJson", "metricsJson",
+            "metadataJson", "finalPnL", "episodeLength", "totalReward",
+            "finalStatus", "finalBalance", "tradesExecuted", "postsCreated",
+            "aiJudgeReward", "isTrainingData", "isEvaluation", "usedInTraining",
+            "startTime", "endTime", "durationMs", "windowHours",
+            "createdAt", "updatedAt"
+        ) VALUES (
+            %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
+        )
+        ''',
+        values
+    )
+    conn.commit()
+    cur.close()
+    return True
+
+
+def import_trajectories(
+    source_dir: Path,
+    dry_run: bool = False,
+    verbose: bool = False,
+) -> ImportStats:
+    """
+    Import all trajectories from source directory to database.
+    
+    Args:
+        source_dir: Directory containing trajectory JSON files
+        dry_run: If True, validate but don't insert
+        verbose: If True, log each trajectory
+        
+    Returns:
+        Import statistics
+    """
+    stats = ImportStats()
+    
+    # Find trajectory files
+    traj_dir = source_dir / "trajectories"
+    if not traj_dir.exists():
+        logger.warning(f"Trajectories directory not found: {traj_dir}")
+        traj_dir = source_dir
+    
+    json_files = list(traj_dir.glob("*.json"))
+    stats.total_files = len(json_files)
+    
+    if stats.total_files == 0:
+        logger.warning(f"No JSON files found in {traj_dir}")
+        return stats
+    
+    logger.info(f"Found {stats.total_files} trajectory files in {traj_dir}")
+    
+    # Get database connection (skip if dry run)
+    conn = None
+    if not dry_run:
+        conn = get_db_connection()
+        logger.info("Connected to database")
+    
+    # Process each file
+    for file_path in json_files:
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                data = json.load(f)
+            
+            # Handle wrapped format (trajectory key) vs direct format
+            traj_data = data.get("trajectory", data)
+            trajectory_id = traj_data.get("trajectoryId", file_path.stem)
+            
+            # Validate
+            is_valid, issues = validate_trajectory(traj_data)
+            if not is_valid:
+                stats.invalid_trajectories += 1
+                if verbose:
+                    logger.warning(f"Invalid trajectory {trajectory_id}: {issues}")
+                continue
+            
+            stats.valid_trajectories += 1
+            archetype = extract_archetype_from_trajectory(traj_data)
+            stats.record_archetype(archetype)
+            
+            if verbose:
+                logger.info(f"Validated: {trajectory_id} (archetype: {archetype})")
+            
+            if dry_run:
+                continue
+            
+            # Check if exists
+            if check_trajectory_exists(conn, trajectory_id):
+                stats.skipped_existing += 1
+                if verbose:
+                    logger.info(f"Skipped existing: {trajectory_id}")
+                continue
+            
+            # Insert
+            try:
+                insert_trajectory(conn, traj_data)
+                stats.inserted += 1
+                if verbose:
+                    logger.info(f"Inserted: {trajectory_id}")
+            except Exception as e:
+                stats.failed += 1
+                logger.error(f"Failed to insert {trajectory_id}: {e}")
+                
+        except json.JSONDecodeError as e:
+            stats.invalid_trajectories += 1
+            logger.warning(f"Invalid JSON in {file_path}: {e}")
+        except Exception as e:
+            stats.failed += 1
+            logger.error(f"Error processing {file_path}: {e}")
+    
+    if conn:
+        conn.close()
+    
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Import JSON trajectories to PostgreSQL database"
+    )
+    parser.add_argument(
+        "--source",
+        type=Path,
+        default=Path("./training-data-output"),
+        help="Source directory containing trajectory JSON files"
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Validate trajectories without inserting to database"
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Log each trajectory processed"
+    )
+    
+    args = parser.parse_args()
+    
+    if not args.source.exists():
+        logger.error(f"Source directory not found: {args.source}")
+        sys.exit(1)
+    
+    if args.dry_run:
+        logger.info("DRY RUN MODE - No database modifications")
+    
+    stats = import_trajectories(
+        source_dir=args.source,
+        dry_run=args.dry_run,
+        verbose=args.verbose,
+    )
+    
+    # Print summary
+    print("\n" + "=" * 50)
+    print("IMPORT SUMMARY")
+    print("=" * 50)
+    print(f"Total files:          {stats.total_files}")
+    print(f"Valid trajectories:   {stats.valid_trajectories}")
+    print(f"Invalid trajectories: {stats.invalid_trajectories}")
+    
+    if not args.dry_run:
+        print(f"Inserted:             {stats.inserted}")
+        print(f"Skipped (existing):   {stats.skipped_existing}")
+        print(f"Failed:               {stats.failed}")
+    
+    if stats.archetypes_seen:
+        print("\nArchetypes found:")
+        for archetype, count in sorted(stats.archetypes_seen.items()):
+            print(f"  - {archetype}: {count}")
+    
+    if stats.failed > 0:
+        sys.exit(1)
+    
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()
+

package/python/scripts/local-finetune/README.md

@@ -0,0 +1,63 @@
+# Local Fine-Tuning Pipeline
+
+This directory contains scripts to train RL adapters from Babylon simulation logs.
+
+## Workflow
+
+1. **Generate Data:** Run `bun packages/engine/examples/generate-training-data.ts`
+2. **Score & Format:** Run `python ingest_and_score.py`
+3. **Train:** Run `python train_from_csv.py`
+4. **Test:** Run `python test_adapter.py`
+
+## Quick Start
+
+If you do not have a local Postgres database, Atropos server, or vLLM instance running, you can use the **Offline Pipeline**. This generates data to JSON files and uses direct PyTorch/HuggingFace libraries for training.
+
+### Prerequisites
+
+1. `GROQ_API_KEY` or `OPENAI_API_KEY` set in environment.
+2. Python dependencies: `pip install torch transformers peft pandas datasets trl`
+
+### Step 1: Generate Data (TypeScript)
+
+Runs the game simulation in-memory and dumps "Observation -> Action" logs to JSON.
+
+```bash
+# Runs 24 simulated hours
+bun packages/engine/examples/generate-training-data.ts
+```
+
+_Output:_ `training-data-output/trajectories/*.json`
+
+### Step 2: Process & Score (Python)
+
+Converts raw JSON logs into a scored CSV dataset (System/User/Assistant format).
+
+```bash
+cd packages/training/python/scripts/local-finetune
+python ingest_and_score.py
+```
+
+_Output:_ `packages/training/data/scored_trajectories.csv`
+
+### Step 3: Train Model (Python)
+
+Fine-tunes a base model (Qwen2.5-0.5B by default) on your scored data using LoRA.
+
+```bash
+python train_from_csv.py --output ./my-model-v1
+```
+
+### Step 4: Test Inference
+
+Interactively chat with your new LoRA adapter to verify behavior.
+
+```bash
+python test_adapter.py
+```
+
+---
+
+## 🏗️ Production Architecture (Tinker/Atropos)
+
+_For the full cloud-based pipeline involving Postgres, GRPO, and Tinker compute, refer to `scripts/run_full_pipeline.py`._