#!/usr/bin/env python3 """ OTAP: one-shot experimental pipeline (Part A perturbation meta-evaluation). One run produces: results/results.json all raw scores + summary numbers results/results_tables.tex LaTeX rows ready to paste into the paper results/ablations.tex ablation table rows (with --ablations) What this script CAN populate honestly: Table I (perturbation meta-eval, PRA) on a released, programmatically generated corpus of agent task plans. What it CANNOT populate: Table II (human correlation) -- that requires human annotators; a harness is provided (part_b_human_correlation) that computes Spearman/Kendall once you supply annotations.csv. Usage: pip install numpy scipy scikit-learn # required pip install sentence-transformers # recommended encoder python run_all.py # main table python run_all.py --ablations # + ablation table ANTHROPIC_API_KEY=... python run_all.py --llm-judge # + LLM judge baseline Deterministic given --seed. """ import os os.environ["VECLIB_MAXIMUM_THREADS"] = "1" os.environ["NUMEXPR_NUM_THREADS"] = "1" # Optional extra safety import argparse, json, math, os, random, re, sys, time # ... rest of your imports import argparse, json, math, os, random, re, sys, time from dataclasses import dataclass, field from itertools import product from typing import List, Dict, Optional, Tuple import numpy as np from scipy.optimize import linear_sum_assignment from scipy.special import logsumexp # ---------------------------------------------------------------------------- # 1. Trajectory data model # ---------------------------------------------------------------------------- @dataclass class Step: key: str # stable id within the plan action_variants: List[str] # paraphrase surface forms; [0] is canonical context: str effect: str tool: str # tool name or "none" consumes: List[str] = field(default_factory=list) produces: List[str] = field(default_factory=list) variant_idx: int = 0 @property def action(self) -> str: return self.action_variants[self.variant_idx % len(self.action_variants)] @property def text(self) -> str: return f"{self.action}. {self.context}" def clone(self, **kw) -> "Step": s = Step(self.key, list(self.action_variants), self.context, self.effect, self.tool, list(self.consumes), list(self.produces), self.variant_idx) for k, v in kw.items(): setattr(s, k, v) return s @dataclass class Trajectory: task: str steps: List[Step] tag: str = "original" # perturbation family tag # -- dependency graph via artifact tracing (edges only point backwards # in execution order: a consumer links to the *nearest earlier* producer; # if none exists the precondition is unmet and no edge is created) ----- def edges(self) -> List[Tuple[int, int]]: E = [] for j, s in enumerate(self.steps): for art in s.consumes: for i in range(j - 1, -1, -1): if art in self.steps[i].produces: E.append((i, j)) break return sorted(set(E)) def signed_dep_matrix(self) -> np.ndarray: """Signed, normalized dependency-distance matrix D (Eq. 2 of paper).""" n = len(self.steps) E = self.edges() adj = [[] for _ in range(n)] for i, j in E: adj[i].append(j) INF = 10 ** 9 hop = np.full((n, n), INF, dtype=float) for src in range(n): # BFS shortest directed hops hop[src, src] = 0 q = [src] while q: u = q.pop(0) for v in adj[u]: if hop[src, v] > hop[src, u] + 1: hop[src, v] = hop[src, u] + 1 q.append(v) H = max(1.0, float(hop[hop < INF].max())) # DAG height (max path) D = np.zeros((n, n)) for i in range(n): for k in range(n): if i == k: continue if hop[i, k] < INF: D[i, k] = +hop[i, k] / H elif hop[k, i] < INF: D[i, k] = -hop[k, i] / H # else incomparable -> 0 return D # ---------------------------------------------------------------------------- # 2. Task corpus: 8 domain templates x parametric expansion # Every step carries paraphrase variants, tools, artifacts. # ---------------------------------------------------------------------------- def S(key, av, ctx, eff, tool, cons=(), prod=()): return Step(key, av, ctx, eff, tool, list(cons), list(prod)) def hotel_plan(city, budget): t = f"Book a hotel in {city} under ${budget} near the main station." return t, [ S("search", [f"Search for hotels in {city}", f"Look up available hotels in {city}", f"Query lodging options located in {city}"], f"hotels in {city}, budget ${budget}, near main station", "list of candidate hotels", "web_search", [], ["candidates"]), S("filt_price", [f"Filter the hotels by price under ${budget}", f"Remove hotels costing more than ${budget}", f"Keep only hotels priced below ${budget}"], f"price threshold ${budget} per night", "price-filtered hotel list", "none", ["candidates"], ["by_price"]), S("filt_loc", ["Filter the hotels by distance to the main station", "Keep hotels close to the main station", "Discard hotels far from the station area"], f"walking distance to {city} main station", "location-filtered hotel list", "none", ["candidates"], ["by_loc"]), S("rank", ["Rank the remaining hotels by review score", "Sort the shortlisted hotels by rating", "Order the surviving hotels from best to worst reviews"], "combine price and location shortlists, sort by rating", "ranked shortlist", "none", ["by_price", "by_loc"], ["ranked"]), S("book", ["Book the top ranked hotel", "Complete the reservation for the best hotel", "Reserve the highest ranked hotel option"], f"reservation for the chosen hotel in {city}", "confirmed booking", "booking_api", ["ranked"], ["booking"]), ] def flight_plan(orig, dest): t = f"Find and book the cheapest reasonable flight from {orig} to {dest}." return t, [ S("search", [f"Search flights from {orig} to {dest}", f"Look up airfare between {orig} and {dest}", f"Query flight availability {orig} to {dest}"], f"one-way economy, {orig} to {dest}, next month", "flight result list", "web_search", [], ["flights"]), S("compare", ["Compare the airlines on price and duration", "Contrast the candidate flights by cost and travel time", "Evaluate the flight options against each other"], "price, total duration, number of stops", "comparison table of flights", "none", ["flights"], ["compared"]), S("select", ["Select the cheapest acceptable flight", "Pick the lowest cost flight that meets constraints", "Choose the most affordable suitable flight"], "cheapest with at most one stop", "selected flight", "none", ["compared"], ["choice"]), S("book", ["Book the selected flight", "Purchase the ticket for the chosen flight", "Complete checkout for the selected itinerary"], "passenger details and payment", "ticket confirmation", "booking_api", ["choice"], ["ticket"]), ] def research_plan(topic): t = f"Write a short evidence-backed answer about {topic}." return t, [ S("search", [f"Search the web for recent sources on {topic}", f"Look up up-to-date material about {topic}", f"Gather current references regarding {topic}"], f"recent, reputable sources on {topic}", "list of relevant documents", "web_search", [], ["docs"]), S("read", ["Read the top documents and extract key claims", "Go through the retrieved pages and pull out the main facts", "Skim the best sources and note their central claims"], "focus on primary sources and data", "extracted notes", "browser", ["docs"], ["notes"]), S("cross", ["Cross-check the claims across sources", "Verify the extracted facts against independent sources", "Validate each claim using a second source"], "flag contradictions between sources", "verified claim set", "browser", ["notes"], ["verified"]), S("write", ["Write the final answer with citations", "Compose the response citing the verified sources", "Draft the concluding answer with references"], "concise answer, inline citations", "final written answer", "none", ["verified"], ["answer"]), ] def data_plan(dataset): t = f"Analyze the {dataset} dataset and chart revenue by region." return t, [ S("download", [f"Download the {dataset} dataset", f"Fetch the {dataset} data file", f"Retrieve the raw {dataset} dataset"], f"{dataset} CSV from the data store", "raw data file on disk", "python", [], ["raw"]), S("clean", ["Clean the data and fix malformed rows", "Preprocess the dataset and repair bad records", "Sanitize the raw table, dropping corrupt rows"], "handle missing values, normalize headers", "clean dataframe", "python", ["raw"], ["clean"]), S("aggregate", ["Aggregate revenue by region", "Group the records by region and sum revenue", "Compute per-region revenue totals"], "group-by region, sum of revenue column", "aggregated table", "python", ["clean"], ["agg"]), S("plot", ["Plot the revenue by region as a bar chart", "Render a bar chart of regional revenue", "Draw the per-region revenue figure"], "bar chart, labeled axes, sorted bars", "saved chart image", "python", ["agg"], ["chart"]), S("report", ["Summarize the findings alongside the chart", "Write a brief summary to accompany the figure", "Report the main takeaways with the chart attached"], "two-paragraph summary referencing the chart", "final report", "none", ["agg", "chart"], ["report"]), ] def bugfix_plan(component): t = f"Diagnose and fix the failing test in the {component} component." return t, [ S("repro", [f"Run the failing test in {component} to reproduce the error", f"Execute the broken {component} test and capture the failure", f"Reproduce the {component} test failure locally"], "capture stack trace and failing assertion", "reproduced failure with trace", "shell", [], ["trace"]), S("locate", ["Locate the faulty code from the stack trace", "Trace the error back to the offending function", "Identify the source lines responsible for the failure"], "walk the trace to the root frame", "identified faulty function", "editor", ["trace"], ["locus"]), S("fix", ["Edit the code to fix the bug", "Apply a patch correcting the faulty logic", "Modify the offending function to repair the defect"], "minimal change preserving the interface", "patched source file", "editor", ["locus"], ["patch"]), S("retest", ["Re-run the test suite to confirm the fix", "Execute the tests again and verify they pass", "Run the full suite to validate the patch"], "all tests green including the previously failing one", "passing test run", "shell", ["patch"], ["green"]), ] def shopping_plan(item, price): t = f"Buy a {item} online for under ${price} with good reviews." return t, [ S("search", [f"Search online stores for a {item}", f"Look up {item} listings across shops", f"Query e-commerce sites for a {item}"], f"{item}, budget ${price}", "product listing results", "web_search", [], ["listings"]), S("filt_price", [f"Filter listings to those under ${price}", f"Drop products priced above ${price}", f"Keep only offers below ${price}"], f"maximum price ${price}", "price-filtered listings", "none", ["listings"], ["cheap"]), S("filt_rev", ["Filter listings by review rating above four stars", "Keep products rated over four stars", "Remove poorly reviewed items"], "minimum 4.0 average rating", "review-filtered listings", "none", ["listings"], ["liked"]), S("choose", ["Choose the best remaining product", "Pick the top option among the filtered listings", "Select the strongest candidate product"], "intersect price and review shortlists", "chosen product", "none", ["cheap", "liked"], ["choice"]), S("checkout", ["Add the product to the cart and check out", "Purchase the selected item through checkout", "Complete the order for the chosen product"], "shipping address and payment on file", "order confirmation", "browser", ["choice"], ["order"]), ] def email_plan(person): t = f"Find the latest report from {person} and email them three questions." return t, [ S("find", [f"Search the drive for the latest report by {person}", f"Locate {person}'s most recent report in the drive", f"Look up the newest report authored by {person}"], f"most recent document by {person}", "located report file", "gdrive", [], ["report"]), S("read", ["Read the report and note unclear points", "Go through the document marking ambiguities", "Review the report and list confusing sections"], "focus on methodology and results sections", "list of unclear points", "gdrive", ["report"], ["points"]), S("draft", ["Draft three specific questions about the report", "Write three targeted questions on the unclear points", "Compose three concrete questions for the author"], "one question per unclear point, specific and short", "drafted questions", "none", ["points"], ["questions"]), S("send", [f"Email the questions to {person}", f"Send the drafted questions to {person} by email", f"Deliver the questions to {person} via email"], "polite subject line, questions in body", "sent email", "email", ["questions"], ["sent"]), ] def trip_plan(city): t = f"Plan a one-day itinerary in {city} and check the weather." return t, [ S("sights", [f"Search for top sights in {city}", f"Look up the main attractions of {city}", f"Find the must-see places in {city}"], f"top rated attractions in {city}", "list of candidate sights", "web_search", [], ["sights"]), S("weather", [f"Check the weather forecast for {city}", f"Look up tomorrow's weather in {city}", f"Fetch the {city} forecast"], "hourly forecast for the visit day", "weather forecast", "weather_api", [], ["forecast"]), S("select", ["Select sights compatible with the forecast", "Pick attractions that fit the expected weather", "Choose indoor or outdoor sights based on the forecast"], "prefer indoor sights if rain expected", "selected sights", "none", ["sights", "forecast"], ["selection"]), S("route", ["Order the selected sights into a walking route", "Sequence the chosen sights into an efficient path", "Arrange the picks into a sensible day route"], "minimize backtracking, lunch stop midway", "ordered itinerary", "maps", ["selection"], ["route"]), S("output", ["Write out the final one-day itinerary", "Produce the final schedule for the day", "Present the completed day plan"], "timed plan with travel segments", "final itinerary", "none", ["route"], ["plan"]), ] def finance_plan(company): t = f"Compute {company}'s revenue growth from its last two annual reports." return t, [ S("find", [f"Search for {company}'s two most recent annual reports", f"Locate the latest two annual filings of {company}", f"Look up {company}'s recent yearly reports"], "official filings, last two fiscal years", "links to two reports", "web_search", [], ["reports"]), S("extract", ["Extract total revenue from each report", "Pull the annual revenue figure out of both filings", "Read off total revenue in each document"], "income statement, total revenue line", "two revenue figures", "browser", ["reports"], ["figures"]), S("compute", ["Compute the year-over-year growth rate", "Calculate the percentage revenue growth", "Derive the YoY growth from the two figures"], "(new - old) / old, as a percentage", "growth percentage", "calculator", ["figures"], ["growth"]), S("answer", ["Report the growth rate with both source figures", "State the computed growth citing the two filings", "Answer with the growth number and its sources"], "one-line answer plus the two inputs", "final answer", "none", ["growth"], ["answer"]), ] DOMAIN_TEMPLATES = [ (hotel_plan, [("Kyoto", 150), ("Lisbon", 120), ("Osaka", 130), ("Prague", 110), ("Seoul", 140)]), (flight_plan, [("Boston", "Denver"), ("Paris", "Rome"), ("Tokyo", "Sapporo"), ("Miami", "Austin"), ("Berlin", "Madrid")]), (research_plan, [("heat pump efficiency",), ("urban rewilding",), ("battery recycling",), ("coral bleaching",), ("rail electrification",)]), (data_plan, [("quarterly sales",), ("clinic visits",), ("store inventory",), ("web traffic",), ("shipment logs",)]), (bugfix_plan, [("parser",), ("scheduler",), ("cache layer",), ("auth module",), ("exporter",)]), (shopping_plan, [("mechanical keyboard", 90), ("office chair", 250), ("air purifier", 180), ("standing desk", 400), ("monitor arm", 120)]), (email_plan, [("Dr. Alvarez",), ("the finance lead",), ("Priya",), ("the lab manager",), ("Tomasz",)]), (trip_plan, [("Vienna",), ("Kanazawa",), ("Porto",), ("Tallinn",), ("Valencia",)]), ] IRRELEVANT_POOL = [ S("noise_a", ["Check the social media feed", "Scroll through the social feed", "Open the social media timeline"], "unrelated browsing", "nothing useful", "browser", [], []), S("noise_b", ["Play some background music", "Start a music playlist", "Put on a song"], "unrelated audio", "nothing useful", "none", [], []), S("noise_c", ["Look up the office lunch menu", "Check today's cafeteria menu", "Read the lunch options"], "unrelated food query", "nothing useful", "web_search", [], []), S("noise_d", ["Re-read the original task statement again", "Restate the task once more", "Repeat the task description"], "redundant restatement", "nothing new", "none", [], []), ] TOOL_GROUPS = { "web_search": "search", "browser": "web", "gdrive": "files", "email": "comm", "booking_api": "transact", "python": "compute", "calculator": "compute", "shell": "compute", "editor": "files", "maps": "search", "weather_api": "search", "sql": "compute", "none": "none", } INCOMPATIBLE_SWAP = { # deliberately wrong tool for P5 corruption "web_search": "sql", "browser": "calculator", "python": "email", "booking_api": "editor", "gdrive": "weather_api", "email": "shell", "shell": "maps", "editor": "booking_api", "calculator": "browser", "maps": "sql", "weather_api": "gdrive", "none": "sql", } def tool_dist(a: str, b: str) -> float: if a == b: return 0.0 ga, gb = TOOL_GROUPS.get(a, a), TOOL_GROUPS.get(b, b) if ga == gb: return 0.3 # substitutable (calculator vs python) if "none" in (a, b): return 0.6 # reasoning step vs tool step return 1.0 # incompatible def build_corpus(n_expansions: int) -> List[Trajectory]: corpus = [] for fn, params in DOMAIN_TEMPLATES: for args in params[:n_expansions]: task, steps = fn(*args) corpus.append(Trajectory(task, steps, tag="original")) return corpus # ---------------------------------------------------------------------------- # 3. Perturbation generators (P1..P5) with known ground-truth quality tiers # tier 0: original, P1 paraphrase, P2 valid reorder, P3 merge (equivalent) # tier 1: P4 causal inversion (worse) # tier 2: P5 corruption (delete critical / hallucinate / bad tool) (worst) # ---------------------------------------------------------------------------- TIER = {"original": 0, "P1": 0, "P2": 0, "P3": 0, "P4": 1, "P5": 2} def p1_paraphrase(tr: Trajectory, rng: random.Random) -> Trajectory: steps = [s.clone(variant_idx=rng.choice([1, 2])) for s in tr.steps] return Trajectory(tr.task, steps, "P1") def incomparable_pairs(tr: Trajectory) -> List[Tuple[int, int]]: D = tr.signed_dep_matrix() n = len(tr.steps) return [(i, j) for i in range(n) for j in range(i + 1, n) if D[i, j] == 0] def p2_valid_reorder(tr: Trajectory, rng: random.Random) -> Optional[Trajectory]: pairs = incomparable_pairs(tr) # only swaps that keep every consumer after its producer ok = [] for i, j in pairs: steps = list(tr.steps) steps[i], steps[j] = steps[j], steps[i] cand = Trajectory(tr.task, steps, "P2") if len(cand.edges()) == len(tr.edges()): ok.append(cand) return rng.choice(ok) if ok else None def p3_merge(tr: Trajectory, rng: random.Random) -> Optional[Trajectory]: idxs = [i for i in range(len(tr.steps) - 1)] rng.shuffle(idxs) for i in idxs: a, b = tr.steps[i], tr.steps[i + 1] internal = set(a.produces) & set(b.consumes) merged = Step( a.key + "+" + b.key, [f"{a.action_variants[k % len(a.action_variants)]} and " f"{b.action_variants[k % len(b.action_variants)][0].lower()}" f"{b.action_variants[k % len(b.action_variants)][1:]}" for k in range(3)], f"{a.context}; {b.context}", b.effect, a.tool if a.tool != "none" else b.tool, sorted((set(a.consumes) | set(b.consumes)) - internal), sorted(set(a.produces) | set(b.produces)), ) steps = tr.steps[:i] + [merged] + tr.steps[i + 2:] cand = Trajectory(tr.task, steps, "P3") if all(len(s.consumes) == 0 or True for s in steps): return cand return None def p4_causal_inversion(tr: Trajectory, rng: random.Random) -> Optional[Trajectory]: E = tr.edges() rng.shuffle(E) for (i, j) in E: steps = list(tr.steps) steps[i], steps[j] = steps[j], steps[i] cand = Trajectory(tr.task, steps, "P4") if len(cand.edges()) < len(tr.edges()): # a precondition is now unmet return cand return None def p5_corruptions(tr: Trajectory, rng: random.Random) -> List[Trajectory]: out = [] # (a) delete a critical step (its product is consumed later) consumed_later = set(a for s in tr.steps for a in s.consumes) crit = [i for i, s in enumerate(tr.steps) if set(s.produces) & consumed_later] if crit: i = rng.choice(crit) out.append(Trajectory(tr.task, tr.steps[:i] + tr.steps[i + 1:], "P5")) # (b) hallucinate: insert 2 irrelevant steps at random positions steps = list(tr.steps) for noise in rng.sample(IRRELEVANT_POOL, 2): steps.insert(rng.randrange(len(steps) + 1), noise.clone()) out.append(Trajectory(tr.task, steps, "P5")) # (c) incompatible tool swap on 2 tool-using steps steps = [s.clone() for s in tr.steps] tool_idx = [i for i, s in enumerate(steps) if s.tool != "none"] for i in rng.sample(tool_idx, min(2, len(tool_idx))): steps[i].tool = INCOMPATIBLE_SWAP.get(steps[i].tool, "sql") out.append(Trajectory(tr.task, steps, "P5")) return out def make_variants(tr: Trajectory, rng: random.Random) -> List[Trajectory]: vs = [tr] vs += [p1_paraphrase(tr, rng) for _ in range(2)] for _ in range(2): v = p2_valid_reorder(tr, rng) if v: vs.append(v) v = p3_merge(tr, rng) if v: vs.append(v) for _ in range(2): v = p4_causal_inversion(tr, rng) if v: vs.append(v) vs += p5_corruptions(tr, rng) return vs # ---------------------------------------------------------------------------- # 4. Encoders # ---------------------------------------------------------------------------- class Encoder: """sentence-transformers if available, else char+word TF-IDF fallback.""" def __init__(self, mode: str, all_texts: List[str]): self.mode = mode if mode == "sbert": from sentence_transformers import SentenceTransformer # FORCE: device="cpu" bypasses buggy/unstable macOS MPS float16 math optimizations self.model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu") self.cache: Dict[str, np.ndarray] = {} else: from sklearn.feature_extraction.text import TfidfVectorizer self.vec_w = TfidfVectorizer(ngram_range=(1, 2), sublinear_tf=True) self.vec_c = TfidfVectorizer(analyzer="char_wb", ngram_range=(3, 5), sublinear_tf=True) self.vec_w.fit(all_texts) self.vec_c.fit(all_texts) self.cache = {} def encode(self, texts: List[str]) -> np.ndarray: missing = [t for t in texts if t not in self.cache] if missing: if self.mode == "sbert": em = self.model.encode(missing, normalize_embeddings=True, show_progress_bar=False) for t, e in zip(missing, em): # DETAILED CLEAN: Force inf/nan values down to exactly 0.0 to prevent matrix multiplication expansions v = np.nan_to_num(np.asarray(e, dtype=float), nan=0.0, posinf=0.0, neginf=0.0) norm = np.linalg.norm(v) if norm > 1e-9: v = v / norm self.cache[t] = v else: w = self.vec_w.transform(missing).toarray() c = self.vec_c.transform(missing).toarray() em = np.hstack([w, c]) em /= (np.linalg.norm(em, axis=1, keepdims=True) + 1e-12) for t, e in zip(missing, em): self.cache[t] = np.nan_to_num(e, nan=0.0, posinf=0.0, neginf=0.0) return np.stack([self.cache[t] for t in texts]) def cos_dist(A: np.ndarray, B: np.ndarray) -> np.ndarray: A_clean = np.nan_to_num(A, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) B_clean = np.nan_to_num(B, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) # Mute the false BLAS multithreading signals with np.errstate(all='ignore'): prod = A_clean @ B_clean.T return np.clip(1.0 - prod, 0.0, 2.0) / 2.0 # normalized to [0,1] # ---------------------------------------------------------------------------- # 5. OTAP: unbalanced fused Gromov-Wasserstein solver # ---------------------------------------------------------------------------- @dataclass class OtapConfig: alpha: float = 0.35 # action beta: float = 0.20 # context gamma: float = 0.25 # effect delta: float = 0.20 # tool (eta/state = 0: no env states in this corpus) theta: float = 0.30 # structure vs feature trade-off lam: float = 1.0 # KL marginal relaxation strength eps: float = 0.05 # entropic regularization n_outer: int = 20 n_sink: int = 200 signed_D: bool = True balanced: bool = False hungarian_transport: bool = False # ablation (iv) whole_step_embed: bool = False # ablation (v) use_tool: bool = True # ablation (vi) def node_cost(P: Trajectory, R: Trajectory, enc: Encoder, cfg: OtapConfig) -> np.ndarray: if cfg.whole_step_embed: Ep = enc.encode([s.text + " -> " + s.effect for s in P.steps]) Er = enc.encode([s.text + " -> " + s.effect for s in R.steps]) C = cos_dist(Ep, Er) else: Ap, Ar = enc.encode([s.action for s in P.steps]), enc.encode([s.action for s in R.steps]) Cp, Cr = enc.encode([s.context for s in P.steps]), enc.encode([s.context for s in R.steps]) Fp, Fr = enc.encode([s.effect for s in P.steps]), enc.encode([s.effect for s in R.steps]) C = cfg.alpha * cos_dist(Ap, Ar) + cfg.beta * cos_dist(Cp, Cr) \ + cfg.gamma * cos_dist(Fp, Fr) C /= (cfg.alpha + cfg.beta + cfg.gamma) C *= (1 - cfg.delta) if cfg.use_tool: T = np.array([[tool_dist(a.tool, b.tool) for b in R.steps] for a in P.steps]) C = C + cfg.delta * T return np.clip(C, 0, 1) def unbalanced_sinkhorn_log(M, mu, nu, eps, lam, n_iter): """Log-domain unbalanced Sinkhorn (KL marginal penalties).""" fi = lam / (lam + eps) log_mu, log_nu = np.log(mu + 1e-300), np.log(nu + 1e-300) f = np.zeros(len(mu)); g = np.zeros(len(nu)) Meps = -M / eps for _ in range(n_iter): f = fi * eps * (log_mu - logsumexp(Meps + g[None, :] / eps, axis=1)) g = fi * eps * (log_nu - logsumexp(Meps + f[:, None] / eps, axis=0)) T = np.exp(Meps + f[:, None] / eps + g[None, :] / eps) return T def gw_grad(Dp, Dr, T): """G_ij = sum_kl (Dp_ik - Dr_jl)^2 T_kl (squared loss).""" p = T.sum(axis=1); q = T.sum(axis=0) Dp_clean = np.nan_to_num(Dp, nan=0.0, posinf=0.0, neginf=0.0) Dr_clean = np.nan_to_num(Dr, nan=0.0, posinf=0.0, neginf=0.0) T_clean = np.nan_to_num(T, nan=0.0, posinf=0.0, neginf=0.0) return ((Dp_clean ** 2) @ p)[:, None] + (q @ (Dr_clean ** 2).T)[None, :] - 2 * Dp_clean @ T_clean @ Dr_clean.T def kl_mass(a, b): return float(np.sum(np.where(a > 0, a * np.log((a + 1e-300) / (b + 1e-300)), 0) - a + b)) def otap_loss(P: Trajectory, R: Trajectory, enc: Encoder, cfg: OtapConfig): n, m = len(P.steps), len(R.steps) mu = np.full(n, 1.0 / n); nu = np.full(m, 1.0 / m) C = node_cost(P, R, enc, cfg) Dp, Dr = P.signed_dep_matrix(), R.signed_dep_matrix() if not cfg.signed_D: Dp, Dr = np.abs(Dp), np.abs(Dr) if cfg.hungarian_transport: # ablation (iv): hard 1-1 big = max(n, m) Cf = np.ones((big, big)) Cf[:n, :m] = (1 - cfg.theta) * C ri, ci = linear_sum_assignment(Cf) T = np.zeros((n, m)) for i, j in zip(ri, ci): if i < n and j < m: T[i, j] = 1.0 / big else: lam = 1e6 if cfg.balanced else cfg.lam T = np.outer(mu, nu) for _ in range(cfg.n_outer): M = (1 - cfg.theta) * C + cfg.theta * gw_grad(Dp, Dr, T) Tn = unbalanced_sinkhorn_log(M, mu, nu, cfg.eps, lam, cfg.n_sink) T = 0.5 * T + 0.5 * Tn feat = float(np.sum(C * T)) gw = float(np.sum(gw_grad(Dp, Dr, T) * T)) kl = cfg.lam * (kl_mass(T.sum(1), mu) + kl_mass(T.sum(0), nu)) loss = (1 - cfg.theta) * feat + cfg.theta * gw + kl prec = float(T.sum()) / 1.0 rec = float(T.sum()) / 1.0 return loss, {"feat": feat, "gw": gw, "kl": kl, "precision": min(1.0, prec), "recall": min(1.0, rec)} def otap_score(P, refs: List[Trajectory], enc, cfg, temp=0.1) -> float: losses = [otap_loss(P, R, enc, cfg)[0] for R in refs] L = -temp * logsumexp(np.array([-l / temp for l in losses]) - math.log(len(losses))) return math.exp(-2.0 * L) # ---------------------------------------------------------------------------- # 6. Baselines # ---------------------------------------------------------------------------- def exact_match_score(P: Trajectory, R: Trajectory) -> float: n, m = len(P.steps), len(R.steps) hits = sum(1 for i in range(min(n, m)) if P.steps[i].text == R.steps[i].text) return hits / max(n, m) def _ngrams(tokens, n): return [tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)] def bleu_score(P: Trajectory, R: Trajectory) -> float: hyp = " ".join(s.text for s in P.steps).lower().split() ref = " ".join(s.text for s in R.steps).lower().split() if not hyp or not ref: return 0.0 logs = [] for n in range(1, 5): h, r = _ngrams(hyp, n), _ngrams(ref, n) if not h: logs.append(math.log(1e-9)); continue rc: Dict[tuple, int] = {} for g in r: rc[g] = rc.get(g, 0) + 1 match = 0 for g in set(h): match += min(h.count(g), rc.get(g, 0)) p = (match + 1) / (len(h) + 1) # add-1 smoothing logs.append(math.log(p)) bp = min(1.0, math.exp(1 - len(ref) / max(1, len(hyp)))) return bp * math.exp(sum(logs) / 4) def greedy_f1_score(P, R, enc) -> float: # BERTScore-style Ep = enc.encode([s.text for s in P.steps]) Er = enc.encode([s.text for s in R.steps]) Ep_clean = np.nan_to_num(Ep, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) Er_clean = np.nan_to_num(Er, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) # Mute the false BLAS multithreading signals with np.errstate(all='ignore'): S = Ep_clean @ Er_clean.T prec = float(S.max(axis=1).mean()) if S.size > 0 else 0.0 rec = float(S.max(axis=0).mean()) if S.size > 0 else 0.0 return 0.0 if prec + rec == 0 else 2 * prec * rec / (prec + rec) def hungarian_score(P, R, enc) -> float: Ep = enc.encode([s.text for s in P.steps]) Er = enc.encode([s.text for s in R.steps]) C = cos_dist(Ep, Er) n, m = C.shape big = max(n, m) Cf = np.ones((big, big)); Cf[:n, :m] = C ri, ci = linear_sum_assignment(Cf) return 1.0 - Cf[ri, ci].sum() / big def llm_judge_score(P: Trajectory, R: Trajectory, model="claude-sonnet-4-6") -> Optional[float]: key = os.environ.get("ANTHROPIC_API_KEY") if not key: return None import urllib.request fmt = lambda t: "\n".join(f"{i+1}. {s.text} [tool: {s.tool}]" for i, s in enumerate(t.steps)) prompt = (f"Task: {P.task}\n\nReference plan:\n{fmt(R)}\n\n" f"Candidate plan:\n{fmt(P)}\n\n" "Rate how well the candidate plan solves the task relative to the " "reference, considering correctness of steps, ordering/causality, " "tool choice, and absence of missing or irrelevant steps. Valid " "alternative orderings and granularities deserve full credit. " "Answer with ONLY an integer 1-10.") body = json.dumps({"model": model, "max_tokens": 8, "messages": [{"role": "user", "content": prompt}]}).encode() req = urllib.request.Request( "https://api.anthropic.com/v1/messages", data=body, headers={"x-api-key": key, "anthropic-version": "2023-06-01", "content-type": "application/json"}) try: with urllib.request.urlopen(req, timeout=60) as r: out = json.loads(r.read()) txt = "".join(b.get("text", "") for b in out.get("content", [])) mnum = re.search(r"\d+", txt) return float(mnum.group()) / 10.0 if mnum else None except Exception as e: print(f" [llm-judge] call failed: {e}", file=sys.stderr) return None # ---------------------------------------------------------------------------- # 7. Part A meta-evaluation: pairwise ranking accuracy per family # ---------------------------------------------------------------------------- FAMS = ["P1", "P2", "P3", "P4", "P5"] def pra_by_family(records: List[dict], metric: str) -> Dict[str, float]: by_task: Dict[int, List[dict]] = {} for r in records: by_task.setdefault(r["task_id"], []).append(r) res = {} for fam in FAMS: ok = tot = 0 for recs in by_task.values(): xs = [r for r in recs if r["tag"] == fam] for x in xs: sx = x["scores"][metric] if fam in ("P1", "P2", "P3"): others = [r for r in recs if r["tier"] >= 1] for o in others: tot += 1; ok += int(sx > o["scores"][metric]) elif fam == "P4": for o in [r for r in recs if r["tier"] == 0]: tot += 1; ok += int(sx < o["scores"][metric]) for o in [r for r in recs if r["tier"] == 2]: tot += 1; ok += int(sx > o["scores"][metric]) else: for o in [r for r in recs if r["tier"] <= 1]: tot += 1; ok += int(sx < o["scores"][metric]) res[fam] = 100.0 * ok / tot if tot else float("nan") return res # ---------------------------------------------------------------------------- # 8. Part B harness (requires human annotations; NOT run automatically) # ---------------------------------------------------------------------------- def part_b_human_correlation(annotations_csv: str, results_json: str): from scipy.stats import spearmanr, kendalltau import csv with open(results_json) as f: data = json.load(f) scores = {r["uid"]: r["scores"] for r in data["records"]} ann = {} with open(annotations_csv) as f: for row in csv.DictReader(f): ann[row["trajectory_uid"]] = float(row["human_score"]) uids = [u for u in ann if u in scores] metrics = list(next(iter(scores.values())).keys()) print(f"{'metric':24s} Spearman Kendall (N={len(uids)})") for mname in metrics: xs = [scores[u][mname] for u in uids] ys = [ann[u] for u in uids] print(f"{mname:24s} {spearmanr(xs, ys).statistic:+.3f} " f"{kendalltau(xs, ys).statistic:+.3f}") # ---------------------------------------------------------------------------- # 9. Main # ---------------------------------------------------------------------------- def latex_row(name, d): return (f"{name:24s} & " + " & ".join( f"{d[f]:.1f}" if not math.isnan(d[f]) else "--" for f in FAMS) + r" \\") def main(): ap = argparse.ArgumentParser() ap.add_argument("--seed", type=int, default=7) ap.add_argument("--n-expansions", type=int, default=5, help="task instances per domain template (max 5)") ap.add_argument("--encoder", choices=["auto", "sbert", "tfidf"], default="auto") ap.add_argument("--llm-judge", action="store_true") ap.add_argument("--ablations", action="store_true") ap.add_argument("--outdir", default="results") args = ap.parse_args() rng = random.Random(args.seed) np.random.seed(args.seed) os.makedirs(args.outdir, exist_ok=True) mode = args.encoder if mode == "auto": try: import sentence_transformers # noqa mode = "sbert" except ImportError: mode = "tfidf" print("[warn] sentence-transformers not installed; falling back to " "TF-IDF embeddings. Install it for the paper-quality run.") corpus = build_corpus(args.n_expansions) print(f"Corpus: {len(corpus)} tasks across {len(DOMAIN_TEMPLATES)} domains " f"| encoder = {mode}") all_variants: List[Tuple[int, Trajectory]] = [] for tid, tr in enumerate(corpus): for v in make_variants(tr, rng): all_variants.append((tid, v)) print(f"Generated {len(all_variants)} trajectories " f"({sum(1 for _, v in all_variants if v.tag != 'original')} perturbed)") texts = set() for _, v in all_variants: for s in v.steps: for k in range(len(s.action_variants)): texts.add(s.action_variants[k]) texts.update([s.context, s.effect, s.text, s.text + " -> " + s.effect]) enc = Encoder(mode, sorted(texts)) cfg = OtapConfig() records = [] t0 = time.time() for idx, (tid, v) in enumerate(all_variants): R = corpus[tid] sc = { "exact": exact_match_score(v, R), "bleu": bleu_score(v, R), "greedy_f1": greedy_f1_score(v, R, enc), "hungarian": hungarian_score(v, R, enc), "otap": otap_score(v, [R], enc, cfg), } if args.llm_judge: j = llm_judge_score(v, R) if j is not None: sc["llm_judge"] = j records.append({"uid": f"t{tid}_{v.tag}_{idx}", "task_id": tid, "tag": v.tag, "tier": TIER[v.tag], "scores": sc}) print(f"Scored all trajectories in {time.time()-t0:.1f}s") metric_names = ["exact", "bleu", "greedy_f1", "hungarian"] if args.llm_judge and any("llm_judge" in r["scores"] for r in records): records = [r for r in records if "llm_judge" in r["scores"]] metric_names.append("llm_judge") metric_names.append("otap") pretty = {"exact": "Exact step match", "bleu": "BLEU", "greedy_f1": "Embedding greedy F1", "hungarian": "Embedding + Hungarian", "llm_judge": "LLM-as-a-judge", "otap": r"\otap{} (ours)"} print("\n=== Table I: pairwise ranking accuracy (%) per family ===") header = f"{'metric':24s} " + " ".join(f"{f:>5s}" for f in FAMS) print(header); print("-" * len(header)) table_rows = [] summary = {} for mname in metric_names: d = pra_by_family(records, mname) summary[mname] = d print(f"{pretty[mname]:24s} " + " ".join(f"{d[f]:5.1f}" for f in FAMS)) table_rows.append(latex_row(pretty[mname], d)) # mean score per family (separability diagnostic) print("\n=== Mean metric score per family (diagnostic) ===") fams_all = ["original"] + FAMS print(f"{'metric':24s} " + " ".join(f"{f:>8s}" for f in fams_all)) for mname in metric_names: row = [] for f in fams_all: vals = [r["scores"][mname] for r in records if r["tag"] == f] row.append(np.mean(vals) if vals else float("nan")) print(f"{pretty[mname]:24s} " + " ".join(f"{v:8.3f}" for v in row)) tex = ["% Table I rows (auto-generated by run_all.py; encoder=" + mode + ")"] tex += table_rows with open(os.path.join(args.outdir, "results_tables.tex"), "w") as f: f.write("\n".join(tex) + "\n") # ---------------- ablations ---------------- if args.ablations: abls = { r"\otap{} full": OtapConfig(), r"-- structure ($\theta=0$)": OtapConfig(theta=0.0), "-- unsigned $D$": OtapConfig(signed_D=False), "-- balanced marginals": OtapConfig(balanced=True), "-- Hungarian transport": OtapConfig(hungarian_transport=True), "-- whole-step embedding": OtapConfig(whole_step_embed=True), "-- no tool term": OtapConfig(use_tool=False), } print("\n=== Ablations (PRA %) ===") abl_rows = [] for name, acfg in abls.items(): for r_ in records: tid = r_["task_id"] v = next(v for t2, v in all_variants if t2 == tid and f"t{tid}_{v.tag}_" in r_["uid"] or True) for r_, (tid, v) in zip(records, all_variants[:len(records)]): r_["scores"]["_abl"] = otap_score(v, [corpus[tid]], enc, acfg) d = pra_by_family(records, "_abl") print(f"{name:32s} " + " ".join(f"{d[f]:5.1f}" for f in FAMS)) abl_rows.append(latex_row(name, d)) with open(os.path.join(args.outdir, "ablations.tex"), "w") as f: f.write("\n".join(abl_rows) + "\n") with open(os.path.join(args.outdir, "results.json"), "w") as f: json.dump({"encoder": mode, "seed": args.seed, "n_tasks": len(corpus), "n_trajectories": len(records), "summary_pra": summary, "records": records}, f, indent=1) print(f"\nWrote {args.outdir}/results.json and " f"{args.outdir}/results_tables.tex") if __name__ == "__main__": main()