run_codex_40_task_experiment.py

#!/usr/bin/env python3
"""Run the four-suite Codex-centered 40-task experiment matrix."""

from __future__ import annotations

import argparse
import math
import subprocess
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
DEFAULT_RUNS_ROOT = ROOT / "lmvs_runs"
DEFAULT_DEMO_BUNDLE_ROOT = ROOT / "lmvs_demo_bundles"
DEFAULT_DEMO_MEMORY_ROOT = ROOT / "lmvs_memory"

from libero_sdk.suites import canonical_suite_name


DEFAULT_SUITES = ["libero_spatial", "libero_object", "libero_goal", "libero_10"]


def split_csv(text: str) -> list[str]:
    return [item.strip() for item in text.split(",") if item.strip()]


def computed_gate_thresholds(
    *,
    suite_count: int,
    task_count: int,
    init_count: int,
    min_success_rate: float,
    max_few_shot_demos: int,
) -> dict[str, int]:
    expected_trials = suite_count * task_count * init_count
    expected_tasks = suite_count * task_count
    return {
        "expected_trials": expected_trials,
        "expected_tasks": expected_tasks,
        "min_memory_reuse_rows": max(1, math.ceil(expected_trials * 0.25)),
        "min_skill_reuse_rows": max(1, math.ceil(expected_trials * 0.25)),
        "min_same_task_prior_items": 1 if expected_trials > expected_tasks else 0,
        "min_cross_task_memory_items": 1 if task_count > 1 else 0,
        "min_cross_suite_memory_items": 1 if suite_count > 1 else 0,
        "min_tasks_with_success": math.ceil(expected_tasks * min_success_rate),
        "min_attempts_per_task": max(1, init_count),
        "max_few_shot_demos": max_few_shot_demos,
    }


