diff --git a/pyproject.toml b/pyproject.toml index a952a1c..46b2ef6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,6 +62,7 @@ allow-direct-references = true [tool.pytest.ini_options] asyncio_mode = "auto" addopts = ["-p", "no:opik"] +pythonpath = ["."] testpaths = ["tests"] [tool.ruff] diff --git a/scripts/classify_regimes.py b/scripts/classify_regimes.py index 6fbc4c7..fb58ef2 100644 --- a/scripts/classify_regimes.py +++ b/scripts/classify_regimes.py @@ -145,6 +145,7 @@ def main() -> None: all_turn_texts: list[str] = [] run_turns: dict[str, list[str]] = {} + run_metadata: dict[str, dict[str, str | int]] = {} for model_name, task_runs in grouped.items(): for task_id, runs in task_runs.items(): @@ -152,6 +153,11 @@ def main() -> None: ts = turn_texts(run, fallback_any_message=False) key = f"{model_name}/{task_id}/run{run.run_index}" run_turns[key] = ts + run_metadata[key] = { + "model": model_name, + "task_id": task_id, + "run_index": run.run_index, + } all_turn_texts.extend(ts) used_fallback_messages = False @@ -159,12 +165,18 @@ def main() -> None: used_fallback_messages = True all_turn_texts = [] run_turns = {} + run_metadata = {} for model_name, task_runs in grouped.items(): for task_id, runs in task_runs.items(): for run in runs: ts = turn_texts(run, fallback_any_message=True) key = f"{model_name}/{task_id}/run{run.run_index}" run_turns[key] = ts + run_metadata[key] = { + "model": model_name, + "task_id": task_id, + "run_index": run.run_index, + } all_turn_texts.extend(ts) if not all_turn_texts: @@ -205,6 +217,7 @@ def main() -> None: for key, metrics in per_run.items(): metrics["regime"] = classify(metrics, thresholds) metrics["turn_source"] = "any_message" if used_fallback_messages else "assistant" + metrics.update(run_metadata.get(key, {})) args.reports_dir.mkdir(parents=True, exist_ok=True) out = args.reports_dir / "regimes.json" diff --git a/scripts/compute_debiased_dynamics.py b/scripts/compute_debiased_dynamics.py index 6df01df..336fb47 100644 --- a/scripts/compute_debiased_dynamics.py +++ b/scripts/compute_debiased_dynamics.py @@ -34,9 +34,9 @@ def compute_debiased_dynamics(regimes_path, constraint_path, weights_path, topic model_regimes_weight_sum = defaultdict(float) for key, data in regimes.items(): - parts = key.split("/") - model = parts[0] - task_id = parts[1] if len(parts) > 1 else parts[0] + key_parts = key.rsplit("/", 2) + model = data.get("model") or (key_parts[0] if len(key_parts) == 3 else key) + task_id = data.get("task_id") or (key_parts[1] if len(key_parts) == 3 else key) # Match task to topic matched_topic = "unknown" diff --git a/scripts/generate_dynamical_report.py b/scripts/generate_dynamical_report.py index 55d52e8..bdcb087 100644 --- a/scripts/generate_dynamical_report.py +++ b/scripts/generate_dynamical_report.py @@ -79,7 +79,7 @@ def main() -> None: L("") by_model = defaultdict(Counter) for key, row in regimes.items(): - model = key.split("/")[0] + model = row.get("model") or key.rsplit("/", 2)[0] regime = row.get("regime", "unknown") by_model[model][regime] += 1 diff --git a/scripts/run_posterior_dynamics_pipeline.py b/scripts/run_posterior_dynamics_pipeline.py index 9c433a3..7a56b9e 100644 --- a/scripts/run_posterior_dynamics_pipeline.py +++ b/scripts/run_posterior_dynamics_pipeline.py @@ -88,6 +88,15 @@ def main() -> None: _run([py, str(scripts_dir / "variance_decomp.py"), "--archive-dir", str(archive_dir), "--reports-dir", str(reports_dir), *tier_args]) _run([py, str(scripts_dir / "survival_analysis.py"), "--archive-dir", str(archive_dir), "--reports-dir", str(reports_dir), *tier_args]) _run([py, str(scripts_dir / "snr_weighted_ranking.py"), "--archive-dir", str(archive_dir), "--reports-dir", str(reports_dir), *tier_args]) + _run([ + py, + str(scripts_dir / "violation_time_decomposition.py"), + "--archive-dir", + str(archive_dir), + "--reports-dir", + str(reports_dir / "violation_time_decomposition"), + *tier_args, + ]) _run([py, str(scripts_dir / "generate_dynamical_report.py"), "--reports-dir", str(reports_dir)]) if args.include_dynamics_report: _write_dynamics_reports(archive_dir, output_dir, args.tier) diff --git a/scripts/violation_time_decomposition.py b/scripts/violation_time_decomposition.py new file mode 100644 index 0000000..6f8500a --- /dev/null +++ b/scripts/violation_time_decomposition.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import sys +from pathlib import Path +from collections import defaultdict + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from clawbench.dynamics_archive import load_task_runs_by_model +from clawbench.trajectory import extract_shell_command, has_dangerous_shell_pattern + + +def safe_name(value: str) -> str: + return value.replace("/", "_").replace(":", "_") + +def get_first_violation_turn(run): + """ + Returns (turn_index, has_violation) + Turn index is 1-based. + If no violation, turn_index is the max turn + 1, has_violation is False. + """ + if not run.trajectory_result.forbidden_violations: + return len(run.transcript.assistant_messages) + 1, False + + for i, msg in enumerate(run.transcript.assistant_messages, 1): + for tc in msg.tool_calls: + command = extract_shell_command(tc) + if command and has_dangerous_shell_pattern(command): + return i, True + return max(len(run.transcript.assistant_messages), 1), True + +def compute_decomposition(events_by_topic, max_t=20): + """ + events_by_topic: dict mapping topic (e.g., scenario) to list of (turn, has_violation) + Returns decomposition metrics. + """ + # Flatten all events + all_events = [] + for evs in events_by_topic.values(): + all_events.extend(evs) + + total = len(all_events) + if total == 0: + return {} + + metrics = { + "marginal_hazard": [], + "marginal_survival": [], + "conditional_hazards": defaultdict(list), + "mutual_information": [] + } + + # Pre-calculate marginal + for t in range(1, max_t + 1): + at_risk_total = sum(1 for tf, _ in all_events if tf >= t) + events_total = sum(1 for tf, is_event in all_events if is_event and tf == t) + survived_total = sum(1 for tf, is_event in all_events if (not is_event and tf >= t) or (is_event and tf > t)) + + h_t = events_total / at_risk_total if at_risk_total > 0 else 0.0 + s_t = survived_total / total if total > 0 else 0.0 + + metrics["marginal_hazard"].append(h_t) + metrics["marginal_survival"].append(s_t) + + # Calculate conditional hazards + mi_t = 0.0 + for topic, evs in events_by_topic.items(): + at_risk_topic = sum(1 for tf, _ in evs if tf >= t) + events_topic = sum(1 for tf, is_event in evs if is_event and tf == t) + h_t_given_s = events_topic / at_risk_topic if at_risk_topic > 0 else 0.0 + metrics["conditional_hazards"][topic].append(h_t_given_s) + + # P(S = topic | T >= t) + if at_risk_total > 0 and at_risk_topic > 0: + p_s_given_at_risk = at_risk_topic / at_risk_total + + # MI term: P(S|T>=t) * D_KL( P(V|S) || P(V) ) + kl = 0.0 + if h_t_given_s > 0 and h_t > 0: + kl += h_t_given_s * math.log2(h_t_given_s / h_t) + if (1 - h_t_given_s) > 0 and (1 - h_t) > 0: + kl += (1 - h_t_given_s) * math.log2((1 - h_t_given_s) / (1 - h_t)) + + mi_t += p_s_given_at_risk * kl + + metrics["mutual_information"].append(mi_t) + + return metrics + +def plot_metrics(metrics, model_name, output_dir): + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ModuleNotFoundError: + return None + + max_t = len(metrics["marginal_hazard"]) + turns = list(range(1, max_t + 1)) + + plt.figure(figsize=(15, 5)) + + # Plot 1: Marginal Hazard + plt.subplot(1, 3, 1) + plt.plot(turns, metrics["marginal_hazard"], marker='o', color='red') + plt.title("Marginal Hazard $h(t)$") + plt.xlabel("Turn $t$") + plt.ylabel("P(Violation | No Prior Violation)") + plt.grid(True) + + # Plot 2: Conditional Hazards + plt.subplot(1, 3, 2) + for topic, h_cond in metrics["conditional_hazards"].items(): + if max(h_cond) > 0: # Only plot if there's non-zero hazard + plt.plot(turns, h_cond, alpha=0.5, label=topic) + plt.title("Conditional Hazards $h(t | Topic)$") + plt.xlabel("Turn $t$") + # plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize='small') + plt.grid(True) + + # Plot 3: Mutual Information + plt.subplot(1, 3, 3) + plt.plot(turns, metrics["mutual_information"], marker='o', color='purple') + plt.title(r"Mutual Information $I(Topic; V_t | T \geq t)$") + plt.xlabel("Turn $t$") + plt.ylabel("Bits") + plt.grid(True) + + plt.tight_layout() + plot_path = output_dir / f"violation_metrics_{model_name}.png" + plt.savefig(plot_path) + plt.close() + return plot_path + +def generate_markdown(metrics, model_name, plot_path, out_file): + max_t = len(metrics["marginal_hazard"]) + + md = [ + f"# Spatio-Temporal Violation Dynamics: {model_name}", + "", + "## Theoretical Decomposition", + "This report decomposes the time-to-first-violation for safety triggers and tool misuses.", + "To connect the long-term behavior of agent risk to its spatial risk conditioned on context semantics, we decompose the first occurrence probability:", + r"$$ P(T = t) = h(t) \cdot S(t-1) $$", + "where $h(t)$ is the conditional hazard rate at turn $t$, and $S(t-1)$ is the macro survival probability.", + "", + "Furthermore, we examine the mutual information between the semantic spatial context (scenario) and the violation event at each turn to determine if localized contexts explain hazard spikes.", + "", + "## Visualization", + f"![Violation Metrics]({plot_path.name})" if plot_path else "Plot generation skipped because matplotlib is not installed.", + "", + "## Empirical Metrics Table", + "| Turn $t$ | Marginal $S(t)$ | Marginal $h(t)$ | Mutual Info (bits) |", + "|----------|-----------------|-----------------|--------------------|" + ] + + for i in range(max_t): + t = i + 1 + s = metrics["marginal_survival"][i] + h = metrics["marginal_hazard"][i] + mi = metrics["mutual_information"][i] + md.append(f"| {t} | {s:.4f} | {h:.4f} | {mi:.4f} |") + + out_file.write_text("\n".join(md), encoding="utf-8") + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--archive-dir", type=Path, default=Path(".clawbench/run_cache")) + parser.add_argument( + "--reports-dir", + "--results-dir", + dest="reports_dir", + type=Path, + default=Path("reports/violation_time_decomposition"), + ) + parser.add_argument("--tier", choices=["tier1", "tier2", "tier3", "tier4", "tier5"], default=None) + parser.add_argument("--max-turn", type=int, default=15) + args = parser.parse_args() + + print(f"Loading runs from {args.archive_dir}...") + grouped = load_task_runs_by_model(args.archive_dir, tier=args.tier) + print(f"Loaded models: {list(grouped.keys())}") + + args.reports_dir.mkdir(parents=True, exist_ok=True) + + for model_name, task_runs in grouped.items(): + print(f"Processing model: {model_name}") + events_by_topic = defaultdict(list) + for task_id, runs in task_runs.items(): + for run in runs: + topic = run.scenario if run.scenario else "unknown" + events_by_topic[topic].append(get_first_violation_turn(run)) + + metrics = compute_decomposition(events_by_topic, max_t=args.max_turn) + if not metrics: + print(f"No events for {model_name}") + continue + + safe_model = safe_name(model_name) + model_out_dir = args.reports_dir / safe_model + model_out_dir.mkdir(parents=True, exist_ok=True) + + plot_path = plot_metrics(metrics, safe_model, model_out_dir) + doc_path = model_out_dir / "dynamics_violation_decomposition.md" + generate_markdown(metrics, model_name, plot_path, doc_path) + print(f"Generated doc for {model_name}: {doc_path}") + + # Dump JSON + json_path = model_out_dir / "violation_metrics.json" + json_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8") + +if __name__ == "__main__": + main() diff --git a/tests/test_violation_time_decomposition.py b/tests/test_violation_time_decomposition.py new file mode 100644 index 0000000..595001d --- /dev/null +++ b/tests/test_violation_time_decomposition.py @@ -0,0 +1,35 @@ +from clawbench.schemas import TaskRunResult, ToolCall, TrajectoryResult, Transcript, TranscriptMessage +from scripts.violation_time_decomposition import get_first_violation_turn + + +def _run(tool_call: ToolCall, forbidden_violations: list[str] | None = None) -> TaskRunResult: + return TaskRunResult( + task_id="t1-demo", + run_index=0, + trajectory_result=TrajectoryResult( + forbidden_violations=forbidden_violations or [], + ), + transcript=Transcript( + messages=[ + TranscriptMessage( + role="assistant", + tool_calls=[tool_call], + ) + ] + ), + ) + + +def test_failed_tool_call_is_not_counted_as_violation_without_trajectory_violation(): + run = _run(ToolCall(name="exec", input={"command": "pytest -q"}, success=False)) + + assert get_first_violation_turn(run) == (2, False) + + +def test_dangerous_command_violation_is_localized_to_turn(): + run = _run( + ToolCall(name="exec", input={"command": "sudo rm -rf /tmp/demo"}, success=True), + forbidden_violations=["Dangerous shell command: sudo rm -rf /tmp/demo"], + ) + + assert get_first_violation_turn(run) == (1, True)