run_full_experiment.py

#!/usr/bin/env python3

import argparse
import json
import sys
from datetime import datetime
from pathlib import Path

# Load environment variables from .env file
from dotenv import load_dotenv
load_dotenv()

from src.environment.fileworld import get_task, get_all_tasks, get_tasks_by_complexity, TaskComplexity
from src.models.openai_client import OpenAIClient, MockOpenAIClient, DeterminismConfig
from src.agents.react_agent import ReActAgent
from src.experiments.runner import ExperimentRunner, ExperimentConfig, compute_experiment_metrics
from src.experiments.analysis import analyze_results, generate_summary_report


def run_full_experiment(use_api: bool = True, num_tasks: int = 30):
    print("=" * 70)
    print("TRAJECTORY STABILITY EXPERIMENT - FULL RUN")
    print("=" * 70)
    print(f"Start time: {datetime.now().isoformat()}")
    print()

    # Select tasks
    all_tasks = get_all_tasks()
    task_ids = [t.task_id for t in all_tasks[:num_tasks]]

    # Configure experiment
    config = ExperimentConfig(
        agent_name="gpt-4o-mini" if use_api else "mock-agent",
        perturbation_types=["MEM-REORDER", "OBS-PARAPHRASE", "CONTEXT-INJECT"],
        seeds=[42, 123, 456, 789, 1000],
        task_ids=task_ids,
        max_steps=20,
        output_dir="results/full_experiment"
    )

    # Calculate totals
    num_baseline = len(task_ids)
    num_perturbed = len(task_ids) * len(config.perturbation_types) * len(config.seeds)
    total_runs = num_baseline + num_perturbed

    print(f"Configuration:")
    print(f"  Model: {config.agent_name}")
    print(f"  Tasks: {len(task_ids)}")
    print(f"  Perturbation types: {config.perturbation_types}")
    print(f"  Seeds per perturbation: {config.seeds}")
    print(f"  Max steps per episode: {config.max_steps}")
    print()
    print(f"Total runs: {num_baseline} baseline + {num_perturbed} perturbed = {total_runs}")
    print(f"Estimated API calls: ~{total_runs * 10} (assuming ~10 steps/episode avg)")
    print()

    # Create client
    if use_api:
        print("Using OpenAI API (gpt-4o-mini)")
        print("  temperature=0, top_p=1, seed=42 for determinism")
        client = OpenAIClient(DeterminismConfig())
    else:
        print("Using Mock Client (no API calls)")
        client = MockOpenAIClient(responses={
            "": "Thought: I should navigate toward my goal.\nAction: go lobby"
        })

    print()
    print("-" * 70)
    print("RUNNING EXPERIMENT")
    print("-" * 70)
    print()

    # Run experiment
    runner = ExperimentRunner(config, client=client)
    results = runner.run(verbose=True)

    print()
    print("-" * 70)
    print("COMPUTING METRICS")
    print("-" * 70)
    print()

    # Compute metrics
    metrics = compute_experiment_metrics(results)

    # Save metrics
    metrics_path = Path(config.output_dir) / f"metrics_{runner.experiment_id}.json"
    with open(metrics_path, "w") as f:
        json.dump(metrics, f, indent=2, default=str)
    print(f"Metrics saved to: {metrics_path}")

    print()
    print("-" * 70)
    print("ANALYSIS")
    print("-" * 70)
    print()

    # Analyze results
    analysis = analyze_results(results)
    report = generate_summary_report(analysis)
    print(report)

    # Save analysis
    analysis_path = Path(config.output_dir) / f"analysis_{runner.experiment_id}.json"
    with open(analysis_path, "w") as f:
        json.dump(analysis, f, indent=2, default=str)
    print(f"\nAnalysis saved to: {analysis_path}")

    # Print key findings
    print()
    print("=" * 70)
    print("KEY FINDINGS")
    print("=" * 70)
    print()

    if metrics["aggregate"]["mean_tdr"] is not None:
        print(f"Mean TDR: {metrics['aggregate']['mean_tdr']:.3f} (+/- {metrics['aggregate']['std_tdr']:.3f})")
    if metrics["aggregate"]["mean_dos"] is not None:
        print(f"Mean DOS: {metrics['aggregate']['mean_dos']:.1f} (+/- {metrics['aggregate']['std_dos']:.1f})")

    # Count outcome flips
    outcome_flips = {"baseline_fail_pert_success": 0, "baseline_success_pert_fail": 0}
    for task_id in results.task_ids:
        baseline_success = results.baseline_trajectories[task_id].success
        for traj in results.perturbed_trajectories[task_id]:
            if baseline_success and not traj.success:
                outcome_flips["baseline_success_pert_fail"] += 1
            elif not baseline_success and traj.success:
                outcome_flips["baseline_fail_pert_success"] += 1

    print(f"\nOutcome Flips:")
    print(f"  Baseline FAIL -> Perturbed SUCCESS: {outcome_flips['baseline_fail_pert_success']}")
    print(f"  Baseline SUCCESS -> Perturbed FAIL: {outcome_flips['baseline_success_pert_fail']}")

    # API usage
    if use_api and hasattr(client, 'call_records'):
        total_tokens = sum(r.usage.get("total_tokens", 0) for r in client.call_records)
        print(f"\nAPI Usage:")
        print(f"  Total API calls: {len(client.call_records)}")
        print(f"  Total tokens: {total_tokens:,}")

        # Check reproducibility
        repro = client.verify_reproducibility()
        if repro["is_reproducible"]:
            print(f"  Reproducibility: VERIFIED (single fingerprint)")
        else:
            print(f"  Reproducibility: WARNING - {repro['fingerprint_count']} different fingerprints")

    print()
    print(f"End time: {datetime.now().isoformat()}")
    print("=" * 70)

    return results, metrics, analysis


def main():
    parser = argparse.ArgumentParser(
        description="Run full trajectory stability experiment"
    )
    parser.add_argument(
        "--use-api",
        action="store_true",
        help="Use real OpenAI API (requires OPENAI_API_KEY)"
    )
    parser.add_argument(
        "--mock",
        action="store_true",
        help="Use mock client (no API calls, for testing)"
    )
    parser.add_argument(
        "--num-tasks",
        type=int,
        default=30,
        help="Number of tasks to run (default: 30)"
    )

    args = parser.parse_args()

    # Determine whether to use API
    use_api = args.use_api and not args.mock

    if not use_api and not args.mock:
        print("Note: Running with mock client. Use --use-api for real experiments.")
        print()

    run_full_experiment(use_api=use_api, num_tasks=args.num_tasks)


if __name__ == "__main__":
    main()