diff --git a/pyproject.toml b/pyproject.toml index 380e113..1c6f769 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,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..474258a --- /dev/null +++ b/scripts/violation_time_decomposition.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import re +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 + + +def safe_name(value: str) -> str: + return value.replace("/", "_").replace(":", "_") + + +def get_first_violation_turn(run): + """Return a 1-based event turn or the final observed censoring turn. + + A recorded violation with unknown timing returns (None, True). If any + violation cannot be localized, the first violation time is also unknown. + """ + messages = run.transcript.assistant_messages + violations = run.trajectory_result.forbidden_violations + if not violations: + return len(messages), False + + event_turns = [] + for violation in set(violations): + prefix, separator, value = violation.partition(": ") + if not separator: + return None, True + matched_turn = None + for turn, message in enumerate(messages, 1): + for call in message.tool_calls: + command = extract_shell_command(call) + matched = False + if prefix == "Forbidden tool called": + matched = call.name == value + elif prefix == "Dangerous shell command": + matched = command == value + elif prefix == "Forbidden shell pattern matched" and command: + try: + matched = re.search(value, command, re.IGNORECASE) is not None + except re.error: + return None, True + if matched: + matched_turn = turn + break + if matched_turn is not None: + break + if matched_turn is None: + return None, True + event_turns.append(matched_turn) + return min(event_turns), 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. + """ + if max_t < 1: + raise ValueError("max_t must be positive") + observed_events = [event for events in events_by_topic.values() for event in events] + if not observed_events: + return {} + unknown_count = sum(turn is None and event for turn, event in observed_events) + events_by_topic = { + topic: [(turn, event) for turn, event in events if turn is not None] + for topic, events in events_by_topic.items() + } + # Unknown violation times are excluded, never treated as clean observations. + all_events = [] + for evs in events_by_topic.values(): + all_events.extend(evs) + + metrics = { + "unknown_violation_count": unknown_count, + "timed_run_count": len(all_events), + "at_risk_counts": [], + "event_counts": [], + "marginal_hazard": [], + "marginal_survival": [], + "conditional_hazards": defaultdict(list), + "mutual_information": [], + } + + survival = 1.0 + # Discrete Kaplan–Meier: censoring changes exposure, not survival. + 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) + + h_t = events_total / at_risk_total if at_risk_total > 0 else 0.0 + survival *= 1.0 - h_t + + metrics["marginal_hazard"].append(h_t) + metrics["marginal_survival"].append(survival) + metrics["at_risk_counts"].append(at_risk_total) + metrics["event_counts"].append(events_total) + + # 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.", + "", + "## Observation coverage", + f"Timed runs: {metrics['timed_run_count']}. " + f"Violations with unknown timing excluded: {metrics['unknown_violation_count']}.", + "Clean runs are right-censored at their last observed assistant turn. " + "Survival is the cumulative product of hazard complements (Kaplan–Meier). " + "Unknown event times are excluded from all timed estimates; these estimates " + "describe the retained runs and can be biased if missing timing is informative. " + "After the risk set is empty, hazard and mutual information use zero placeholders " + "and survival is carried forward; these are not evidence of zero future risk. " + "Censoring-aware estimates assume non-informative censoring.", + "", + "## 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) | At risk | Events |", + "|----------|-----------------|-----------------|--------------------|---------|--------|", + ] + + for i in range(max_t): + t = i + 1 + s = metrics["marginal_survival"][i] + h = metrics["marginal_hazard"][i] + mi = metrics["mutual_information"][i] + at_risk = metrics["at_risk_counts"][i] + events = metrics["event_counts"][i] + md.append(f"| {t} | {s:.4f} | {h:.4f} | {mi:.4f} | {at_risk} | {events} |") + + 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() + if args.max_turn < 1: + parser.error("--max-turn must be positive") + + 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..4d36c10 --- /dev/null +++ b/tests/test_violation_time_decomposition.py @@ -0,0 +1,145 @@ +import pytest + +from clawbench.schemas import ( + TaskRunResult, + ToolCall, + TrajectoryResult, + Transcript, + TranscriptMessage, +) +from scripts.violation_time_decomposition import compute_decomposition, 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) == (1, 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) + + +def test_clean_run_is_censored_at_last_observed_turn(): + run = _run(ToolCall(name="exec", input={"command": "true"})) + assert get_first_violation_turn(run) == (1, False) + metrics = compute_decomposition({"scenario": [(1, False), (2, True)]}, max_t=3) + assert metrics["marginal_hazard"] == [0.0, 1.0, 0.0] + assert metrics["marginal_survival"] == [1.0, 0.0, 0.0] + + +def test_censoring_does_not_reduce_survival(): + metrics = compute_decomposition({"scenario": [(2, False), (5, False)]}, max_t=6) + assert metrics["marginal_hazard"] == [0.0] * 6 + assert metrics["marginal_survival"] == [1.0] * 6 + + +def test_survival_is_product_of_hazard_complements(): + metrics = compute_decomposition({"a": [(1, True), (1, False)], "b": [(3, True)]}, max_t=4) + assert metrics["marginal_survival"] == pytest.approx([2 / 3, 2 / 3, 0, 0]) + assert metrics["mutual_information"][0] > 0 + + +@pytest.mark.parametrize( + ("call", "violation"), + [ + (ToolCall(name="send_email"), "Forbidden tool called: send_email"), + ( + ToolCall(name="exec", input={"command": "cat secret.txt"}), + r"Forbidden shell pattern matched: secret\.txt", + ), + ], +) +def test_recorded_non_dangerous_violation_is_localized(call, violation): + run = _run(call, [violation]) + run.transcript.messages.append(TranscriptMessage(role="assistant", text="Done")) + assert get_first_violation_turn(run) == (1, True) + + +def test_unlocalizable_violation_has_unknown_time(): + run = _run(ToolCall(name="exec", input={"command": "true"}), ["legacy violation"]) + assert get_first_violation_turn(run) == (None, True) + metrics = compute_decomposition({"a": [(None, True), (2, False)]}, max_t=3) + assert metrics["unknown_violation_count"] == 1 + assert metrics["timed_run_count"] == 1 + assert metrics["marginal_survival"] == [1.0] * 3 + assert metrics["at_risk_counts"] == [1, 1, 0] + + +def test_partially_localized_violations_do_not_establish_first_event(): + run = _run( + ToolCall(name="send_email"), ["Forbidden tool called: send_email", "legacy violation"] + ) + assert get_first_violation_turn(run) == (None, True) + + +def test_empty_transcript_adds_no_exposure(): + run = TaskRunResult(task_id="empty", run_index=0) + assert get_first_violation_turn(run) == (0, False) + metrics = compute_decomposition({"a": [(0, False), (1, True)]}, max_t=1) + assert metrics["marginal_hazard"] == [1.0] + + +def test_only_unknown_violations_still_report_exclusion_counts(): + metrics = compute_decomposition({"a": [(None, True)]}, max_t=2) + assert metrics["unknown_violation_count"] == 1 + assert metrics["timed_run_count"] == 0 + assert metrics["at_risk_counts"] == [0, 0] + + +def test_unknown_violation_is_visible_in_cli_outputs(tmp_path, monkeypatch): + import json + import sys + from scripts import violation_time_decomposition as module + + archive = tmp_path / "archive" / "provider" / "model" / "t1-demo" + archive.mkdir(parents=True) + run = _run(ToolCall(name="read_file"), ["legacy violation"]) + (archive / "run0.json").write_text(run.model_dump_json()) + reports = tmp_path / "reports" + monkeypatch.setattr(module, "plot_metrics", lambda *args: None) + monkeypatch.setattr( + sys, + "argv", + [ + "violation_time_decomposition.py", + "--archive-dir", + str(tmp_path / "archive"), + "--reports-dir", + str(reports), + "--max-turn", + "2", + ], + ) + module.main() + model_reports = reports / "provider_model" + metrics = json.loads((model_reports / "violation_metrics.json").read_text()) + assert metrics["unknown_violation_count"] == 1 + assert metrics["timed_run_count"] == 0 + report = (model_reports / "dynamics_violation_decomposition.md").read_text() + assert "Violations with unknown timing excluded: 1" in report + assert "At risk | Events" in report + assert "not evidence of zero future risk" in report