@elizaos/training 2.0.0-alpha.11

package/python/scripts/merge_trajectories.py

@@ -0,0 +1,318 @@
+#!/usr/bin/env python3
+"""
+Merge Trajectories from Multiple Workers
+
+Combines trajectory files from parallel generation workers into a single
+output directory, handling deduplication and validation.
+
+Usage:
+    python scripts/merge_trajectories.py ./training-data-output
+    python scripts/merge_trajectories.py ./training-data-output --output ./merged
+    python scripts/merge_trajectories.py ./training-data-output --validate
+
+Requirements:
+    - Trajectory JSON files from generate_dataset.sh
+"""
+
+import argparse
+import hashlib
+import json
+import logging
+import os
+import shutil
+import sys
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Set
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MergeStats:
+    """Statistics for merge operation"""
+    total_files_found: int = 0
+    valid_trajectories: int = 0
+    duplicate_trajectories: int = 0
+    invalid_trajectories: int = 0
+    merged_trajectories: int = 0
+    archetypes: Dict[str, int] = None
+    
+    def __post_init__(self):
+        if self.archetypes is None:
+            self.archetypes = {}
+    
+    def record_archetype(self, archetype: str):
+        self.archetypes[archetype] = self.archetypes.get(archetype, 0) + 1
+
+
+def generate_content_hash(data: Dict) -> str:
+    """Generate a hash of trajectory content for deduplication"""
+    # Use trajectory ID if available
+    if "trajectoryId" in data:
+        return data["trajectoryId"]
+    
+    # Otherwise hash the content
+    content = json.dumps(data, sort_keys=True)
+    return hashlib.sha256(content.encode()).hexdigest()[:16]
+
+
+def validate_trajectory(data: Dict) -> tuple[bool, List[str]]:
+    """
+    Validate trajectory data structure.
+    
+    Returns:
+        (is_valid, list_of_issues)
+    """
+    issues = []
+    
+    # Handle wrapped format
+    if "trajectory" in data:
+        data = data["trajectory"]
+    
+    # Check required fields
+    required = ["trajectoryId", "agentId"]
+    for field in required:
+        if not data.get(field):
+            issues.append(f"Missing field: {field}")
+    
+    # Check steps
+    steps = data.get("stepsJson", "[]")
+    if isinstance(steps, str):
+        try:
+            steps = json.loads(steps)
+        except json.JSONDecodeError:
+            issues.append("Invalid stepsJson")
+            steps = []
+    
+    if len(steps) == 0:
+        issues.append("No steps in trajectory")
+    
+    return len(issues) == 0, issues
+
+
+def extract_archetype(data: Dict) -> str:
+    """Extract archetype from trajectory data"""
+    if "trajectory" in data:
+        data = data["trajectory"]
+    
+    archetype = data.get("archetype")
+    if archetype and archetype != "default":
+        return archetype
+    
+    # Try to extract from steps
+    steps = data.get("stepsJson", "[]")
+    if isinstance(steps, str):
+        try:
+            steps = json.loads(steps)
+        except json.JSONDecodeError:
+            return "default"
+    
+    for step in steps:
+        action = step.get("action", {})
+        params = action.get("parameters", {})
+        if params.get("archetype"):
+            return params["archetype"]
+    
+    return "default"
+
+
+def find_trajectory_files(source_dir: Path) -> List[Path]:
+    """Find all trajectory JSON files in source directory"""
+    files = []
+    
+    # Check for batch_N directories
+    batch_dirs = list(source_dir.glob("batch_*/trajectories"))
+    
+    if batch_dirs:
+        for batch_dir in batch_dirs:
+            files.extend(batch_dir.glob("*.json"))
+    else:
+        # Check for direct trajectories directory
+        traj_dir = source_dir / "trajectories"
+        if traj_dir.exists():
+            files.extend(traj_dir.glob("*.json"))
+        else:
+            # Check source dir itself
+            files.extend(source_dir.glob("*.json"))
+    
+    return sorted(files)
+
+
+def merge_trajectories(
+    source_dir: Path,
+    output_dir: Path,
+    validate: bool = True,
+    dry_run: bool = False,
+) -> MergeStats:
+    """
+    Merge trajectories from multiple workers.
+    
+    Args:
+        source_dir: Directory containing batch_N subdirectories
+        output_dir: Output directory for merged trajectories
+        validate: Whether to validate trajectories before merging
+        dry_run: If True, don't actually copy files
+    
+    Returns:
+        Merge statistics
+    """
+    stats = MergeStats()
+    seen_hashes: Set[str] = set()
+    
+    # Find all trajectory files
+    files = find_trajectory_files(source_dir)
+    stats.total_files_found = len(files)
+    
+    if stats.total_files_found == 0:
+        logger.warning(f"No trajectory files found in {source_dir}")
+        return stats
+    
+    logger.info(f"Found {stats.total_files_found} trajectory files")
+    
+    # Create output directory
+    output_traj_dir = output_dir / "trajectories"
+    if not dry_run:
+        output_traj_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Process each file
+    for file_path in files:
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                data = json.load(f)
+        except json.JSONDecodeError as e:
+            logger.warning(f"Invalid JSON in {file_path}: {e}")
+            stats.invalid_trajectories += 1
+            continue
+        
+        # Validate if requested
+        if validate:
+            is_valid, issues = validate_trajectory(data)
+            if not is_valid:
+                logger.debug(f"Invalid trajectory {file_path}: {issues}")
+                stats.invalid_trajectories += 1
+                continue
+        
+        stats.valid_trajectories += 1
+        
+        # Check for duplicates
+        content_hash = generate_content_hash(data)
+        if content_hash in seen_hashes:
+            stats.duplicate_trajectories += 1
+            continue
+        seen_hashes.add(content_hash)
+        
+        # Record archetype
+        archetype = extract_archetype(data)
+        stats.record_archetype(archetype)
+        
+        # Copy to output
+        if not dry_run:
+            output_file = output_traj_dir / file_path.name
+            
+            # Handle name collisions
+            if output_file.exists():
+                base = file_path.stem
+                suffix = file_path.suffix
+                counter = 1
+                while output_file.exists():
+                    output_file = output_traj_dir / f"{base}_{counter}{suffix}"
+                    counter += 1
+            
+            shutil.copy2(file_path, output_file)
+        
+        stats.merged_trajectories += 1
+    
+    return stats
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Merge trajectories from multiple generation workers"
+    )
+    parser.add_argument(
+        "source_dir",
+        type=Path,
+        help="Source directory containing batch_N subdirectories"
+    )
+    parser.add_argument(
+        "--output", "-o",
+        type=Path,
+        default=None,
+        help="Output directory (default: source_dir/merged)"
+    )
+    parser.add_argument(
+        "--validate",
+        action="store_true",
+        help="Validate trajectories before merging"
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Show what would be done without copying files"
+    )
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Verbose output"
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    if not args.source_dir.exists():
+        logger.error(f"Source directory not found: {args.source_dir}")
+        sys.exit(1)
+    
+    output_dir = args.output or (args.source_dir / "merged")
+    
+    if args.dry_run:
+        logger.info("DRY RUN MODE - No files will be copied")
+    
+    logger.info(f"Merging trajectories from {args.source_dir}")
+    logger.info(f"Output directory: {output_dir}")
+    
+    stats = merge_trajectories(
+        source_dir=args.source_dir,
+        output_dir=output_dir,
+        validate=args.validate,
+        dry_run=args.dry_run,
+    )
+    
+    # Print summary
+    print("\n" + "=" * 50)
+    print("MERGE SUMMARY")
+    print("=" * 50)
+    print(f"Total files found:      {stats.total_files_found}")
+    print(f"Valid trajectories:     {stats.valid_trajectories}")
+    print(f"Invalid trajectories:   {stats.invalid_trajectories}")
+    print(f"Duplicate trajectories: {stats.duplicate_trajectories}")
+    print(f"Merged trajectories:    {stats.merged_trajectories}")
+    
+    if stats.archetypes:
+        print("\nArchetype distribution:")
+        for archetype, count in sorted(stats.archetypes.items()):
+            pct = (count / stats.merged_trajectories * 100) if stats.merged_trajectories > 0 else 0
+            print(f"  {archetype}: {count} ({pct:.1f}%)")
+    
+    if not args.dry_run:
+        print(f"\nMerged trajectories saved to: {output_dir / 'trajectories'}")
+    
+    if stats.invalid_trajectories > stats.total_files_found * 0.1:
+        logger.warning("More than 10% of trajectories are invalid")
+        sys.exit(1)
+    
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()
+
+

