| 1 | #!/usr/bin/env python3
|
| 2 | """Run baseline, generic-review, or OpenClasp Shield τ³ evaluations."""
|
| 3 |
|
| 4 | from __future__ import annotations
|
| 5 |
|
| 6 | import argparse
|
| 7 | import json
|
| 8 | import os
|
| 9 | from pathlib import Path
|
| 10 |
|
| 11 | from tau2.data_model.simulation import TextRunConfig
|
| 12 | from tau2.runner import run_domain
|
| 13 |
|
| 14 | from openclasp_agent import close_shield_cases, register_openclasp_agents
|
| 15 | from openclasp_client import OpenClaspMcpClient
|
| 16 |
|
| 17 |
|
| 18 | def parse_args() -> argparse.Namespace:
|
| 19 | parser = argparse.ArgumentParser()
|
| 20 | parser.add_argument(
|
| 21 | "--mode", choices=["baseline", "generic-review", "shield"], required=True
|
| 22 | )
|
| 23 | parser.add_argument("--domain", default="airline")
|
| 24 | parser.add_argument("--task-ids", nargs="+")
|
| 25 | parser.add_argument("--num-tasks", type=int)
|
| 26 | parser.add_argument("--num-trials", type=int, default=1)
|
| 27 | parser.add_argument("--agent-llm", default="openai/gpt-4.1")
|
| 28 | parser.add_argument("--user-llm", default="openai/gpt-4.1")
|
| 29 | parser.add_argument("--generic-review-model")
|
| 30 | parser.add_argument("--max-concurrency", type=int, default=1)
|
| 31 | parser.add_argument("--seed", type=int, default=300)
|
| 32 | parser.add_argument("--temperature", type=float, default=0.0)
|
| 33 | parser.add_argument("--retrieval-config")
|
| 34 | parser.add_argument("--trigger", choices=["decisions", "actions", "all"], default="decisions")
|
| 35 | parser.add_argument("--save-to")
|
| 36 | parser.add_argument("--summary-file")
|
| 37 | return parser.parse_args()
|
| 38 |
|
| 39 |
|
| 40 | def main() -> None:
|
| 41 | args = parse_args()
|
| 42 | save_to = args.save_to or f"openclasp_{args.mode.replace('-', '_')}_{args.domain}"
|
| 43 | summary_file = Path(args.summary_file or f"{save_to}_summary.json")
|
| 44 | os.environ["OPENCLASP_TAU_DOMAIN"] = args.domain
|
| 45 | os.environ["OPENCLASP_TAU_TRIGGER"] = args.trigger
|
| 46 | if args.generic_review_model:
|
| 47 | os.environ["OPENCLASP_GENERIC_REVIEW_MODEL"] = args.generic_review_model
|
| 48 |
|
| 49 | register_openclasp_agents()
|
| 50 | agent = {
|
| 51 | "baseline": "llm_agent",
|
| 52 | "generic-review": "openclasp_generic_review_agent",
|
| 53 | "shield": "openclasp_shield_agent",
|
| 54 | }[args.mode]
|
| 55 | if args.mode == "shield":
|
| 56 | token = os.getenv("OPENCLASP_AGENT_TOKEN", "")
|
| 57 | client = OpenClaspMcpClient(
|
| 58 | os.getenv("OPENCLASP_URL", "https://openclasp.dev"), token
|
| 59 | )
|
| 60 | client.call_tool("openclasp_shield_list_cases", {})
|
| 61 | else:
|
| 62 | client = None
|
| 63 |
|
| 64 | config = TextRunConfig(
|
| 65 | domain=args.domain,
|
| 66 | agent=agent,
|
| 67 | llm_agent=args.agent_llm,
|
| 68 | llm_user=args.user_llm,
|
| 69 | llm_args_agent={"temperature": args.temperature},
|
| 70 | llm_args_user={"temperature": args.temperature},
|
| 71 | task_ids=args.task_ids,
|
| 72 | num_tasks=args.num_tasks,
|
| 73 | num_trials=args.num_trials,
|
| 74 | max_concurrency=args.max_concurrency,
|
| 75 | workers=0,
|
| 76 | seed=args.seed,
|
| 77 | retrieval_config=args.retrieval_config,
|
| 78 | save_to=save_to,
|
| 79 | )
|
| 80 | results = run_domain(config)
|
| 81 | shield_cases = close_shield_cases(results, client) if client else []
|
| 82 | rewards = [
|
| 83 | simulation.reward_info.reward
|
| 84 | for simulation in results.simulations
|
| 85 | if simulation.reward_info and simulation.reward_info.reward is not None
|
| 86 | ]
|
| 87 | summary = {
|
| 88 | "mode": args.mode,
|
| 89 | "domain": args.domain,
|
| 90 | "agent": agent,
|
| 91 | "agentModel": args.agent_llm,
|
| 92 | "userModel": args.user_llm,
|
| 93 | "seed": args.seed,
|
| 94 | "numTrials": args.num_trials,
|
| 95 | "simulationCount": len(results.simulations),
|
| 96 | "averageReward": sum(rewards) / len(rewards) if rewards else None,
|
| 97 | "passCount": sum(reward >= 0.999 for reward in rewards),
|
| 98 | "rewards": rewards,
|
| 99 | "simulations": [
|
| 100 | {
|
| 101 | "simulationId": simulation.id,
|
| 102 | "taskId": simulation.task_id,
|
| 103 | "trial": simulation.trial,
|
| 104 | "seed": simulation.seed,
|
| 105 | "reward": simulation.reward_info.reward if simulation.reward_info else None,
|
| 106 | "termination": str(simulation.termination_reason.value),
|
| 107 | "agentCost": simulation.agent_cost,
|
| 108 | "duration": simulation.duration,
|
| 109 | }
|
| 110 | for simulation in results.simulations
|
| 111 | ],
|
| 112 | "shieldCases": shield_cases,
|
| 113 | }
|
| 114 | summary_file.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
|
| 115 | print(f"OpenClasp experiment summary: {summary_file.resolve()}")
|
| 116 |
|
| 117 |
|
| 118 | if __name__ == "__main__":
|
| 119 | main()
|