def run_cmd(cmd: list[str], dry_run: bool) -> int:
    print("+ " + " ".join(cmd), flush=True)
    if dry_run:
        return 0
    return subprocess.run(cmd, cwd=str(ROOT), check=False).returncode


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--suites", default=",".join(DEFAULT_SUITES))
    ap.add_argument("--tasks", default="0,1,2,3,4,5,6,7,8,9")
    ap.add_argument("--init-states", default="0,1,2")
    ap.add_argument("--out-dir", default=str(DEFAULT_RUNS_ROOT))
    ap.add_argument("--demo-bundle-dir", default=str(DEFAULT_DEMO_BUNDLE_ROOT))
    ap.add_argument("--demo-memory-dir", default=str(DEFAULT_DEMO_MEMORY_ROOT))
    ap.add_argument("--tag", default="")
    ap.add_argument(
        "--extract-demo-memory",
        dest="extract_demo_memory",
        action="store_true",
        default=True,
        help=(
            "extract human-observable few-shot demo memory before each suite "
            "(default)"
        ),
    )
    ap.add_argument(
        "--no-extract-demo-memory",
        dest="extract_demo_memory",
        action="store_false",
        help=(
            "skip few-shot demo memory extraction for autonomous exploration "
            "or ablation runs"
        ),
    )
    ap.add_argument("--num-demos", type=int, default=3)
    ap.add_argument("--max-turns", type=int, default=10)
    ap.add_argument("--horizon", type=int, default=8000)
    ap.add_argument("--per-call-timeout", type=int, default=240)
    ap.add_argument("--camera-h", type=int, default=256)
    ap.add_argument("--camera-w", type=int, default=256)
    ap.add_argument("--codex-bin", default="codex")
    ap.add_argument(
        "--codex-transport",
        choices=["cli", "api", "subagent", "subagent_file", "sdk_python", "sdk_typescript"],
        default="subagent_file",
    )
    ap.add_argument("--codex-api-base-url", default=None)
    ap.add_argument("--codex-api-key-env", default="CODEX_API_KEY")
    ap.add_argument("--codex-api-model", default=None)
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--skip-gate", action="store_true")
    ap.add_argument(
        "--skip-demo-safety-checks",
        action="store_true",
        help=(
            "skip human-observable demo memory/bundle safety checks; "
            "do not use for final 40-task claims"
        ),
    )
    args = ap.parse_args()

    suites = [canonical_suite_name(s) for s in split_csv(args.suites)]
    task_items = split_csv(args.tasks)
    init_items = split_csv(args.init_states)
    gate_thresholds = computed_gate_thresholds(
        suite_count=len(suites),
        task_count=len(task_items),
        init_count=len(init_items),
        min_success_rate=0.95,
        max_few_shot_demos=args.num_demos if args.extract_demo_memory else 0,
    )
    tag = args.tag or time.strftime("%Y%m%d_%H%M%S")
    out_dir = Path(args.out_dir)
    out_paths: list[str] = []

    for suite in suites:
        if args.extract_demo_memory:
            cmd = [
                sys.executable,
                "scripts/extract_demo_knowledge.py",
                "--suite", suite,
                "--tasks", args.tasks.replace(",", "-") if args.tasks == "0,1,2,3,4,5,6,7,8,9" else args.tasks,
                "--num-demos", str(args.num_demos),
                "--codex-bin", args.codex_bin,
                "--api-backend", "codex-api" if args.codex_transport == "api" else "codex",
                "--demo-bundle-dir", args.demo_bundle_dir,
            ]
            if args.codex_transport == "api":
                if args.codex_api_base_url:
                    cmd.extend(["--codex-api-base-url", args.codex_api_base_url])
                cmd.extend(["--codex-api-key-env", args.codex_api_key_env])
                if args.codex_api_model:
                    cmd.extend(["--codex-api-model", args.codex_api_model])
            code = run_cmd(cmd, args.dry_run)
            if code != 0:
                return code

            if not args.skip_demo_safety_checks:
                for safety_cmd in [
                    [
                        sys.executable,
                        "scripts/check_human_observable_demo_memory.py",
                        args.demo_memory_dir,
                    ],
                    [
                        sys.executable,
                        "scripts/check_human_observable_demo_bundles.py",
                        args.demo_bundle_dir,
                    ],
                ]:
                    code = run_cmd(safety_cmd, args.dry_run)
                    if code != 0:
                        return code

        out_path = out_dir / f"codex40_{suite}_{tag}.json"
        out_paths.append(str(out_path))
        cmd = [
            sys.executable,
            "scripts/run_libero_object_sweep.py",
            "--suite", suite,
            "--tasks", args.tasks,
            "--init-states", args.init_states,
            "--max-turns", str(args.max_turns),
            "--horizon", str(args.horizon),
            "--policy-mode", "codex-brain",
            "--codex-tool-profile", "generic",
            "--stop-on-codex-health-failure",
            "--codex-bin", args.codex_bin,
            "--codex-transport", args.codex_transport,
            "--per-call-timeout", str(args.per_call_timeout),
            "--camera-h", str(args.camera_h),
            "--camera-w", str(args.camera_w),
            "--few-shot-demos", str(args.num_demos if args.extract_demo_memory else 0),
            "--learning-stage", (
                "fewshot_autonomous_eval"
                if args.extract_demo_memory else "autonomous_eval"
            ),
            "--learning-run-id", tag,
            "--out", str(out_path),
        ]
        if args.codex_transport == "api":
            if args.codex_api_base_url:
                cmd.extend(["--codex-api-base-url", args.codex_api_base_url])
            cmd.extend(["--codex-api-key-env", args.codex_api_key_env])
            if args.codex_api_model:
                cmd.extend(["--codex-api-model", args.codex_api_model])
        code = run_cmd(cmd, args.dry_run)
        if code != 0:
            return code

    if not args.skip_gate:
        cmd = [
            sys.executable,
            "scripts/check_codex_40_task_gate.py",
            *out_paths,
            "--expected-suites", ",".join(suites),
            "--expected-tasks", args.tasks,
            "--expected-inits", args.init_states,
            "--min-success-rate", "0.95",
        ]
        if not args.skip_demo_safety_checks:
            cmd.extend([
                "--demo-memory-dir", args.demo_memory_dir,
                "--demo-bundle-dir", args.demo_bundle_dir,
            ])
        code = run_cmd(cmd, args.dry_run)
        if code != 0:
            return code

        health_cmd = [
            sys.executable,
            "scripts/summarize_codex_health.py",
            *out_paths,
            "--fail-on-usage-limit",
            "--fail-on-missing-agent-io",
            "--fail-on-empty-response",
        ]
        code = run_cmd(health_cmd, args.dry_run)
        if code != 0:
            return code

        evidence_cmd = [
            sys.executable,
            "scripts/check_trial_evidence.py",
            *out_paths,
        ]
        code = run_cmd(evidence_cmd, args.dry_run)
        if code != 0:
            return code

        context_cmd = [
            sys.executable,
            "scripts/check_context_usage.py",
            *out_paths,
        ]
        code = run_cmd(context_cmd, args.dry_run)
        if code != 0:
            return code

        reuse_cmd = [
            sys.executable,
            "scripts/check_reuse_evidence.py",
            *out_paths,
            "--require-complete-evidence",
            "--min-memory-reuse-rows",
            str(gate_thresholds["min_memory_reuse_rows"]),
            "--min-skill-reuse-rows",
            str(gate_thresholds["min_skill_reuse_rows"]),
            "--min-same-task-prior-items",
            str(gate_thresholds["min_same_task_prior_items"]),
            "--min-cross-task-memory-items",
            str(gate_thresholds["min_cross_task_memory_items"]),
            "--min-cross-suite-memory-items",
            str(gate_thresholds["min_cross_suite_memory_items"]),
        ]
        code = run_cmd(reuse_cmd, args.dry_run)
        if code != 0:
            return code

        learning_cmd = [
            sys.executable,
            "scripts/check_learning_curve.py",
            *out_paths,
            "--require-complete-evidence",
            "--min-autonomous-trials",
            str(gate_thresholds["expected_trials"]),
            "--min-rows-with-few-shot-metadata",
            str(gate_thresholds["expected_trials"]),
            "--min-tasks-with-success",
            str(gate_thresholds["min_tasks_with_success"]),
            "--min-attempts-per-task",
            str(gate_thresholds["min_attempts_per_task"]),
            "--max-few-shot-demos",
            str(gate_thresholds["max_few_shot_demos"]),
        ]
        code = run_cmd(learning_cmd, args.dry_run)
        if code != 0:
            return code

        report_path = out_dir / f"codex40_report_{tag}.json"
        report_cmd = [
            sys.executable,
            "scripts/report_codex_matrix.py",
            *out_paths,
            "--require-complete-evidence",
            "--json",
            "--out", str(report_path),
        ]
        code = run_cmd(report_cmd, args.dry_run)
        if code != 0:
            return code

        eval_memory_cmd = [
            sys.executable,
            "scripts/update_evaluation_memory.py",
            *out_paths,
            "--memory-root", args.demo_memory_dir,
            "--label", tag,
            "--report-json", str(report_path),
            "--require-complete-evidence",
        ]
        code = run_cmd(eval_memory_cmd, args.dry_run)
        if code != 0:
            return code

        skill_perf_cmd = [
            sys.executable,
            "scripts/update_skill_performance.py",
            *out_paths,
            "--memory-root", args.demo_memory_dir,
            "--require-complete-evidence",
        ]
        code = run_cmd(skill_perf_cmd, args.dry_run)
        if code != 0:
            return code

        perception_perf_cmd = [
            sys.executable,
            "scripts/update_perception_performance.py",
            *out_paths,
            "--memory-root", args.demo_memory_dir,
            "--require-complete-evidence",
        ]
        code = run_cmd(perception_perf_cmd, args.dry_run)
        if code != 0:
            return code

    return 0


if __name__ == "__main__":
    raise SystemExit(main())