package/python/scripts/run_ab_test.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""
+A/B Test Runner for Babylon Agent Training
+
+Compares a trained model against a baseline model using standardized evaluation scenarios.
+
+Usage:
+    # Compare trained model against baseline
+    python scripts/run_ab_test.py \
+        --model-a Qwen/Qwen3-30B \
+        --model-b ./trained_models/final_model \
+        --archetypes trader degen
+
+    # Quick test with fewer scenarios
+    python scripts/run_ab_test.py \
+        --model-a Qwen/Qwen2.5-0.5B-Instruct \
+        --model-b ./trained_models/final_model \
+        --num-runs 1
+
+    # Full test suite
+    python scripts/run_ab_test.py \
+        --model-a Qwen/Qwen3-30B \
+        --model-b ./trained_models/final_model \
+        --num-runs 5 \
+        --output-dir ./ab_results
+"""
+
+import argparse
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+# Add src to path for imports
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.training.ab_testing import ABTestRunner, EVAL_SCENARIOS
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+async def main():
+    parser = argparse.ArgumentParser(
+        description="Run A/B test comparing two models",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    
+    parser.add_argument(
+        "--model-a",
+        required=True,
+        help="Baseline model (path or HuggingFace name)",
+    )
+    parser.add_argument(
+        "--model-b", 
+        required=True,
+        help="Trained model to compare (path or HuggingFace name)",
+    )
+    parser.add_argument(
+        "--archetypes",
+        nargs="+",
+        choices=list(EVAL_SCENARIOS.keys()),
+        help="Archetypes to test (default: all)",
+    )
+    parser.add_argument(
+        "--num-runs",
+        type=int,
+        default=3,
+        help="Number of runs per scenario for statistical significance (default: 3)",
+    )
+    parser.add_argument(
+        "--vllm-url",
+        default="http://localhost:9001/v1",
+        help="vLLM server URL (default: http://localhost:9001/v1)",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="./ab_test_results",
+        help="Directory to save results (default: ./ab_test_results)",
+    )
+    parser.add_argument(
+        "--verbose", "-v",
+        action="store_true",
+        help="Enable verbose logging",
+    )
+    
+    args = parser.parse_args()
+    
+    if args.verbose:
+        logging.getLogger().setLevel(logging.DEBUG)
+    
+    # Filter scenarios by archetype if specified
+    scenarios = EVAL_SCENARIOS
+    if args.archetypes:
+        scenarios = {k: v for k, v in EVAL_SCENARIOS.items() if k in args.archetypes}
+        logger.info(f"Testing archetypes: {', '.join(args.archetypes)}")
+    else:
+        logger.info(f"Testing all archetypes: {', '.join(EVAL_SCENARIOS.keys())}")
+    
+    logger.info(f"Model A (baseline): {args.model_a}")
+    logger.info(f"Model B (trained):  {args.model_b}")
+    logger.info(f"Runs per scenario:  {args.num_runs}")
+    
+    # Create runner
+    runner = ABTestRunner(
+        model_a=args.model_a,
+        model_b=args.model_b,
+        scenarios=scenarios,
+        vllm_url=args.vllm_url,
+        num_runs_per_scenario=args.num_runs,
+        output_dir=args.output_dir,
+    )
+    
+    # Run tests
+    logger.info("Starting A/B test...")
+    result = await runner.run()
+    
+    # Print summary
+    print()
+    print(result.summary())
+    
+    # Return exit code based on result
+    if result.model_b_wins > result.model_a_wins:
+        logger.info("Trained model outperforms baseline!")
+        return 0
+    elif result.model_a_wins > result.model_b_wins:
+        logger.warning("Baseline model outperforms trained model")
+        return 1
+    else:
+        logger.info("Models performed equally")
+        return 0
+
+
+if __name__ == "__main__":
+    exit_code = asyncio.run(main())
+    sys.exit(exit_code)
+
+
+