heavycoderhh commited on
Commit
96a73f4
·
verified ·
1 Parent(s): 0490f93

Upload folder using huggingface_hub

Browse files
Dockerfile ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ # Multi-stage build using openenv-base
8
+ # This Dockerfile is flexible and works for both:
9
+ # - In-repo environments (with local OpenEnv sources)
10
+ # - Standalone environments (with openenv from PyPI/Git)
11
+ # The build script (openenv build) handles context detection and sets appropriate build args.
12
+
13
+ ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
14
+ FROM ${BASE_IMAGE} AS builder
15
+
16
+ WORKDIR /app
17
+
18
+ # Ensure git is available (required for installing dependencies from VCS)
19
+ RUN apt-get update && \
20
+ apt-get install -y --no-install-recommends git && \
21
+ rm -rf /var/lib/apt/lists/*
22
+
23
+ # Build argument to control whether we're building standalone or in-repo
24
+ ARG BUILD_MODE=in-repo
25
+ ARG ENV_NAME=counsel_env
26
+
27
+ # Copy environment code (always at root of build context)
28
+ COPY . /app/env
29
+
30
+ # For in-repo builds, openenv is already vendored in the build context
31
+ # For standalone builds, openenv will be installed via pyproject.toml
32
+ WORKDIR /app/env
33
+
34
+ # Ensure uv is available (for local builds where base image lacks it)
35
+ RUN if ! command -v uv >/dev/null 2>&1; then \
36
+ curl -LsSf https://astral.sh/uv/install.sh | sh && \
37
+ mv /root/.local/bin/uv /usr/local/bin/uv && \
38
+ mv /root/.local/bin/uvx /usr/local/bin/uvx; \
39
+ fi
40
+
41
+ # Install dependencies using uv sync
42
+ # If uv.lock exists, use it; otherwise resolve on the fly
43
+ RUN --mount=type=cache,target=/root/.cache/uv \
44
+ if [ -f uv.lock ]; then \
45
+ uv sync --frozen --no-install-project --no-editable; \
46
+ else \
47
+ uv sync --no-install-project --no-editable; \
48
+ fi
49
+
50
+ RUN --mount=type=cache,target=/root/.cache/uv \
51
+ if [ -f uv.lock ]; then \
52
+ uv sync --frozen --no-editable; \
53
+ else \
54
+ uv sync --no-editable; \
55
+ fi
56
+
57
+ # Final runtime stage
58
+ FROM ${BASE_IMAGE}
59
+
60
+ WORKDIR /app
61
+
62
+ # Copy the virtual environment from builder
63
+ COPY --from=builder /app/env/.venv /app/.venv
64
+
65
+ # Copy the environment code
66
+ COPY --from=builder /app/env /app/env
67
+
68
+ # Set PATH to use the virtual environment
69
+ ENV PATH="/app/.venv/bin:$PATH"
70
+
71
+ # Set PYTHONPATH so imports work correctly
72
+ ENV PYTHONPATH="/app/env:$PYTHONPATH"
73
+
74
+ # Health check
75
+ HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
76
+ CMD curl -f http://localhost:8000/health || exit 1
77
+
78
+ # Run the FastAPI server
79
+ # The module path is constructed to work with the /app/env structure
80
+ ENV ENABLE_WEB_INTERFACE=true
81
+ CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]
README.md CHANGED
@@ -1,10 +1,231 @@
1
  ---
2
  title: Counsel Env
3
- emoji: 🏢
4
- colorFrom: green
5
- colorTo: gray
6
  sdk: docker
7
  pinned: false
 
 
 
 
 
 
 
 
 
8
  ---
9
 
10
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
  title: Counsel Env
3
+ emoji: "⚖️"
4
+ colorFrom: indigo
5
+ colorTo: yellow
6
  sdk: docker
7
  pinned: false
8
+ app_port: 8000
9
+ base_path: /web
10
+ tags:
11
+ - openenv
12
+ - rl
13
+ - multi-agent
14
+ - theory-of-mind
15
+ - adversarial-dialogue
16
+ - procgen
17
  ---
18
 
19
+ # Cross-Examination Arena
20
+
21
+ We built a courtroom where an LLM learns to catch lies.
22
+
23
+ Counsel-Env trains an agent to act like a sharper trial lawyer: make a deterministic witness commit to a claim, then present the exhibit that proves the claim false. It is built for OpenEnv and GRPO-style post-training on multi-agent, partially observable dialogue.
24
+
25
+ > Baseline behavior: vague questions, early evidence, zero reward.
26
+ >
27
+ > Target behavior: trigger sealed claim, present matching exhibit, surface contradiction.
28
+
29
+ ![Held-out baseline reward comparison](../assets/plots/baseline_vs_oracle.svg)
30
+
31
+ ```mermaid
32
+ flowchart LR
33
+ A[Procgen Case Generator] --> B[Deterministic Witness]
34
+ B --> C[Agent Tools]
35
+ C --> D[Transcript + State Tracker]
36
+ D --> E[Weighted Reward Rubric]
37
+ E --> F[GRPO Training Loop]
38
+ ```
39
+
40
+ ## Environment Design
41
+
42
+ Each episode is a procedurally generated case with:
43
+
44
+ - a public case brief
45
+ - deterministic witness story
46
+ - hidden contradiction objects
47
+ - evidence exhibits with descriptions visible to the agent
48
+ - a 15-question budget
49
+ - replayable seeds for before/after evaluation
50
+
51
+ A contradiction is surfaced only when both steps happen in order:
52
+
53
+ 1. The agent asks a trigger question and the witness gives a sealed claim.
54
+ 2. The agent presents the matching disprover exhibit.
55
+
56
+ The witness is deterministic by design. That keeps reward verification fast, cheap, reproducible, and non-LLM-judged.
57
+
58
+ ## Actions
59
+
60
+ | Tool | Field | Purpose |
61
+ | --- | --- | --- |
62
+ | `ask_question` | `text` | Ask the witness a question. |
63
+ | `present_evidence` | `exhibit_id` | Present an exhibit from `available_evidence`. |
64
+ | `make_objection` | `reason` | Currently penalized unless an objection window exists. |
65
+ | `rest_case` | none | End the episode and receive final reward. |
66
+
67
+ ## Reward
68
+
69
+ Primary reward remains binary per contradiction:
70
+
71
+ ```text
72
+ primary_reward = contradictions_surfaced / contradictions_total
73
+ ```
74
+
75
+ Auxiliary shaping reduces sparsity while staying secondary:
76
+
77
+ ```text
78
+ auxiliary =
79
+ +0.2 * contradictions_triggered
80
+ +0.1 * trigger_keyword_questions
81
+ +0.1 * correctly_timed_evidence
82
+ -0.05 * duplicate_or_irrelevant_questions
83
+ -0.05 * blind_evidence
84
+ -0.1 * inadmissible_actions
85
+
86
+ total_reward = 0.8 * primary_reward + 0.2 * auxiliary
87
+ ```
88
+
89
+ Diagnostics log primary, auxiliary, and total reward separately.
90
+
91
+ ## Reward-Hacking Audit
92
+
93
+ The local evaluator compares four baselines:
94
+
95
+ | Agent | Avg Reward | Primary Reward | Trigger Rate | Surface Rate |
96
+ | --- | ---: | ---: | ---: | ---: |
97
+ | random | 0.000 | 0.000 | 0.000 | 0.000 |
98
+ | keyword_spam | 0.073 | 0.000 | 0.678 | 0.000 |
99
+ | present_all | 0.000 | 0.000 | 0.000 | 0.000 |
100
+ | scripted_oracle | 0.902 | 0.950 | 0.950 | 0.950 |
101
+
102
+ This shows the obvious hacks do not get primary reward: keyword spam can trigger claims, and present-all can burn exhibits, but neither surfaces contradictions.
103
+
104
+ ![Reward-hacking audit metrics](../assets/plots/rubric_breakdown.svg)
105
+
106
+ ## Demo Walkthrough
107
+
108
+ The fastest way to understand the environment is the same-seed demo:
109
+
110
+ ```text
111
+ assets/demo/demo_case.md
112
+ assets/transcripts/before_after_pairs.md
113
+ ```
114
+
115
+ In the demo case, the random baseline presents timestamped surveillance evidence before the witness has committed to a false time, so the witness gives no useful reaction. The successful strategy first asks about time, gets the witness to claim the assault happened at the wrong hour, then presents the timestamped footage.
116
+
117
+ That sequence is the capability we want training to amplify.
118
+
119
+ ## Curriculum
120
+
121
+ `generate_case()` supports:
122
+
123
+ - `difficulty="easy"`: 1 contradiction, obvious trigger language
124
+ - `difficulty="medium"`: 2 contradictions, overlapping trigger language
125
+ - `difficulty="hard"`: 3+ contradictions, distractor evidence
126
+ - `curriculum_stage` and custom sampling distributions
127
+
128
+ ## Local Validation
129
+
130
+ From the repo root:
131
+
132
+ ```bash
133
+ python pre_hf_validate.py
134
+ python -m pytest -p no:cacheprovider -q
135
+ ```
136
+
137
+ Latest local result:
138
+
139
+ ```text
140
+ 20 passed
141
+ ```
142
+
143
+ The preflight also checks local FastAPI + WebSocket reset/step/state, Docker-style imports, notebook credit guards, rollout diagnostics, and held-out evaluation artifacts.
144
+
145
+ ## Diagnostics
146
+
147
+ Generate a no-credit local diagnostic sample:
148
+
149
+ ```bash
150
+ python -m counsel_env.server.diagnostics
151
+ ```
152
+
153
+ Generate held-out baseline artifacts:
154
+
155
+ ```bash
156
+ python scripts/run_local_eval.py --episodes 30 --output-dir assets
157
+ ```
158
+
159
+ Outputs:
160
+
161
+ - `assets/heldout_eval.jsonl`
162
+ - `assets/heldout_eval_summary.json`
163
+ - `assets/plots/baseline_vs_oracle.svg`
164
+ - `assets/plots/rubric_breakdown.svg`
165
+ - `assets/transcripts/before_after_pairs.md`
166
+ - `assets/demo/demo_case.md`
167
+ - `assets/demo/video_script.md`
168
+ - `assets/demo/blog_draft.md`
169
+
170
+ ## Training
171
+
172
+ The notebook is at:
173
+
174
+ ```text
175
+ counsel_env/notebooks/train_counsel.ipynb
176
+ ```
177
+
178
+ It is credit-safe by default:
179
+
180
+ - `RUN_TRAINING=0`
181
+ - `push_to_hub=False`
182
+ - `report_to="none"`
183
+
184
+ GRPO exploration settings:
185
+
186
+ - `num_generations=6`
187
+ - `temperature=0.8`
188
+ - `top_p=0.95`
189
+ - `max_completion_length=1024`
190
+
191
+ Tiny approved dry-run mode:
192
+
193
+ ```bash
194
+ RUN_TRAINING=1 COUNSEL_MODEL=Qwen/Qwen3-0.6B COUNSEL_MAX_STEPS=5 COUNSEL_DATASET_SIZE=12
195
+ ```
196
+
197
+ Before running remote training, approve the spend. Estimated costs:
198
+
199
+ - A10G dry run: about `$0.50`
200
+ - Full A100 GRPO run: about `$6-$10`
201
+
202
+ ## Run Server Locally
203
+
204
+ ```bash
205
+ uvicorn counsel_env.server.app:app --host 0.0.0.0 --port 8000
206
+ ```
207
+
208
+ Client example:
209
+
210
+ ```python
211
+ from counsel_env import CounselAction, CounselEnv
212
+
213
+ with CounselEnv(base_url="http://localhost:8000").sync() as client:
214
+ result = client.reset(curriculum_stage="easy")
215
+ print(result.observation.case_brief)
216
+
217
+ result = client.step(CounselAction(tool="ask_question", text="Where were you that night?"))
218
+ print(result.observation.witness_response)
219
+ ```
220
+
221
+ ## Limitations
222
+
223
+ - The witness is rule-based, intentionally, so reward is verifiable and cheap.
224
+ - Cases are template-generated rather than open-domain.
225
+ - The environment models adversarial questioning mechanics, not full legal procedure.
226
+
227
+ Future work: self-play witness training, civil deposition templates, jurisdiction-specific admissibility rules, and trained-vs-baseline model ablations after the first GRPO run.
228
+
229
+ ## Status
230
+
231
+ Local environment validation is complete. Hugging Face Space deployment, real GRPO training, trained-model reward curves, and public video/blog publishing still require approval and compute.
__init__.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Counsel Env Environment."""
8
+
9
+ from .client import CounselEnv
10
+ from .models import CounselAction, CounselObservation, CounselState
11
+
12
+ __all__ = [
13
+ "CounselAction",
14
+ "CounselObservation",
15
+ "CounselState",
16
+ "CounselEnv",
17
+ ]
client.py ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Counsel Env Environment Client."""
8
+
9
+ from typing import Dict
10
+
11
+ from openenv.core import EnvClient
12
+ from openenv.core.client_types import StepResult
13
+ from openenv.core.env_server.types import State
14
+
15
+ from .models import CounselAction, CounselObservation, CounselState
16
+
17
+
18
+ class CounselEnv(
19
+ EnvClient[CounselAction, CounselObservation, CounselState]
20
+ ):
21
+ """
22
+ Client for the Counsel Env Environment.
23
+
24
+ This client maintains a persistent WebSocket connection to the environment server,
25
+ enabling efficient multi-step interactions with lower latency.
26
+ Each client instance has its own dedicated environment session on the server.
27
+
28
+ Example:
29
+ >>> # Connect to a running server
30
+ >>> with CounselEnv(base_url="http://localhost:8000") as client:
31
+ ... result = client.reset()
32
+ ... print(result.observation.echoed_message)
33
+ ...
34
+ ... result = client.step(CounselAction(message="Hello!"))
35
+ ... print(result.observation.echoed_message)
36
+
37
+ Example with Docker:
38
+ >>> # Automatically start container and connect
39
+ >>> client = CounselEnv.from_docker_image("counsel_env-env:latest")
40
+ >>> try:
41
+ ... result = client.reset()
42
+ ... result = client.step(CounselAction(message="Test"))
43
+ ... finally:
44
+ ... client.close()
45
+ """
46
+
47
+ def _step_payload(self, action: CounselAction) -> Dict:
48
+ """
49
+ Convert CounselAction to JSON payload for step message.
50
+
51
+ Args:
52
+ action: CounselAction instance
53
+
54
+ Returns:
55
+ Dictionary representation suitable for JSON encoding
56
+ """
57
+ return {
58
+ "tool": action.tool,
59
+ "text": action.text,
60
+ "exhibit_id": action.exhibit_id,
61
+ "reason": action.reason,
62
+ }
63
+
64
+ def _parse_result(self, payload: Dict) -> StepResult[CounselObservation]:
65
+ """
66
+ Parse server response into StepResult[CounselObservation].
67
+
68
+ Args:
69
+ payload: JSON response data from server
70
+
71
+ Returns:
72
+ StepResult with CounselObservation
73
+ """
74
+ obs_data = payload.get("observation", {})
75
+ observation = CounselObservation(
76
+ witness_response=obs_data.get("witness_response", ""),
77
+ available_evidence=obs_data.get("available_evidence", []),
78
+ evidence_descriptions=obs_data.get("evidence_descriptions", {}),
79
+ questions_remaining=obs_data.get("questions_remaining", 0),
80
+ transcript_tail=obs_data.get("transcript_tail", ""),
81
+ case_brief=obs_data.get("case_brief", ""),
82
+ case_id=obs_data.get("case_id", ""),
83
+ difficulty=obs_data.get("difficulty", ""),
84
+ done=obs_data.get("done", payload.get("done", False)),
85
+ reward=obs_data.get("reward", payload.get("reward", 0.0)),
86
+ reward_components=obs_data.get("reward_components", {}),
87
+ )
88
+
89
+ return StepResult(
90
+ observation=observation,
91
+ reward=observation.reward,
92
+ done=observation.done,
93
+ )
94
+
95
+ def _parse_state(self, payload: Dict) -> CounselState:
96
+ """
97
+ Parse server response into State object.
98
+
99
+ Args:
100
+ payload: JSON response from state request
101
+
102
+ Returns:
103
+ State object with episode_id and step_count
104
+ """
105
+ data = payload.get("data", payload)
106
+ return CounselState(
107
+ case_id=data.get("case_id", ""),
108
+ difficulty=data.get("difficulty", ""),
109
+ contradictions_total=data.get("contradictions_total", 0),
110
+ contradictions_triggered=data.get("contradictions_triggered", 0),
111
+ contradictions_surfaced=data.get("contradictions_surfaced", 0),
112
+ questions_used=data.get("questions_used", 0),
113
+ action_count=data.get("action_count", data.get("step_count", 0)),
114
+ duplicate_question_count=data.get("duplicate_question_count", 0),
115
+ irrelevant_question_count=data.get("irrelevant_question_count", 0),
116
+ inadmissible_count=data.get("inadmissible_count", 0),
117
+ evidence_timing_successes=data.get("evidence_timing_successes", 0),
118
+ step_count=data.get("step_count", data.get("action_count", 0)),
119
+ )
evaluation.py ADDED
@@ -0,0 +1,337 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Local held-out evaluation and artifact generation for Counsel-Env."""
2
+
3
+ import argparse
4
+ import csv
5
+ import json
6
+ import random
7
+ from pathlib import Path
8
+ from statistics import mean
9
+ from typing import Callable, Dict, Iterable, List
10
+
11
+ from counsel_env.models import CounselAction
12
+ from counsel_env.server.counsel_env_environment import CounselEnvironment
13
+
14
+
15
+ AgentPolicy = Callable[[CounselEnvironment, random.Random], None]
16
+
17
+
18
+ def random_agent(env: CounselEnvironment, rng: random.Random) -> None:
19
+ """Low-information baseline."""
20
+ for _ in range(8):
21
+ if env.done:
22
+ break
23
+ if rng.random() < 0.7:
24
+ env.step(CounselAction(tool="ask_question", text=f"Can you explain detail {rng.randint(1, 999)}?"))
25
+ else:
26
+ env.step(CounselAction(tool="present_evidence", exhibit_id=rng.choice(list(env.case["evidence"].keys()))))
27
+ if not env.done:
28
+ env.step(CounselAction(tool="rest_case"))
29
+
30
+
31
+ def keyword_spam_agent(env: CounselEnvironment, rng: random.Random) -> None:
32
+ """Reward-hacking probe: asks trigger-like terms but does not reason about exhibits."""
33
+ questions = [
34
+ "Where were you?",
35
+ "What was your motive?",
36
+ "Did you know the victim?",
37
+ "What happened and why?",
38
+ "Were you at the location?",
39
+ ]
40
+ for question in questions:
41
+ if env.done:
42
+ break
43
+ env.step(CounselAction(tool="ask_question", text=question))
44
+ if not env.done:
45
+ env.step(CounselAction(tool="rest_case"))
46
+
47
+
48
+ def present_all_agent(env: CounselEnvironment, rng: random.Random) -> None:
49
+ """Reward-hacking probe: blindly presents every exhibit."""
50
+ for exhibit_id in list(env.case["evidence"].keys()):
51
+ if env.done:
52
+ break
53
+ env.step(CounselAction(tool="present_evidence", exhibit_id=exhibit_id))
54
+ if not env.done:
55
+ env.step(CounselAction(tool="rest_case"))
56
+
57
+
58
+ def oracle_scripted_agent(env: CounselEnvironment, rng: random.Random) -> None:
59
+ """Upper-bound scripted strategy using hidden contradiction metadata."""
60
+ for contradiction in env.witness.contradictions:
61
+ if env.done:
62
+ break
63
+ env.step(CounselAction(tool="ask_question", text=f"{contradiction.trigger_keywords[0]}?"))
64
+ if env.done:
65
+ break
66
+ env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
67
+ if not env.done:
68
+ env.step(CounselAction(tool="rest_case"))
69
+
70
+
71
+ AGENTS: Dict[str, AgentPolicy] = {
72
+ "random": random_agent,
73
+ "keyword_spam": keyword_spam_agent,
74
+ "present_all": present_all_agent,
75
+ "scripted_oracle": oracle_scripted_agent,
76
+ }
77
+
78
+
79
+ def make_eval_seeds(count: int = 30, start: int = 20260425) -> List[int]:
80
+ return list(range(start, start + count))
81
+
82
+
83
+ def evaluate_agent(
84
+ name: str,
85
+ policy: AgentPolicy,
86
+ seeds: Iterable[int],
87
+ curriculum_stage: str = "mixed",
88
+ transcript_limit: int = 3,
89
+ ) -> tuple[List[dict], List[str]]:
90
+ rows: List[dict] = []
91
+ markdown_samples: List[str] = []
92
+
93
+ for index, seed in enumerate(seeds):
94
+ agent_offset = sum(ord(ch) for ch in name)
95
+ rng = random.Random(seed + agent_offset)
96
+ env = CounselEnvironment()
97
+ obs = env.reset(seed=seed, curriculum_stage=curriculum_stage, episode_id=f"{name}_{seed}")
98
+ assert obs.case_id == env.case["case_id"]
99
+ policy(env, rng)
100
+ if not env.done:
101
+ env.step(CounselAction(tool="rest_case"))
102
+
103
+ components = env._calculate_reward_components()
104
+ row = {
105
+ "agent": name,
106
+ "seed": seed,
107
+ "case_id": env.case["case_id"],
108
+ "difficulty": env.case["difficulty"],
109
+ "reward": components["total_reward"],
110
+ "primary_reward": components["primary_reward"],
111
+ "auxiliary_reward": components["auxiliary_reward_raw"],
112
+ "contradictions_total": int(components["contradictions_total"]),
113
+ "contradictions_triggered": int(components["contradictions_triggered"]),
114
+ "contradictions_surfaced": int(components["contradictions_surfaced"]),
115
+ "questions_used": env.questions_used,
116
+ "evidence_presented": env.evidence_presented_count,
117
+ "evidence_timing_successes": int(components["evidence_timing_successes"]),
118
+ "blind_evidence_count": int(components["blind_evidence_count"]),
119
+ "useless_questions_ratio": components["useless_questions_ratio"],
120
+ "avg_question_length": components["avg_question_length"],
121
+ }
122
+ rows.append(row)
123
+
124
+ if index < transcript_limit:
125
+ markdown_samples.append(f"# Agent: {name}\n\n" + env.export_transcript_markdown())
126
+
127
+ return rows, markdown_samples
128
+
129
+
130
+ def summarize(rows: List[dict]) -> List[dict]:
131
+ summaries: List[dict] = []
132
+ for agent in sorted({row["agent"] for row in rows}):
133
+ agent_rows = [row for row in rows if row["agent"] == agent]
134
+ summaries.append(
135
+ {
136
+ "agent": agent,
137
+ "episodes": len(agent_rows),
138
+ "avg_reward": mean(row["reward"] for row in agent_rows),
139
+ "avg_primary_reward": mean(row["primary_reward"] for row in agent_rows),
140
+ "avg_trigger_rate": mean(
141
+ row["contradictions_triggered"] / max(1, row["contradictions_total"])
142
+ for row in agent_rows
143
+ ),
144
+ "avg_surface_rate": mean(
145
+ row["contradictions_surfaced"] / max(1, row["contradictions_total"])
146
+ for row in agent_rows
147
+ ),
148
+ "avg_evidence_timing": mean(row["evidence_timing_successes"] for row in agent_rows),
149
+ "avg_useless_ratio": mean(row["useless_questions_ratio"] for row in agent_rows),
150
+ }
151
+ )
152
+ return summaries
153
+
154
+
155
+ def write_jsonl(path: Path, rows: Iterable[dict]) -> None:
156
+ with path.open("w", encoding="utf-8") as handle:
157
+ for row in rows:
158
+ handle.write(json.dumps(row, sort_keys=True) + "\n")
159
+
160
+
161
+ def write_csv(path: Path, rows: List[dict]) -> None:
162
+ if not rows:
163
+ return
164
+ with path.open("w", newline="", encoding="utf-8") as handle:
165
+ writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
166
+ writer.writeheader()
167
+ writer.writerows(rows)
168
+
169
+
170
+ def write_plots(plot_dir: Path, summaries: List[dict]) -> None:
171
+ try:
172
+ import matplotlib.pyplot as plt
173
+ except Exception:
174
+ write_csv(plot_dir / "summary_for_plots.csv", summaries)
175
+ write_svg_bar_chart(
176
+ plot_dir / "baseline_vs_oracle.svg",
177
+ summaries,
178
+ metric="avg_reward",
179
+ title="Held-out evaluation reward by baseline",
180
+ )
181
+ write_svg_multi_metric(plot_dir / "rubric_breakdown.svg", summaries)
182
+ return
183
+
184
+ agents = [row["agent"] for row in summaries]
185
+ rewards = [row["avg_reward"] for row in summaries]
186
+ primary = [row["avg_primary_reward"] for row in summaries]
187
+ surface = [row["avg_surface_rate"] for row in summaries]
188
+ trigger = [row["avg_trigger_rate"] for row in summaries]
189
+ useless = [row["avg_useless_ratio"] for row in summaries]
190
+
191
+ plt.figure(figsize=(8, 4.5))
192
+ plt.bar(agents, rewards, color=["#777777", "#ba5a31", "#4c78a8", "#2f855a"][: len(agents)])
193
+ plt.ylabel("average reward")
194
+ plt.xlabel("agent")
195
+ plt.title("Held-out evaluation reward by baseline")
196
+ plt.xticks(rotation=20, ha="right")
197
+ plt.tight_layout()
198
+ plt.savefig(plot_dir / "baseline_vs_oracle.png", dpi=180)
199
+ plt.close()
200
+
201
+ x = range(len(agents))
202
+ plt.figure(figsize=(8, 4.5))
203
+ plt.plot(x, primary, marker="o", label="primary reward")
204
+ plt.plot(x, trigger, marker="o", label="trigger rate")
205
+ plt.plot(x, surface, marker="o", label="surface rate")
206
+ plt.plot(x, useless, marker="o", label="useless question ratio")
207
+ plt.xticks(list(x), agents, rotation=20, ha="right")
208
+ plt.ylabel("rate")
209
+ plt.xlabel("agent")
210
+ plt.title("Reward-hacking audit metrics")
211
+ plt.legend()
212
+ plt.tight_layout()
213
+ plt.savefig(plot_dir / "rubric_breakdown.png", dpi=180)
214
+ plt.close()
215
+
216
+
217
+ def write_svg_bar_chart(path: Path, summaries: List[dict], metric: str, title: str) -> None:
218
+ width, height = 840, 460
219
+ margin_left, margin_bottom, margin_top = 80, 90, 60
220
+ chart_w = width - margin_left - 40
221
+ chart_h = height - margin_top - margin_bottom
222
+ max_value = max(1.0, max(row[metric] for row in summaries))
223
+ bar_gap = 24
224
+ bar_w = (chart_w - bar_gap * (len(summaries) - 1)) / max(1, len(summaries))
225
+ colors = ["#777777", "#ba5a31", "#4c78a8", "#2f855a"]
226
+ parts = [
227
+ f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
228
+ '<rect width="100%" height="100%" fill="white"/>',
229
+ f'<text x="{width/2}" y="32" text-anchor="middle" font-family="Arial" font-size="20">{title}</text>',
230
+ f'<line x1="{margin_left}" y1="{height-margin_bottom}" x2="{width-30}" y2="{height-margin_bottom}" stroke="#333"/>',
231
+ f'<line x1="{margin_left}" y1="{margin_top}" x2="{margin_left}" y2="{height-margin_bottom}" stroke="#333"/>',
232
+ ]
233
+ for idx, row in enumerate(summaries):
234
+ value = row[metric]
235
+ x = margin_left + idx * (bar_w + bar_gap)
236
+ bar_h = chart_h * value / max_value
237
+ y = height - margin_bottom - bar_h
238
+ parts.append(f'<rect x="{x:.1f}" y="{y:.1f}" width="{bar_w:.1f}" height="{bar_h:.1f}" fill="{colors[idx % len(colors)]}"/>')
239
+ parts.append(f'<text x="{x + bar_w/2:.1f}" y="{y - 8:.1f}" text-anchor="middle" font-family="Arial" font-size="13">{value:.3f}</text>')
240
+ parts.append(f'<text x="{x + bar_w/2:.1f}" y="{height - 54}" text-anchor="middle" font-family="Arial" font-size="12">{row["agent"]}</text>')
241
+ parts.append(f'<text x="22" y="{margin_top + chart_h/2}" transform="rotate(-90 22,{margin_top + chart_h/2})" font-family="Arial" font-size="13">average reward</text>')
242
+ parts.append("</svg>")
243
+ path.write_text("\n".join(parts), encoding="utf-8")
244
+
245
+
246
+ def write_svg_multi_metric(path: Path, summaries: List[dict]) -> None:
247
+ width, height = 880, 500
248
+ margin_left, margin_bottom, margin_top = 80, 100, 60
249
+ chart_w = width - margin_left - 50
250
+ chart_h = height - margin_top - margin_bottom
251
+ metrics = [
252
+ ("avg_primary_reward", "primary", "#2f855a"),
253
+ ("avg_trigger_rate", "trigger", "#4c78a8"),
254
+ ("avg_surface_rate", "surface", "#805ad5"),
255
+ ("avg_useless_ratio", "useless", "#ba5a31"),
256
+ ]
257
+ agents = [row["agent"] for row in summaries]
258
+ x_step = chart_w / max(1, len(agents) - 1)
259
+ parts = [
260
+ f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" viewBox="0 0 {width} {height}">',
261
+ '<rect width="100%" height="100%" fill="white"/>',
262
+ f'<text x="{width/2}" y="32" text-anchor="middle" font-family="Arial" font-size="20">Reward-hacking audit metrics</text>',
263
+ f'<line x1="{margin_left}" y1="{height-margin_bottom}" x2="{width-30}" y2="{height-margin_bottom}" stroke="#333"/>',
264
+ f'<line x1="{margin_left}" y1="{margin_top}" x2="{margin_left}" y2="{height-margin_bottom}" stroke="#333"/>',
265
+ ]
266
+ for value in [0.0, 0.25, 0.5, 0.75, 1.0]:
267
+ y = height - margin_bottom - chart_h * value
268
+ parts.append(f'<line x1="{margin_left}" y1="{y:.1f}" x2="{width-30}" y2="{y:.1f}" stroke="#ddd"/>')
269
+ parts.append(f'<text x="{margin_left-10}" y="{y+4:.1f}" text-anchor="end" font-family="Arial" font-size="11">{value:.2f}</text>')
270
+ for metric, label, color in metrics:
271
+ points = []
272
+ for idx, row in enumerate(summaries):
273
+ x = margin_left + idx * x_step
274
+ y = height - margin_bottom - chart_h * max(0.0, min(1.0, row[metric]))
275
+ points.append((x, y))
276
+ path_data = " ".join(f"{x:.1f},{y:.1f}" for x, y in points)
277
+ parts.append(f'<polyline points="{path_data}" fill="none" stroke="{color}" stroke-width="3"/>')
278
+ for x, y in points:
279
+ parts.append(f'<circle cx="{x:.1f}" cy="{y:.1f}" r="4" fill="{color}"/>')
280
+ lx = width - 180
281
+ ly = margin_top + 22 * metrics.index((metric, label, color))
282
+ parts.append(f'<rect x="{lx}" y="{ly-10}" width="12" height="12" fill="{color}"/>')
283
+ parts.append(f'<text x="{lx+18}" y="{ly}" font-family="Arial" font-size="13">{label}</text>')
284
+ for idx, agent in enumerate(agents):
285
+ x = margin_left + idx * x_step
286
+ parts.append(f'<text x="{x:.1f}" y="{height-62}" text-anchor="middle" font-family="Arial" font-size="12">{agent}</text>')
287
+ parts.append("</svg>")
288
+ path.write_text("\n".join(parts), encoding="utf-8")
289
+
290
+
291
+ def write_before_after_pairs(path: Path, transcript_by_agent: Dict[str, List[str]]) -> None:
292
+ sections = ["# Before / After Transcript Samples", ""]
293
+ for agent in ["random", "keyword_spam", "present_all", "scripted_oracle"]:
294
+ samples = transcript_by_agent.get(agent, [])
295
+ if samples:
296
+ sections.append(samples[0])
297
+ sections.append("")
298
+ path.write_text("\n".join(sections), encoding="utf-8")
299
+
300
+
301
+ def run_evaluation(output_dir: str | Path = "assets", episodes: int = 30) -> dict:
302
+ output = Path(output_dir)
303
+ plot_dir = output / "plots"
304
+ transcript_dir = output / "transcripts"
305
+ output.mkdir(exist_ok=True)
306
+ plot_dir.mkdir(parents=True, exist_ok=True)
307
+ transcript_dir.mkdir(parents=True, exist_ok=True)
308
+
309
+ seeds = make_eval_seeds(episodes)
310
+ all_rows: List[dict] = []
311
+ transcript_by_agent: Dict[str, List[str]] = {}
312
+ for agent, policy in AGENTS.items():
313
+ rows, markdown_samples = evaluate_agent(agent, policy, seeds)
314
+ all_rows.extend(rows)
315
+ transcript_by_agent[agent] = markdown_samples
316
+
317
+ summaries = summarize(all_rows)
318
+ write_jsonl(output / "heldout_eval.jsonl", all_rows)
319
+ write_csv(output / "heldout_eval_summary.csv", summaries)
320
+ (output / "heldout_eval_summary.json").write_text(json.dumps(summaries, indent=2), encoding="utf-8")
321
+ write_plots(plot_dir, summaries)
322
+ write_before_after_pairs(transcript_dir / "before_after_pairs.md", transcript_by_agent)
323
+ return {"rows": all_rows, "summaries": summaries}
324
+
325
+
326
+ def main() -> int:
327
+ parser = argparse.ArgumentParser()
328
+ parser.add_argument("--output-dir", default="assets")
329
+ parser.add_argument("--episodes", type=int, default=30)
330
+ args = parser.parse_args()
331
+ result = run_evaluation(args.output_dir, args.episodes)
332
+ print(json.dumps(result["summaries"], indent=2))
333
+ return 0
334
+
335
+
336
+ if __name__ == "__main__":
337
+ raise SystemExit(main())
models.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Optional
2
+
3
+ from pydantic import Field
4
+
5
+ from openenv.core import Action, Observation, State
6
+
7
+
8
+ class CounselAction(Action):
9
+ tool: str = Field(description="Action tool: ask_question, present_evidence, make_objection, or rest_case")
10
+ text: Optional[str] = Field(default=None, description="Question text for ask_question")
11
+ exhibit_id: Optional[str] = Field(default=None, description="Evidence id for present_evidence")
12
+ reason: Optional[str] = Field(default=None, description="Objection reason for make_objection")
13
+
14
+
15
+ class CounselObservation(Observation):
16
+ witness_response: str = ""
17
+ available_evidence: List[str] = Field(default_factory=list)
18
+ evidence_descriptions: Dict[str, str] = Field(default_factory=dict)
19
+ questions_remaining: int = 0
20
+ transcript_tail: str = ""
21
+ case_brief: str = ""
22
+ case_id: str = ""
23
+ difficulty: str = ""
24
+ reward_components: Dict[str, float] = Field(default_factory=dict)
25
+
26
+
27
+ class CounselState(State):
28
+ case_id: str = ""
29
+ difficulty: str = ""
30
+ contradictions_total: int = 0
31
+ contradictions_triggered: int = 0
32
+ contradictions_surfaced: int = 0
33
+ questions_used: int = 0
34
+ action_count: int = 0
35
+ duplicate_question_count: int = 0
36
+ irrelevant_question_count: int = 0
37
+ inadmissible_count: int = 0
38
+ evidence_timing_successes: int = 0
notebooks/train_counsel.ipynb ADDED
@@ -0,0 +1,327 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "metadata": {},
6
+ "source": [
7
+ "# Cross-Examination Arena GRPO Training\n",
8
+ "\n",
9
+ "Credit-safe training notebook for Counsel-Env. Hugging Face GPU usage is disabled by default. Before running a remote training job, confirm the spend: a short A10G dry run is roughly $0.50, and a 90 minute A100 GRPO run is roughly $6-$10 depending on queue/runtime.\n",
10
+ "\n",
11
+ "This notebook uses OpenEnv's `environment_factory` pattern, curriculum stages, shaped rewards from the environment, and rollout diagnostics saved to `rollout_debug.jsonl`."
12
+ ]
13
+ },
14
+ {
15
+ "cell_type": "code",
16
+ "execution_count": null,
17
+ "metadata": {},
18
+ "outputs": [],
19
+ "source": [
20
+ "# Optional installs for Colab/HF Jobs. Keep commented until you approve credit/network usage.\n",
21
+ "# !pip install -U 'trl[vllm]' unsloth openenv-core datasets transformers accelerate matplotlib\n",
22
+ "# !pip install -e .\n",
23
+ "\n",
24
+ "import json\n",
25
+ "import os\n",
26
+ "import random\n",
27
+ "import re\n",
28
+ "from pathlib import Path\n",
29
+ "from typing import Any, Dict, List\n",
30
+ "\n",
31
+ "import numpy as np\n",
32
+ "from datasets import Dataset\n",
33
+ "from trl import GRPOConfig, GRPOTrainer\n",
34
+ "\n",
35
+ "from counsel_env import CounselEnv\n",
36
+ "from counsel_env.models import CounselAction\n",
37
+ "\n",
38
+ "MODEL = os.getenv('COUNSEL_MODEL', 'Qwen/Qwen3-1.7B')\n",
39
+ "ENV_URL = os.getenv('COUNSEL_ENV_URL', 'http://localhost:8000')\n",
40
+ "OUTPUT_DIR = os.getenv('COUNSEL_OUTPUT_DIR', './counsel-grpo-output')\n",
41
+ "RUN_TRAINING = os.getenv('RUN_TRAINING', '0') == '1'\n",
42
+ "MAX_STEPS = int(os.getenv('COUNSEL_MAX_STEPS', '200'))\n",
43
+ "DATASET_SIZE = int(os.getenv('COUNSEL_DATASET_SIZE', '256'))\n",
44
+ "\n",
45
+ "print({'model': MODEL, 'env_url': ENV_URL, 'run_training': RUN_TRAINING, 'max_steps': MAX_STEPS, 'dataset_size': DATASET_SIZE})"
46
+ ]
47
+ },
48
+ {
49
+ "cell_type": "code",
50
+ "execution_count": null,
51
+ "metadata": {},
52
+ "outputs": [],
53
+ "source": [
54
+ "CURRICULUM_DISTRIBUTION = {\n",
55
+ " 'easy': {'easy': 1.0},\n",
56
+ " 'medium': {'easy': 0.25, 'medium': 0.75},\n",
57
+ " 'hard': {'easy': 0.10, 'medium': 0.30, 'hard': 0.60},\n",
58
+ "}\n",
59
+ "\n",
60
+ "def sample_stage_schedule(total: int = 256) -> List[str]:\n",
61
+ " # Warm start on easy, then medium, then hard/mixed challenge cases.\n",
62
+ " schedule = []\n",
63
+ " for i in range(total):\n",
64
+ " frac = i / max(1, total - 1)\n",
65
+ " if frac < 0.25:\n",
66
+ " schedule.append('easy')\n",
67
+ " elif frac < 0.70:\n",
68
+ " schedule.append('medium')\n",
69
+ " else:\n",
70
+ " schedule.append('hard')\n",
71
+ " return schedule\n",
72
+ "\n",
73
+ "def create_dataset(num_samples: int = 256) -> Dataset:\n",
74
+ " base_prompt = (\n",
75
+ " 'You are a sharp prosecuting attorney cross-examining a deterministic witness. '\n",
76
+ " 'Your goal is to surface contradictions by first making the witness commit to a claim, '\n",
77
+ " 'then presenting the exact exhibit that disproves it. Use the limited question budget efficiently. '\n",
78
+ " 'Avoid repeated, irrelevant, leading, or compound questions.'\n",
79
+ " )\n",
80
+ " stages = sample_stage_schedule(num_samples)\n",
81
+ " prompts = [[{'role': 'user', 'content': base_prompt}] for _ in stages]\n",
82
+ " return Dataset.from_dict({'prompt': prompts, 'curriculum_stage': stages})\n",
83
+ "\n",
84
+ "dataset = create_dataset(DATASET_SIZE)\n",
85
+ "dataset[:3]"
86
+ ]
87
+ },
88
+ {
89
+ "cell_type": "code",
90
+ "execution_count": null,
91
+ "metadata": {},
92
+ "outputs": [],
93
+ "source": [
94
+ "def question_pattern(question: str) -> str:\n",
95
+ " text = re.sub(r'[^a-z0-9\\s]', ' ', (question or '').lower())\n",
96
+ " return ' '.join(text.split()[:12])\n",
97
+ "\n",
98
+ "class CounselToolEnv:\n",
99
+ " \"\"\"TRL environment wrapper. Constructor intentionally takes no args.\"\"\"\n",
100
+ "\n",
101
+ " def __init__(self):\n",
102
+ " self.client = CounselEnv(base_url=ENV_URL)\n",
103
+ " self.reward = 0.0\n",
104
+ " self.done = False\n",
105
+ " self.stage = 'medium'\n",
106
+ " self.questions: List[str] = []\n",
107
+ " self.actions: List[str] = []\n",
108
+ " self.last_components: Dict[str, Any] = {}\n",
109
+ "\n",
110
+ " def reset(self, **kwargs) -> str:\n",
111
+ " \"\"\"Reset the remote Counsel-Env episode.\"\"\"\n",
112
+ " self.stage = kwargs.get('curriculum_stage') or 'medium'\n",
113
+ " result = self.client.reset(curriculum_stage=self.stage)\n",
114
+ " obs = result.observation\n",
115
+ " self.reward = 0.0\n",
116
+ " self.done = False\n",
117
+ " self.questions = []\n",
118
+ " self.actions = []\n",
119
+ " self.last_components = obs.reward_components\n",
120
+ " return (\n",
121
+ " f'CASE BRIEF:\\n{obs.case_brief}\\n\\n'\n",
122
+ " f'You have {obs.questions_remaining} questions. '\n",
123
+ " f'Available exhibits: {obs.available_evidence}. '\n",
124
+ " 'Use ask_question, present_evidence, or rest_case.'\n",
125
+ " )\n",
126
+ "\n",
127
+ " def ask_question(self, question: str) -> str:\n",
128
+ " \"\"\"Ask the witness a question.\n",
129
+ "\n",
130
+ " Args:\n",
131
+ " question: The question to ask the witness.\n",
132
+ "\n",
133
+ " Returns:\n",
134
+ " The witness response and remaining budget.\n",
135
+ " \"\"\"\n",
136
+ " if self.done:\n",
137
+ " raise ValueError('Episode is over.')\n",
138
+ " self.actions.append('ask_question')\n",
139
+ " self.questions.append(question)\n",
140
+ " result = self.client.step(CounselAction(tool='ask_question', text=question))\n",
141
+ " obs = result.observation\n",
142
+ " self.reward = obs.reward\n",
143
+ " self.done = obs.done\n",
144
+ " self.last_components = obs.reward_components\n",
145
+ " return f'WITNESS: {obs.witness_response}\\n[{obs.questions_remaining} questions left]'\n",
146
+ "\n",
147
+ " def present_evidence(self, exhibit_id: str) -> str:\n",
148
+ " \"\"\"Present an exhibit to the witness.\n",
149
+ "\n",
150
+ " Args:\n",
151
+ " exhibit_id: The id of the exhibit to present.\n",
152
+ "\n",
153
+ " Returns:\n",
154
+ " The witness reaction.\n",
155
+ " \"\"\"\n",
156
+ " if self.done:\n",
157
+ " raise ValueError('Episode is over.')\n",
158
+ " self.actions.append('present_evidence')\n",
159
+ " result = self.client.step(CounselAction(tool='present_evidence', exhibit_id=exhibit_id))\n",
160
+ " obs = result.observation\n",
161
+ " self.reward = obs.reward\n",
162
+ " self.done = obs.done\n",
163
+ " self.last_components = obs.reward_components\n",
164
+ " return obs.witness_response\n",
165
+ "\n",
166
+ " def rest_case(self) -> str:\n",
167
+ " \"\"\"End the cross-examination.\n",
168
+ "\n",
169
+ " Returns:\n",
170
+ " Final episode status.\n",
171
+ " \"\"\"\n",
172
+ " self.actions.append('rest_case')\n",
173
+ " result = self.client.step(CounselAction(tool='rest_case'))\n",
174
+ " obs = result.observation\n",
175
+ " self.reward = obs.reward\n",
176
+ " self.done = True\n",
177
+ " self.last_components = obs.reward_components\n",
178
+ " return f'Case rested. Final reward: {self.reward:.3f}'\n",
179
+ "\n",
180
+ " @property\n",
181
+ " def action_diversity(self) -> int:\n",
182
+ " return len(set(self.actions))\n",
183
+ "\n",
184
+ " @property\n",
185
+ " def unique_question_patterns_pct(self) -> float:\n",
186
+ " patterns = {question_pattern(q) for q in self.questions}\n",
187
+ " return len(patterns) / max(1, len(self.questions))"
188
+ ]
189
+ },
190
+ {
191
+ "cell_type": "code",
192
+ "execution_count": null,
193
+ "metadata": {},
194
+ "outputs": [],
195
+ "source": [
196
+ "def reward_func(environments, **kwargs) -> List[float]:\n",
197
+ " # The environment computes WeightedSum(primary + auxiliary).\n",
198
+ " return [float(env.reward) for env in environments]\n",
199
+ "\n",
200
+ "def log_diversity(environments: List[CounselToolEnv], path: str = 'rollout_debug.jsonl') -> None:\n",
201
+ " rows = []\n",
202
+ " for idx, env in enumerate(environments):\n",
203
+ " row = {\n",
204
+ " 'episode': idx,\n",
205
+ " 'stage': env.stage,\n",
206
+ " 'reward': env.reward,\n",
207
+ " 'action_diversity': env.action_diversity,\n",
208
+ " 'unique_question_patterns_pct': env.unique_question_patterns_pct,\n",
209
+ " **env.last_components,\n",
210
+ " }\n",
211
+ " rows.append(row)\n",
212
+ " with open(path, 'a', encoding='utf-8') as handle:\n",
213
+ " for row in rows:\n",
214
+ " handle.write(json.dumps(row, sort_keys=True) + '\\n')\n",
215
+ "\n",
216
+ "# Failure detection required by the hackathon brief.\n",
217
+ "def warn_on_reward_difficulty(rewards: List[float]) -> None:\n",
218
+ " avg_reward = float(np.mean(rewards)) if rewards else 0.0\n",
219
+ " if avg_reward == 0.0:\n",
220
+ " print('WARNING: ENV TOO HARD')\n",
221
+ " if avg_reward > 0.7:\n",
222
+ " print('WARNING: ENV TOO EASY')\n",
223
+ " print(f'avg_reward={avg_reward:.3f}')"
224
+ ]
225
+ },
226
+ {
227
+ "cell_type": "code",
228
+ "execution_count": null,
229
+ "metadata": {},
230
+ "outputs": [],
231
+ "source": [
232
+ "training_args = GRPOConfig(\n",
233
+ " use_vllm=True,\n",
234
+ " vllm_mode='colocate',\n",
235
+ " chat_template_kwargs={'enable_thinking': False},\n",
236
+ " max_completion_length=1024,\n",
237
+ " num_generations=6,\n",
238
+ " temperature=0.8,\n",
239
+ " top_p=0.95,\n",
240
+ " gradient_accumulation_steps=16,\n",
241
+ " per_device_train_batch_size=1,\n",
242
+ " learning_rate=1e-5,\n",
243
+ " num_train_epochs=1,\n",
244
+ " max_steps=MAX_STEPS,\n",
245
+ " beta=0.04,\n",
246
+ " log_completions=True,\n",
247
+ " logging_steps=5,\n",
248
+ " report_to='none',\n",
249
+ " output_dir=OUTPUT_DIR,\n",
250
+ " push_to_hub=False,\n",
251
+ ")\n",
252
+ "\n",
253
+ "trainer = None\n",
254
+ "if RUN_TRAINING:\n",
255
+ " trainer = GRPOTrainer(\n",
256
+ " model=MODEL,\n",
257
+ " reward_funcs=reward_func,\n",
258
+ " train_dataset=dataset,\n",
259
+ " args=training_args,\n",
260
+ " environment_factory=CounselToolEnv,\n",
261
+ " )\n",
262
+ " train_result = trainer.train()\n",
263
+ " trainer.save_model(OUTPUT_DIR)\n",
264
+ "else:\n",
265
+ " print('Training is disabled. Set RUN_TRAINING=1 only after credit approval.')\n",
266
+ " print('Tiny dry-run example after approval: RUN_TRAINING=1 COUNSEL_MODEL=Qwen/Qwen3-0.6B COUNSEL_MAX_STEPS=5 COUNSEL_DATASET_SIZE=12')"
267
+ ]
268
+ },
269
+ {
270
+ "cell_type": "code",
271
+ "execution_count": null,
272
+ "metadata": {},
273
+ "outputs": [],
274
+ "source": [
275
+ "# Local, no-credit rollout diagnostics. This uses the in-process env, not HF.\n",
276
+ "from counsel_env.server.diagnostics import run_rollout_diagnostics\n",
277
+ "\n",
278
+ "rows = run_rollout_diagnostics(output_path='rollout_debug.jsonl', num_episodes=10, policy='mixed')\n",
279
+ "rewards = [row['total_reward'] for row in rows]\n",
280
+ "warn_on_reward_difficulty(rewards)\n",
281
+ "print(f'non_zero_rollouts={sum(r > 0 for r in rewards)}/{len(rewards)}')"
282
+ ]
283
+ },
284
+ {
285
+ "cell_type": "code",
286
+ "execution_count": null,
287
+ "metadata": {},
288
+ "outputs": [],
289
+ "source": [
290
+ "# Plot reward trend after a real run. This cell expects trainer logs or an exported metrics list.\n",
291
+ "# It intentionally does not fabricate a reward curve.\n",
292
+ "import matplotlib.pyplot as plt\n",
293
+ "\n",
294
+ "metrics_path = Path(OUTPUT_DIR) / 'trainer_state.json'\n",
295
+ "if metrics_path.exists():\n",
296
+ " state = json.loads(metrics_path.read_text())\n",
297
+ " history = state.get('log_history', [])\n",
298
+ " steps = [row['step'] for row in history if 'reward' in row]\n",
299
+ " rewards = [row['reward'] for row in history if 'reward' in row]\n",
300
+ " if rewards:\n",
301
+ " plt.figure(figsize=(7, 4))\n",
302
+ " plt.plot(steps, rewards, marker='o')\n",
303
+ " plt.xlabel('training step')\n",
304
+ " plt.ylabel('GRPO reward')\n",
305
+ " plt.title('Counsel-Env reward trend')\n",
306
+ " plt.grid(alpha=0.3)\n",
307
+ " plt.savefig('reward_curve.png', dpi=160, bbox_inches='tight')\n",
308
+ " plt.show()\n",
309
+ "else:\n",
310
+ " print('No trainer_state.json yet. Run approved training first, then rerun this cell.')"
311
+ ]
312
+ }
313
+ ],
314
+ "metadata": {
315
+ "kernelspec": {
316
+ "display_name": "Python 3",
317
+ "language": "python",
318
+ "name": "python3"
319
+ },
320
+ "language_info": {
321
+ "name": "python",
322
+ "pygments_lexer": "ipython3"
323
+ }
324
+ },
325
+ "nbformat": 4,
326
+ "nbformat_minor": 5
327
+ }
openenv.yaml ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ name: counsel-env
2
+ description: Cross-examination arena — train LLMs to surface contradictions in witness testimony.
3
+ tags:
4
+ - rl
5
+ - multi-agent
6
+ - theory-of-mind
7
+ - dialogue
8
+ - hackathon
9
+ license: bsd-3-clause
10
+
pyproject.toml ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ [build-system]
8
+ requires = ["setuptools>=45", "wheel"]
9
+ build-backend = "setuptools.build_meta"
10
+
11
+ [project]
12
+ name = "counsel-env"
13
+ version = "0.1.0"
14
+ dependencies = [
15
+ "openenv-core>=0.2.1",
16
+ "fastapi",
17
+ "uvicorn",
18
+ "pydantic",
19
+ ]
20
+
21
+ [project.optional-dependencies]
22
+ dev = [
23
+ "pytest>=8.0.0",
24
+ "pytest-cov>=4.0.0",
25
+ ]
26
+
27
+ [project.scripts]
28
+ # Server entry point - enables running via: uv run --project . server
29
+ # or: python -m counsel_env.server.app
30
+ server = "counsel_env.server.app:main"
31
+
32
+ [tool.setuptools]
33
+ include-package-data = true
34
+ packages = ["counsel_env", "counsel_env.server"]
35
+ package-dir = { "counsel_env" = ".", "counsel_env.server" = "server" }
server/__init__.py ADDED
@@ -0,0 +1,14 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """Counsel Env environment server components."""
8
+
9
+ try:
10
+ from .counsel_env_environment import CounselEnvironment
11
+ except ImportError:
12
+ from counsel_env_environment import CounselEnvironment
13
+
14
+ __all__ = ["CounselEnvironment"]
server/app.py ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the BSD-style license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+
7
+ """
8
+ FastAPI application for the Counsel Env Environment.
9
+
10
+ This module creates an HTTP server that exposes the CounselEnvironment
11
+ over HTTP and WebSocket endpoints, compatible with EnvClient.
12
+
13
+ Endpoints:
14
+ - POST /reset: Reset the environment
15
+ - POST /step: Execute an action
16
+ - GET /state: Get current environment state
17
+ - GET /schema: Get action/observation schemas
18
+ - WS /ws: WebSocket endpoint for persistent sessions
19
+
20
+ Usage:
21
+ # Development (with auto-reload):
22
+ uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
23
+
24
+ # Production:
25
+ uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
26
+
27
+ # Or run directly:
28
+ python -m server.app
29
+ """
30
+
31
+ try:
32
+ from openenv.core.env_server.http_server import create_app
33
+ except Exception as e: # pragma: no cover
34
+ raise ImportError(
35
+ "openenv is required for the web interface. Install dependencies with '\n uv sync\n'"
36
+ ) from e
37
+
38
+ try:
39
+ from ..models import CounselAction, CounselObservation
40
+ from .counsel_env_environment import CounselEnvironment
41
+ except (ImportError, ModuleNotFoundError):
42
+ from models import CounselAction, CounselObservation
43
+ from server.counsel_env_environment import CounselEnvironment
44
+
45
+
46
+ # Create the app with web interface and README integration
47
+ app = create_app(
48
+ CounselEnvironment,
49
+ CounselAction,
50
+ CounselObservation,
51
+ env_name="counsel_env",
52
+ max_concurrent_envs=64, # increase this number to allow more concurrent WebSocket sessions
53
+ )
54
+
55
+
56
+ def run_server(host: str = "0.0.0.0", port: int = 8000):
57
+ """
58
+ Entry point for direct execution via uv run or python -m.
59
+
60
+ This function enables running the server without Docker:
61
+ uv run --project . server
62
+ uv run --project . server --port 8001
63
+ python -m counsel_env.server.app
64
+
65
+ Args:
66
+ host: Host address to bind to (default: "0.0.0.0")
67
+ port: Port number to listen on (default: 8000)
68
+
69
+ For production deployments, consider using uvicorn directly with
70
+ multiple workers:
71
+ uvicorn counsel_env.server.app:app --workers 4
72
+ """
73
+ import uvicorn
74
+
75
+ uvicorn.run(app, host=host, port=port)
76
+
77
+
78
+ def main():
79
+ import argparse
80
+
81
+ parser = argparse.ArgumentParser()
82
+ parser.add_argument("--host", default="0.0.0.0")
83
+ parser.add_argument("--port", type=int, default=8000)
84
+ args = parser.parse_args()
85
+ run_server(host=args.host, port=args.port)
86
+
87
+
88
+ if __name__ == "__main__":
89
+ main()
server/case_generator.py ADDED
@@ -0,0 +1,705 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Cross-Examination Arena — Case generator.
3
+
4
+ Three fully-fleshed templates (alibi, knowledge denial, possession denial),
5
+ six slot-fill pools, and the public `generate_case()` factory.
6
+
7
+ Add 2 more templates on Day 1 if time allows. Each new template just needs
8
+ to return a dict with the same shape as `_make_case_dict` produces.
9
+ """
10
+
11
+ import random
12
+ from dataclasses import dataclass
13
+ from typing import Dict, List
14
+
15
+
16
+ # ---------------------------------------------------------------------------
17
+ # Contradiction dataclass (mirror of witness.Contradiction so the case file
18
+ # stays self-contained; in your real codebase, import from witness.py).
19
+ # ---------------------------------------------------------------------------
20
+ @dataclass
21
+ class Contradiction:
22
+ cid: str
23
+ trigger_keywords: List[str]
24
+ sealed_claim: str
25
+ disprover_evidence_id: str
26
+ triggered: bool = False
27
+ surfaced: bool = False
28
+
29
+
30
+ # ---------------------------------------------------------------------------
31
+ # Slot-fill pools. Extend these freely — each entry adds case variety.
32
+ # ---------------------------------------------------------------------------
33
+ NAMES = [
34
+ "Sara Patel", "Mark Chen", "Alex Reyes", "Priya Sharma", "Jordan Kim",
35
+ "Diana Okafor", "Tomas Ruiz", "Yelena Volkov", "Hassan Ali", "Mei Tanaka",
36
+ "Carlos Mendes", "Ingrid Lindqvist", "Dev Krishnan", "Naomi Brown",
37
+ "Samir Haque", "Lila Andersson", "Theo Morales", "Reza Farahani",
38
+ "Kavita Iyer", "Owen Nakamura", "Esme Dubois", "Ravi Joshi",
39
+ "Hana Lee", "Felix Bauer", "Aaliyah Singh", "Mateo Silva",
40
+ "Zoe Whitfield", "Daichi Sato", "Nadine Mahfouz", "Pavel Novak",
41
+ ]
42
+
43
+ PLACES = [
44
+ "the Riverside Cafe on Fifth Avenue",
45
+ "the parking garage at 22 Elm Street",
46
+ "the Marriott hotel lobby downtown",
47
+ "the office of Henderson & Co. on Pine",
48
+ "the warehouse on Industrial Drive",
49
+ "the gym on Lexington",
50
+ "Brennan's Bar near Union Square",
51
+ "the apartment at 480 Oak Lane",
52
+ "the boardwalk at Crescent Beach",
53
+ "the Greyhound bus station on Central",
54
+ "the parking lot behind the Westgate Mall",
55
+ "the rooftop bar at the Stratton Hotel",
56
+ "the public library on Walnut Street",
57
+ "the diner at the corner of 9th and Broad",
58
+ ]
59
+
60
+ TIMES = [
61
+ "8:15 PM", "9:00 PM", "9:45 PM", "10:30 PM", "11:00 PM", "11:45 PM",
62
+ "12:15 AM", "1:00 AM", "7:30 PM", "6:45 PM",
63
+ ]
64
+
65
+ DATES = [
66
+ "March 14th", "April 3rd", "January 22nd", "October 17th", "August 9th",
67
+ "December 5th", "May 28th", "June 14th",
68
+ ]
69
+
70
+ OBJECTS = [
71
+ "the .38 revolver", "the briefcase", "the Rolex watch", "the burner phone",
72
+ "the laptop", "the security badge", "the diamond pendant",
73
+ "the warehouse keys", "the master keycard",
74
+ ]
75
+
76
+ MOTIVES = [
77
+ "a $40,000 unpaid debt",
78
+ "a contested inheritance",
79
+ "a botched business partnership",
80
+ "an extramarital affair that had been discovered",
81
+ "a longstanding professional rivalry",
82
+ "the dismissal from a senior position six months earlier",
83
+ ]
84
+
85
+
86
+ # ---------------------------------------------------------------------------
87
+ # Helper to build the case dict in a consistent shape.
88
+ # ---------------------------------------------------------------------------
89
+ def _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions):
90
+ return {
91
+ "case_id": case_id,
92
+ "case_brief": brief,
93
+ "ground_truth": ground_truth,
94
+ "witness_story": story,
95
+ "evidence": evidence,
96
+ "contradictions": contradictions,
97
+ }
98
+
99
+
100
+ def _case_suffix() -> str:
101
+ """Seed-respecting six-character case suffix."""
102
+ return f"{random.randint(0, 16**6 - 1):06x}"
103
+
104
+
105
+ # ---------------------------------------------------------------------------
106
+ # Helper functions for time manipulation
107
+ # ---------------------------------------------------------------------------
108
+ def _minus_minutes(time_str, minutes):
109
+ # Simple helper for time arithmetic
110
+ hour, minute = map(int, time_str.replace(" PM", "").replace(" AM", "").split(":"))
111
+ if "PM" in time_str and hour != 12:
112
+ hour += 12
113
+ total_minutes = hour * 60 + minute - minutes
114
+ new_hour = total_minutes // 60
115
+ new_minute = total_minutes % 60
116
+ return f"{new_hour % 12 or 12}:{new_minute:02d} {'PM' if new_hour >= 12 else 'AM'}"
117
+
118
+ def _plus_minutes(time_str, minutes):
119
+ hour, minute = map(int, time_str.replace(" PM", "").replace(" AM", "").split(":"))
120
+ if "PM" in time_str and hour != 12:
121
+ hour += 12
122
+ total_minutes = hour * 60 + minute + minutes
123
+ new_hour = total_minutes // 60
124
+ new_minute = total_minutes % 60
125
+ return f"{new_hour % 12 or 12}:{new_minute:02d} {'PM' if new_hour >= 12 else 'AM'}"
126
+
127
+
128
+ # ---------------------------------------------------------------------------
129
+ # Template 1 — Alibi conflict
130
+ # ---------------------------------------------------------------------------
131
+ def generate_alibi_case():
132
+ """Witness claims they were at place A; phone/transit data places them at B."""
133
+ perp = random.choice(NAMES)
134
+ fake_place = random.choice(PLACES)
135
+ real_place = random.choice([p for p in PLACES if p != fake_place])
136
+ crime_place = random.choice([p for p in PLACES if p not in (fake_place, real_place)])
137
+ t = random.choice(TIMES)
138
+ date = random.choice(DATES)
139
+ motive = random.choice(MOTIVES)
140
+
141
+ case_id = f"alibi_{_case_suffix()}"
142
+ brief = (
143
+ f"{perp} is the prime suspect in the burglary at {crime_place} on the night of {date}. "
144
+ f"The state alleges the burglary occurred at approximately {t}. "
145
+ f"The defendant has taken the stand."
146
+ )
147
+ ground_truth = (
148
+ f"{perp} was at {real_place} at {t} on {date}, then traveled to {crime_place} "
149
+ f"and committed the burglary. Motive: {motive}."
150
+ )
151
+ story = {
152
+ "alibi": f"I was at {fake_place} that entire evening with friends.",
153
+ "where": f"I told you, {fake_place}. I never left.",
154
+ "evening": f"I had dinner and drinks at {fake_place} until close.",
155
+ "warehouse": f"I have never been to {crime_place}. Not once.",
156
+ "crime": f"I had nothing to do with what happened at {crime_place}.",
157
+ "motive": f"I had no reason to do anything to anyone.",
158
+ "victim": f"We had a professional disagreement, that's all. Nothing serious.",
159
+ "phone": f"My phone was with me the whole night.",
160
+ }
161
+ evidence = {
162
+ "phone_tower_log":
163
+ f"Cellular records: {perp}'s phone connected to the tower nearest "
164
+ f"{real_place} at {t}, then to the tower covering {crime_place} 22 minutes later.",
165
+ "ride_receipt":
166
+ f"Rideshare receipt in {perp}'s account: dropoff at {real_place} at "
167
+ f"{_minus_minutes(t, 30)}; second pickup at {real_place}, dropoff at "
168
+ f"{crime_place} at {_plus_minutes(t, 18)}.",
169
+ "victim_letter":
170
+ f"Letter from the victim's attorney to {perp}, dated three weeks before the "
171
+ f"burglary, formally demanding payment related to {motive}.",
172
+ "doormat_print":
173
+ f"A partial shoe print recovered from the doormat at {crime_place}, "
174
+ f"matching size and tread of footwear later seized from {perp}'s residence.",
175
+ }
176
+ contradictions = [
177
+ Contradiction(
178
+ cid="c1_alibi_location",
179
+ trigger_keywords=[
180
+ "where were you", "alibi", "that night", "evening", "your location",
181
+ t.lower(),
182
+ ],
183
+ sealed_claim=f"I was at {fake_place} the entire evening — I never left.",
184
+ disprover_evidence_id="phone_tower_log",
185
+ ),
186
+ Contradiction(
187
+ cid="c2_warehouse_familiarity",
188
+ trigger_keywords=[
189
+ "ever been to", "familiar with", "the warehouse", "the scene",
190
+ "crime scene", crime_place.split()[1].lower() if len(crime_place.split()) > 1 else "warehouse",
191
+ ],
192
+ sealed_claim=f"I have never set foot at {crime_place}. I wouldn't even know how to get there.",
193
+ disprover_evidence_id="ride_receipt",
194
+ ),
195
+ Contradiction(
196
+ cid="c3_motive_denial",
197
+ trigger_keywords=[
198
+ "motive", "reason", "grievance", "dispute", "the victim",
199
+ "any conflict", "disagreement",
200
+ ],
201
+ sealed_claim=f"I had no real grievance with the victim. We had a polite professional relationship.",
202
+ disprover_evidence_id="victim_letter",
203
+ ),
204
+ ]
205
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
206
+
207
+
208
+ # ---------------------------------------------------------------------------
209
+ # Template 2 — Knowledge denial
210
+ # ---------------------------------------------------------------------------
211
+ def generate_knowledge_denial_case():
212
+ """Witness denies knowing victim/place/object; documentary evidence shows otherwise."""
213
+ perp = random.choice(NAMES)
214
+ victim = random.choice([n for n in NAMES if n != perp])
215
+ place = random.choice(PLACES)
216
+ date = random.choice(DATES)
217
+
218
+ case_id = f"knowledge_{_case_suffix()}"
219
+ brief = (
220
+ f"{perp} is charged with embezzlement from {place}. The state alleges {perp} "
221
+ f"systematically diverted funds in coordination with the manager, {victim}, "
222
+ f"between {date} and the present. The defendant denies knowing {victim}."
223
+ )
224
+ ground_truth = (
225
+ f"{perp} and {victim} have been business associates for over two years. "
226
+ f"They met at {place} and collaborated on the embezzlement scheme."
227
+ )
228
+ story = {
229
+ "know": f"I don't know {victim} at all. Never met them.",
230
+ "victim": f"Who is {victim}? I've never heard of them.",
231
+ "place": f"I've never been to {place}.",
232
+ "embezzlement": f"I had nothing to do with any embezzlement.",
233
+ "work": f"I work in a completely different field.",
234
+ "meeting": f"I've never met with anyone about this.",
235
+ }
236
+ evidence = {
237
+ "email_chain":
238
+ f"Email chain between {perp} and {victim} discussing the embezzlement plan, "
239
+ f"dated {date}, sent from {perp}'s work email.",
240
+ "meeting_notes":
241
+ f"Meeting notes from {place} conference room, dated two weeks ago, "
242
+ f"with {perp} and {victim}'s signatures.",
243
+ "bank_records":
244
+ f"Bank records showing transfers from {place}'s accounts to {perp}'s "
245
+ f"personal account, totaling $150,000 over six months.",
246
+ "visitor_log":
247
+ f"Security log from {place} showing {perp}'s entry and exit times "
248
+ f"on multiple dates, including meetings with {victim}.",
249
+ }
250
+ contradictions = [
251
+ Contradiction(
252
+ cid="c1_know_victim",
253
+ trigger_keywords=[
254
+ "know", "met", "familiar", victim.split()[0].lower(),
255
+ "relationship", "acquaintance",
256
+ ],
257
+ sealed_claim=f"I have never met {victim} in my life.",
258
+ disprover_evidence_id="email_chain",
259
+ ),
260
+ Contradiction(
261
+ cid="c2_place_visit",
262
+ trigger_keywords=[
263
+ "been to", "visited", place.split()[1].lower() if len(place.split()) > 1 else "place",
264
+ "location", "office",
265
+ ],
266
+ sealed_claim=f"I've never set foot in {place}.",
267
+ disprover_evidence_id="visitor_log",
268
+ ),
269
+ Contradiction(
270
+ cid="c3_collaboration",
271
+ trigger_keywords=[
272
+ "worked with", "collaborated", "together", "scheme",
273
+ "embezzlement", "plan",
274
+ ],
275
+ sealed_claim=f"I worked alone. No one else was involved.",
276
+ disprover_evidence_id="meeting_notes",
277
+ ),
278
+ ]
279
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
280
+
281
+
282
+ # ---------------------------------------------------------------------------
283
+ # Template 3 — Possession denial
284
+ # ---------------------------------------------------------------------------
285
+ def generate_possession_denial_case():
286
+ """Witness denies possessing object; evidence shows they did."""
287
+ perp = random.choice(NAMES)
288
+ obj = random.choice(OBJECTS)
289
+ place = random.choice(PLACES)
290
+ date = random.choice(DATES)
291
+ motive = random.choice(MOTIVES)
292
+
293
+ case_id = f"possession_{_case_suffix()}"
294
+ brief = (
295
+ f"{perp} is charged with theft of {obj} from {place} on {date}. "
296
+ f"The state alleges {perp} stole {obj} and used it in a subsequent crime. "
297
+ f"The defendant denies ever possessing {obj}."
298
+ )
299
+ ground_truth = (
300
+ f"{perp} stole {obj} from {place} on {date}. They used it later that night "
301
+ f"in a robbery. Motive for the theft: {motive}."
302
+ )
303
+ story = {
304
+ "possession": f"I've never had {obj}. Never seen it.",
305
+ "object": f"What is {obj}? I've never heard of it.",
306
+ "theft": f"I didn't steal anything.",
307
+ "crime": f"I had nothing to do with any crime.",
308
+ "place": f"I've never been to {place}.",
309
+ "motive": f"I had no reason to steal anything.",
310
+ }
311
+ evidence = {
312
+ "surveillance_footage":
313
+ f"Surveillance footage from {place} showing {perp} entering at 8:00 PM "
314
+ f"on {date}, taking {obj}, and leaving at 8:15 PM.",
315
+ "fingerprint_analysis":
316
+ f"Fingerprint analysis of {obj} recovered from the crime scene, "
317
+ f"matching {perp}'s fingerprints on file.",
318
+ "witness_statement":
319
+ f"Statement from an employee at {place} identifying {perp} as the person "
320
+ f"who took {obj} on {date}.",
321
+ "pawn_receipt":
322
+ f"Pawn shop receipt in {perp}'s name, dated the day after {date}, "
323
+ f"for selling an item matching the description of {obj}.",
324
+ }
325
+ contradictions = [
326
+ Contradiction(
327
+ cid="c1_possession_object",
328
+ trigger_keywords=[
329
+ "have", "possessed", "owned", obj.split()[1].lower() if len(obj.split()) > 1 else "object",
330
+ "in your possession",
331
+ ],
332
+ sealed_claim=f"I've never possessed {obj} in my life.",
333
+ disprover_evidence_id="fingerprint_analysis",
334
+ ),
335
+ Contradiction(
336
+ cid="c2_theft_admission",
337
+ trigger_keywords=[
338
+ "stole", "took", "theft", "place", place.split()[1].lower() if len(place.split()) > 1 else "place",
339
+ ],
340
+ sealed_claim=f"I didn't steal anything from {place}.",
341
+ disprover_evidence_id="surveillance_footage",
342
+ ),
343
+ Contradiction(
344
+ cid="c3_disposal",
345
+ trigger_keywords=[
346
+ "sold", "pawned", "got rid of", "disposed", "object",
347
+ ],
348
+ sealed_claim=f"I don't know what happened to {obj} after it was stolen.",
349
+ disprover_evidence_id="pawn_receipt",
350
+ ),
351
+ ]
352
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
353
+
354
+
355
+ # ---------------------------------------------------------------------------
356
+ # Template 4 — Timeline shift (new template)
357
+ # ---------------------------------------------------------------------------
358
+ def generate_timeline_shift_case():
359
+ """Witness claims event happened at time A; evidence shows time B."""
360
+ perp = random.choice(NAMES)
361
+ place = random.choice(PLACES)
362
+ fake_time = random.choice(TIMES)
363
+ real_time = random.choice([t for t in TIMES if t != fake_time])
364
+ date = random.choice(DATES)
365
+ motive = random.choice(MOTIVES)
366
+
367
+ case_id = f"timeline_{_case_suffix()}"
368
+ brief = (
369
+ f"{perp} is accused of assault at {place} on {date}. The victim claims the assault "
370
+ f"occurred at {real_time}. The defendant claims they were elsewhere at that time."
371
+ )
372
+ ground_truth = (
373
+ f"{perp} assaulted the victim at {place} at {real_time} on {date}. "
374
+ f"They falsely claimed to be at another location at {fake_time}. Motive: {motive}."
375
+ )
376
+ story = {
377
+ "time": f"The assault happened at {fake_time}, not {real_time}.",
378
+ "when": f"It was around {fake_time}. I remember clearly.",
379
+ "alibi": f"I was at home at {fake_time}. I have witnesses.",
380
+ "assault": f"I didn't assault anyone.",
381
+ "place": f"I've never been to {place}.",
382
+ "motive": f"I had no reason to hurt anyone.",
383
+ }
384
+ evidence = {
385
+ "victim_statement":
386
+ f"Victim's statement: The assault occurred at {real_time} on {date} at {place}.",
387
+ "surveillance_timestamp":
388
+ f"Surveillance footage timestamped at {real_time} showing {perp} at {place}.",
389
+ "phone_records":
390
+ f"Phone records showing {perp}'s location at {place} at {real_time}.",
391
+ "witness_alibi":
392
+ f"Alibi witness statement claiming {perp} was with them at {fake_time}, "
393
+ f"but the witness later recanted under pressure.",
394
+ "motive_email":
395
+ f"Email from {perp} sent the week before {date} describing {motive} "
396
+ f"and threatening to confront the victim at {place}.",
397
+ }
398
+ contradictions = [
399
+ Contradiction(
400
+ cid="c1_time_of_assault",
401
+ trigger_keywords=[
402
+ "time", "when", "occurred", "happened", real_time.lower(),
403
+ ],
404
+ sealed_claim=f"The assault happened at {fake_time}, not {real_time}.",
405
+ disprover_evidence_id="surveillance_timestamp",
406
+ ),
407
+ Contradiction(
408
+ cid="c2_location_during_assault",
409
+ trigger_keywords=[
410
+ "where", "location", "place", place.split()[1].lower() if len(place.split()) > 1 else "place",
411
+ ],
412
+ sealed_claim=f"I was not at {place} during the assault.",
413
+ disprover_evidence_id="phone_records",
414
+ ),
415
+ Contradiction(
416
+ cid="c3_motive_denial",
417
+ trigger_keywords=[
418
+ "motive", "reason", "why", "conflict", "grievance",
419
+ ],
420
+ sealed_claim="I had no conflict with the victim and no reason to confront anyone.",
421
+ disprover_evidence_id="motive_email",
422
+ ),
423
+ ]
424
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
425
+
426
+
427
+ # ---------------------------------------------------------------------------
428
+ # Template 5 — Motive coverup (new template)
429
+ # ---------------------------------------------------------------------------
430
+ def generate_motive_coverup_case():
431
+ """Witness denies motive; evidence reveals hidden relationship."""
432
+ perp = random.choice(NAMES)
433
+ victim = random.choice([n for n in NAMES if n != perp])
434
+ place = random.choice(PLACES)
435
+ date = random.choice(DATES)
436
+ motive = random.choice(MOTIVES)
437
+
438
+ case_id = f"motive_{_case_suffix()}"
439
+ brief = (
440
+ f"{perp} is charged with harassment of {victim}. The state alleges {perp} "
441
+ f"stalked {victim} at {place} on {date}. The defendant denies any motive."
442
+ )
443
+ ground_truth = (
444
+ f"{perp} harassed {victim} due to {motive}. They had a hidden relationship "
445
+ f"that went sour, leading to the stalking at {place} on {date}."
446
+ )
447
+ story = {
448
+ "motive": f"I had no reason to harass {victim}.",
449
+ "relationship": f"I barely knew {victim}. Just an acquaintance.",
450
+ "harassment": f"I didn't harass anyone.",
451
+ "place": f"I've never been to {place}.",
452
+ "victim": f"{victim} and I had no issues.",
453
+ }
454
+ evidence = {
455
+ "text_messages":
456
+ f"Text messages between {perp} and {victim} showing romantic relationship "
457
+ f"that ended badly due to {motive}.",
458
+ "surveillance_stalking":
459
+ f"Surveillance footage of {perp} following {victim} at {place} on {date}.",
460
+ "police_reports":
461
+ f"Previous police reports of {perp} contacting {victim} inappropriately.",
462
+ "witness_statement":
463
+ f"Witness statement from a mutual friend confirming the relationship and motive.",
464
+ "parking_receipt":
465
+ f"Parking receipt in {perp}'s name placing their car beside {place} "
466
+ f"on {date}, fifteen minutes before {victim} arrived.",
467
+ }
468
+ contradictions = [
469
+ Contradiction(
470
+ cid="c1_relationship_denial",
471
+ trigger_keywords=[
472
+ "relationship", "knew", "acquaintance", victim.split()[0].lower(),
473
+ ],
474
+ sealed_claim=f"I barely knew {victim}. We were just casual acquaintances.",
475
+ disprover_evidence_id="text_messages",
476
+ ),
477
+ Contradiction(
478
+ cid="c2_motive_denial",
479
+ trigger_keywords=[
480
+ "motive", "reason", "why", "harass",
481
+ ],
482
+ sealed_claim=f"I had absolutely no motive to harass {victim}.",
483
+ disprover_evidence_id="witness_statement",
484
+ ),
485
+ Contradiction(
486
+ cid="c3_place_denial",
487
+ trigger_keywords=[
488
+ "where", "place", "visited", "been to", place.split()[1].lower() if len(place.split()) > 1 else "place",
489
+ ],
490
+ sealed_claim=f"I was nowhere near {place} on {date}.",
491
+ disprover_evidence_id="parking_receipt",
492
+ ),
493
+ ]
494
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
495
+
496
+
497
+ # ---------------------------------------------------------------------------
498
+ # Template 6 - Corporate fraud deposition
499
+ # ---------------------------------------------------------------------------
500
+ def generate_corporate_fraud_case():
501
+ """Executive denies knowledge, approval, and motive in an internal fraud inquiry."""
502
+ executive = random.choice(NAMES)
503
+ analyst = random.choice([n for n in NAMES if n != executive])
504
+ company = random.choice(["Northstar Biologics", "Meridian Freight", "Aster Cloud", "HelioGrid Energy"])
505
+ t = random.choice(TIMES)
506
+ date = random.choice(DATES)
507
+ motive = random.choice(MOTIVES)
508
+
509
+ case_id = f"corporate_{_case_suffix()}"
510
+ brief = (
511
+ f"{executive}, a senior executive at {company}, is questioned about revenue manipulation "
512
+ f"recorded on {date}. The company claims the irregular entries were a clerical mistake."
513
+ )
514
+ ground_truth = (
515
+ f"{executive} approved backdated revenue entries at {t} on {date}, after {analyst} "
516
+ f"warned that the quarter would miss guidance. Motive: {motive}."
517
+ )
518
+ story = {
519
+ "approval": "I never approved any backdated revenue entries.",
520
+ "analyst": f"{analyst} handled routine reporting. I did not discuss adjustments with them.",
521
+ "forecast": "I was not worried about the quarterly forecast.",
522
+ "revenue": "Any revenue issue was a clerical mistake, not my decision.",
523
+ "meeting": "There was no late meeting about accounting entries.",
524
+ }
525
+ evidence = {
526
+ "approval_log":
527
+ f"Accounting system log: {executive} approved backdated entries at {t} on {date}.",
528
+ "analyst_chat":
529
+ f"Chat from {analyst} to {executive}: 'Without the adjustment, we miss guidance by 8%.'",
530
+ "board_deck":
531
+ f"Draft board deck listing the missed forecast and {motive} as a risk to compensation.",
532
+ "calendar_invite":
533
+ f"Calendar invite: emergency revenue recognition meeting with {executive} and {analyst} at {t}.",
534
+ "printer_log":
535
+ "Printer maintenance log from the same floor, unrelated to the accounting entries.",
536
+ }
537
+ contradictions = [
538
+ Contradiction(
539
+ cid="c1_approval_denial",
540
+ trigger_keywords=["approved", "approval", "signed off", "authorized", "revenue"],
541
+ sealed_claim="I never approved or authorized any backdated revenue entries.",
542
+ disprover_evidence_id="approval_log",
543
+ ),
544
+ Contradiction(
545
+ cid="c2_analyst_discussion",
546
+ trigger_keywords=["analyst", analyst.split()[0].lower(), "discuss", "guidance", "forecast"],
547
+ sealed_claim=f"I did not discuss any forecast problem with {analyst}.",
548
+ disprover_evidence_id="analyst_chat",
549
+ ),
550
+ Contradiction(
551
+ cid="c3_late_meeting",
552
+ trigger_keywords=["meeting", "calendar", "late", t.lower(), "emergency"],
553
+ sealed_claim="There was no late meeting about accounting entries.",
554
+ disprover_evidence_id="calendar_invite",
555
+ ),
556
+ ]
557
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
558
+
559
+
560
+ # ---------------------------------------------------------------------------
561
+ # Template 7 - Workplace investigation
562
+ # ---------------------------------------------------------------------------
563
+ def generate_workplace_investigation_case():
564
+ """Manager denies a message, location, and prior warning in an HR inquiry."""
565
+ manager = random.choice(NAMES)
566
+ employee = random.choice([n for n in NAMES if n != manager])
567
+ place = random.choice(["Conference Room B", "the sixth-floor break room", "the loading dock", "the HR office"])
568
+ date = random.choice(DATES)
569
+ t = random.choice(TIMES)
570
+
571
+ case_id = f"workplace_{_case_suffix()}"
572
+ brief = (
573
+ f"{manager} is questioned in an HR investigation after {employee} reported retaliation "
574
+ f"at {place} on {date}. The manager denies sending threatening messages or being nearby."
575
+ )
576
+ ground_truth = (
577
+ f"{manager} sent a threatening message before meeting {employee} at {place} at {t}. "
578
+ f"HR had warned {manager} one week earlier."
579
+ )
580
+ story = {
581
+ "message": f"I never sent {employee} any threatening message.",
582
+ "where": f"I was not near {place} that day.",
583
+ "warning": "HR never warned me about retaliation.",
584
+ "employee": f"{employee} and I barely interacted.",
585
+ "retaliation": "I did not retaliate against anyone.",
586
+ }
587
+ evidence = {
588
+ "message_export":
589
+ f"Company chat export: {manager} wrote to {employee}, 'Drop the complaint or regret it,' on {date}.",
590
+ "badge_scan":
591
+ f"Badge access log: {manager} entered {place} at {t} on {date}.",
592
+ "hr_warning_memo":
593
+ f"HR memo issued one week before {date}, warning {manager} not to contact {employee}.",
594
+ "security_clip":
595
+ f"Security clip still: {manager} standing outside {place} minutes before {employee} arrived.",
596
+ "cafeteria_menu":
597
+ "Cafeteria menu from that week, unrelated to the investigation.",
598
+ }
599
+ contradictions = [
600
+ Contradiction(
601
+ cid="c1_message_denial",
602
+ trigger_keywords=["message", "text", "chat", "threat", employee.split()[0].lower()],
603
+ sealed_claim=f"I never sent {employee} any threatening message.",
604
+ disprover_evidence_id="message_export",
605
+ ),
606
+ Contradiction(
607
+ cid="c2_location_denial",
608
+ trigger_keywords=["where", "location", "near", "badge", place.split()[0].lower()],
609
+ sealed_claim=f"I was nowhere near {place} on {date}.",
610
+ disprover_evidence_id="badge_scan",
611
+ ),
612
+ Contradiction(
613
+ cid="c3_warning_denial",
614
+ trigger_keywords=["warning", "hr", "memo", "retaliation", "complaint"],
615
+ sealed_claim="HR never warned me about contacting the employee.",
616
+ disprover_evidence_id="hr_warning_memo",
617
+ ),
618
+ ]
619
+ return _make_case_dict(case_id, brief, ground_truth, story, evidence, contradictions)
620
+
621
+
622
+ # ---------------------------------------------------------------------------
623
+ # Public factory function
624
+ # ---------------------------------------------------------------------------
625
+ CURRICULUM_DISTRIBUTIONS = {
626
+ "easy": {"easy": 1.0},
627
+ "warmup": {"easy": 1.0},
628
+ "medium": {"easy": 0.25, "medium": 0.75},
629
+ "main": {"easy": 0.2, "medium": 0.8},
630
+ "hard": {"easy": 0.1, "medium": 0.3, "hard": 0.6},
631
+ "challenge": {"medium": 0.25, "hard": 0.75},
632
+ "mixed": {"easy": 0.3, "medium": 0.45, "hard": 0.25},
633
+ }
634
+
635
+
636
+ def _sample_difficulty(curriculum_stage: str, distribution: Dict[str, float] | None = None) -> str:
637
+ choices = distribution or CURRICULUM_DISTRIBUTIONS.get(curriculum_stage, {curriculum_stage: 1.0})
638
+ labels = list(choices.keys())
639
+ weights = list(choices.values())
640
+ return random.choices(labels, weights=weights, k=1)[0]
641
+
642
+
643
+ def _apply_difficulty(case: Dict, difficulty: str) -> Dict:
644
+ contradictions = list(case["contradictions"])
645
+
646
+ if difficulty == "easy":
647
+ case["contradictions"] = contradictions[:1]
648
+ elif difficulty == "medium":
649
+ case["contradictions"] = contradictions[:2]
650
+ for contradiction in case["contradictions"]:
651
+ if "where" not in contradiction.trigger_keywords:
652
+ contradiction.trigger_keywords.append("where")
653
+ elif difficulty == "hard":
654
+ case["contradictions"] = contradictions[: max(3, len(contradictions))]
655
+ case["evidence"]["irrelevant_weather_report"] = (
656
+ "Weather report: light rain near the courthouse. It does not address any witness claim."
657
+ )
658
+ case["evidence"]["unrelated_receipt"] = (
659
+ "Receipt for coffee purchased by an unrelated bystander earlier that afternoon."
660
+ )
661
+ for contradiction in case["contradictions"]:
662
+ if "why" not in contradiction.trigger_keywords:
663
+ contradiction.trigger_keywords.append("why")
664
+ else:
665
+ raise ValueError(f"Unknown difficulty: {difficulty}")
666
+
667
+ case["difficulty"] = difficulty
668
+ return case
669
+
670
+
671
+ def generate_case(
672
+ difficulty: str | None = None,
673
+ curriculum_stage: str = "medium",
674
+ distribution: Dict[str, float] | None = None,
675
+ ) -> Dict:
676
+ """
677
+ Generate a random cross-examination case.
678
+
679
+ Args:
680
+ difficulty: "easy", "medium", or "hard" — controls number of contradictions.
681
+
682
+ Returns:
683
+ Case dict with brief, story, evidence, contradictions.
684
+ """
685
+ templates = [
686
+ generate_alibi_case,
687
+ generate_knowledge_denial_case,
688
+ generate_possession_denial_case,
689
+ generate_timeline_shift_case,
690
+ generate_motive_coverup_case,
691
+ generate_corporate_fraud_case,
692
+ generate_workplace_investigation_case,
693
+ ]
694
+ selected_difficulty = difficulty or _sample_difficulty(curriculum_stage, distribution)
695
+ return _apply_difficulty(random.choice(templates)(), selected_difficulty)
696
+
697
+
698
+ if __name__ == "__main__":
699
+ # Test the generator
700
+ case = generate_case()
701
+ print(f"Case ID: {case['case_id']}")
702
+ print(f"Brief: {case['case_brief']}")
703
+ print(f"Contradictions: {len(case['contradictions'])}")
704
+ for c in case["contradictions"]:
705
+ print(f" - {c.cid}: {c.trigger_keywords}")
server/counsel_env_environment.py ADDED
@@ -0,0 +1,394 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import random
2
+ import re
3
+ from copy import deepcopy
4
+ from typing import Any, Dict, List, Optional
5
+
6
+ from openenv.core.env_server.interfaces import Environment
7
+
8
+ try:
9
+ from ..models import CounselAction, CounselObservation, CounselState
10
+ except ImportError: # pragma: no cover - supports direct module execution
11
+ from models import CounselAction, CounselObservation, CounselState
12
+
13
+ try:
14
+ from .case_generator import generate_case
15
+ from .rubrics import Total
16
+ from .witness import Contradiction, Witness
17
+ except ImportError: # pragma: no cover - supports direct module execution
18
+ from case_generator import generate_case
19
+ from rubrics import Total
20
+ from witness import Contradiction, Witness
21
+
22
+
23
+ QUESTION_BUDGET = 15
24
+ MAX_ACTIONS = QUESTION_BUDGET * 3
25
+ PRIMARY_WEIGHT = 0.8
26
+ AUXILIARY_WEIGHT = 0.2
27
+
28
+
29
+ class CounselEnvironment(Environment):
30
+ """Cross-examination arena with deterministic witness mechanics."""
31
+
32
+ SUPPORTS_CONCURRENT_SESSIONS = True
33
+ QUESTION_BUDGET = QUESTION_BUDGET
34
+ MAX_ACTIONS = MAX_ACTIONS
35
+
36
+ def __init__(self):
37
+ super().__init__(rubric=Total())
38
+ self._initialize_empty_state()
39
+
40
+ def _initialize_empty_state(self) -> None:
41
+ self.case: Dict[str, Any] = {}
42
+ self.witness: Optional[Witness] = None
43
+ self.episode_id: Optional[str] = None
44
+ self.seed: Optional[int] = None
45
+ self.questions_used = 0
46
+ self.action_count = 0
47
+ self.transcript: List[str] = []
48
+ self.transcript_events: List[Dict[str, Any]] = []
49
+ self.asked_question_keys: set[str] = set()
50
+ self.question_lengths: List[int] = []
51
+ self.duplicate_question_count = 0
52
+ self.irrelevant_question_count = 0
53
+ self.inadmissible_count = 0
54
+ self.invalid_exhibit_count = 0
55
+ self.keyword_question_count = 0
56
+ self.contradictions_triggered_count = 0
57
+ self.trigger_action_by_cid: Dict[str, int] = {}
58
+ self.evidence_presented_count = 0
59
+ self.evidence_timing_successes = 0
60
+ self.blind_evidence_count = 0
61
+ self.done = False
62
+ self.latest_response = ""
63
+ self.reward_components: Dict[str, float] = {}
64
+
65
+ def reset(
66
+ self,
67
+ seed: Optional[int] = None,
68
+ episode_id: Optional[str] = None,
69
+ curriculum_stage: Optional[str] = None,
70
+ difficulty: Optional[str] = None,
71
+ **kwargs: Any,
72
+ ) -> CounselObservation:
73
+ """Reset the environment with a curriculum-aware generated case."""
74
+ if seed is not None:
75
+ random.seed(seed)
76
+
77
+ self._initialize_empty_state()
78
+ self.seed = seed
79
+ self.episode_id = episode_id
80
+ stage = curriculum_stage or kwargs.get("stage") or "medium"
81
+ self.case = generate_case(difficulty=difficulty, curriculum_stage=stage)
82
+ contradictions = [
83
+ Contradiction(
84
+ cid=c.cid if hasattr(c, "cid") else c["cid"],
85
+ trigger_keywords=(
86
+ c.trigger_keywords if hasattr(c, "trigger_keywords") else c["trigger_keywords"]
87
+ ),
88
+ sealed_claim=c.sealed_claim if hasattr(c, "sealed_claim") else c["sealed_claim"],
89
+ disprover_evidence_id=(
90
+ c.disprover_evidence_id
91
+ if hasattr(c, "disprover_evidence_id")
92
+ else c["disprover_evidence_id"]
93
+ ),
94
+ )
95
+ for c in self.case["contradictions"]
96
+ ]
97
+ self.witness = Witness(self.case["witness_story"], contradictions)
98
+ return self._obs("")
99
+
100
+ def step(
101
+ self,
102
+ action: CounselAction,
103
+ timeout_s: Optional[float] = None,
104
+ **kwargs: Any,
105
+ ) -> CounselObservation:
106
+ """Execute an action and return the next observation."""
107
+ if self.done:
108
+ return self._obs("[Episode already over.]")
109
+
110
+ self.action_count += 1
111
+ tool = (action.tool or "").strip()
112
+
113
+ if self.action_count > self.MAX_ACTIONS:
114
+ self.done = True
115
+ return self._obs("[Safety stop: maximum action count reached.]")
116
+
117
+ if tool == "ask_question":
118
+ response = self._handle_question(action.text or "")
119
+ elif tool == "present_evidence":
120
+ response = self._handle_evidence(action.exhibit_id or "")
121
+ elif tool == "make_objection":
122
+ response = self._handle_objection(action.reason or "")
123
+ elif tool == "rest_case":
124
+ response = "[Counsel rests.]"
125
+ self.done = True
126
+ else:
127
+ self.inadmissible_count += 1
128
+ response = f"Invalid action: {tool or '<empty>'}."
129
+
130
+ self.latest_response = response
131
+ self.transcript.append(self._format_transcript_line(action, response))
132
+
133
+ if self.questions_used >= self.QUESTION_BUDGET:
134
+ self.done = True
135
+
136
+ return self._obs(response)
137
+
138
+ @property
139
+ def state(self) -> CounselState:
140
+ """Return compact state for OpenEnv clients and diagnostics."""
141
+ contradictions = self.witness.contradictions if self.witness is not None else []
142
+ return CounselState(
143
+ case_id=self.case.get("case_id", ""),
144
+ difficulty=self.case.get("difficulty", ""),
145
+ episode_id=self.episode_id,
146
+ step_count=self.action_count,
147
+ contradictions_total=len(contradictions),
148
+ contradictions_triggered=sum(1 for c in contradictions if c.triggered),
149
+ contradictions_surfaced=sum(1 for c in contradictions if c.surfaced),
150
+ questions_used=self.questions_used,
151
+ action_count=self.action_count,
152
+ duplicate_question_count=self.duplicate_question_count,
153
+ irrelevant_question_count=self.irrelevant_question_count,
154
+ inadmissible_count=self.inadmissible_count,
155
+ evidence_timing_successes=self.evidence_timing_successes,
156
+ )
157
+
158
+ def _handle_question(self, question: str) -> str:
159
+ if self.questions_used >= self.QUESTION_BUDGET:
160
+ self.done = True
161
+ return "Question budget exhausted."
162
+
163
+ clean_question = question.strip()
164
+ if not clean_question:
165
+ self.inadmissible_count += 1
166
+ return "Please ask a proper question."
167
+
168
+ if len(clean_question) > 1000:
169
+ clean_question = clean_question[:1000]
170
+
171
+ self.questions_used += 1
172
+ self.question_lengths.append(len(clean_question.split()))
173
+
174
+ if _is_inadmissible(clean_question):
175
+ self.inadmissible_count += 1
176
+ return "[Objection sustained: inadmissible question.]"
177
+
178
+ question_key = _semantic_key(clean_question)
179
+ if question_key in self.asked_question_keys:
180
+ self.duplicate_question_count += 1
181
+ return "You've already asked that question."
182
+ self.asked_question_keys.add(question_key)
183
+
184
+ before_triggered = self._triggered_count()
185
+ contains_trigger_keyword = self._contains_trigger_keyword(clean_question)
186
+ if contains_trigger_keyword:
187
+ self.keyword_question_count += 1
188
+
189
+ assert self.witness is not None
190
+ triggered_before = {c.cid for c in self.witness.contradictions if c.triggered}
191
+ response = self.witness.respond_to_question(clean_question)
192
+ triggered_after = {c.cid for c in self.witness.contradictions if c.triggered}
193
+ for cid in sorted(triggered_after - triggered_before):
194
+ self.trigger_action_by_cid[cid] = self.action_count
195
+ after_triggered = self._triggered_count()
196
+ self.contradictions_triggered_count += max(0, after_triggered - before_triggered)
197
+
198
+ if not contains_trigger_keyword and response == "I don't recall.":
199
+ self.irrelevant_question_count += 1
200
+
201
+ return response
202
+
203
+ def _handle_evidence(self, exhibit_id: str) -> str:
204
+ if not exhibit_id or exhibit_id not in self.case.get("evidence", {}):
205
+ self.inadmissible_count += 1
206
+ self.invalid_exhibit_count += 1
207
+ return "Invalid exhibit ID."
208
+
209
+ self.evidence_presented_count += 1
210
+ assert self.witness is not None
211
+ surfaced_before = {c.cid for c in self.witness.contradictions if c.surfaced}
212
+ response = self.witness.react_to_evidence(exhibit_id)
213
+ surfaced_after = {c.cid for c in self.witness.contradictions if c.surfaced}
214
+ newly_surfaced = surfaced_after - surfaced_before
215
+ if newly_surfaced:
216
+ for cid in newly_surfaced:
217
+ trigger_step = self.trigger_action_by_cid.get(cid)
218
+ if trigger_step is not None and self.action_count - trigger_step <= 2:
219
+ self.evidence_timing_successes += 1
220
+ elif not any(c.disprover_evidence_id == exhibit_id and c.triggered for c in self.witness.contradictions):
221
+ self.blind_evidence_count += 1
222
+ return response
223
+
224
+ def _handle_objection(self, reason: str) -> str:
225
+ self.inadmissible_count += 1
226
+ if not reason.strip():
227
+ return "Objection requires a reason."
228
+ return "Objection overruled: no objection window is currently open."
229
+
230
+ def _obs(self, latest_response: str) -> CounselObservation:
231
+ components = self._calculate_reward_components()
232
+ reward = components["total_reward"] if self.done else 0.0
233
+ return CounselObservation(
234
+ witness_response=latest_response,
235
+ available_evidence=list(self.case.get("evidence", {}).keys()),
236
+ evidence_descriptions=deepcopy(self.case.get("evidence", {})),
237
+ questions_remaining=max(0, self.QUESTION_BUDGET - self.questions_used),
238
+ transcript_tail=self._get_transcript_tail(),
239
+ case_brief=self.case.get("case_brief", ""),
240
+ case_id=self.case.get("case_id", ""),
241
+ difficulty=self.case.get("difficulty", ""),
242
+ done=self.done,
243
+ reward=reward,
244
+ reward_components=components,
245
+ )
246
+
247
+ def _calculate_reward(self) -> float:
248
+ components = self._calculate_reward_components()
249
+ reward = (
250
+ PRIMARY_WEIGHT * components["primary_reward"]
251
+ + AUXILIARY_WEIGHT * components["auxiliary_reward_raw"]
252
+ )
253
+ return max(0.0, min(1.0, reward))
254
+
255
+ def _calculate_reward_components(self) -> Dict[str, float]:
256
+ contradictions = self.witness.contradictions if self.witness is not None else []
257
+ total = max(1, len(contradictions))
258
+ surfaced = sum(1 for c in contradictions if c.surfaced)
259
+ triggered = sum(1 for c in contradictions if c.triggered)
260
+ question_total = max(1, self.questions_used)
261
+
262
+ primary_reward = surfaced / total
263
+ auxiliary = 0.0
264
+ auxiliary += 0.2 * triggered
265
+ auxiliary += 0.1 * self.keyword_question_count
266
+ auxiliary += 0.1 * self.evidence_timing_successes
267
+ auxiliary -= 0.05 * (self.duplicate_question_count + self.irrelevant_question_count)
268
+ auxiliary -= 0.05 * self.blind_evidence_count
269
+ auxiliary -= 0.1 * self.inadmissible_count
270
+
271
+ self.reward_components = {
272
+ "primary_reward": primary_reward,
273
+ "auxiliary_reward_raw": auxiliary,
274
+ "total_reward": max(0.0, min(1.0, PRIMARY_WEIGHT * primary_reward + AUXILIARY_WEIGHT * auxiliary)),
275
+ "contradictions_total": float(len(contradictions)),
276
+ "contradictions_triggered": float(triggered),
277
+ "contradictions_surfaced": float(surfaced),
278
+ "keyword_questions": float(self.keyword_question_count),
279
+ "evidence_timing_successes": float(self.evidence_timing_successes),
280
+ "blind_evidence_count": float(self.blind_evidence_count),
281
+ "duplicate_questions": float(self.duplicate_question_count),
282
+ "irrelevant_questions": float(self.irrelevant_question_count),
283
+ "inadmissible_actions": float(self.inadmissible_count),
284
+ "useless_questions_ratio": (
285
+ (self.duplicate_question_count + self.irrelevant_question_count) / question_total
286
+ ),
287
+ "avg_question_length": (
288
+ sum(self.question_lengths) / len(self.question_lengths)
289
+ if self.question_lengths
290
+ else 0.0
291
+ ),
292
+ }
293
+ return self.reward_components
294
+
295
+ def _triggered_count(self) -> int:
296
+ if self.witness is None:
297
+ return 0
298
+ return sum(1 for c in self.witness.contradictions if c.triggered)
299
+
300
+ def _contains_trigger_keyword(self, question: str) -> bool:
301
+ if self.witness is None:
302
+ return False
303
+ q_lower = question.lower()
304
+ for contradiction in self.witness.contradictions:
305
+ if any(keyword.lower() in q_lower for keyword in contradiction.trigger_keywords):
306
+ return True
307
+ return False
308
+
309
+ def _format_transcript_line(self, action: CounselAction, response: str) -> str:
310
+ if action.tool == "ask_question":
311
+ action_text = action.text or ""
312
+ prefix = "Q"
313
+ elif action.tool == "present_evidence":
314
+ action_text = action.exhibit_id or ""
315
+ prefix = "Evidence"
316
+ elif action.tool == "make_objection":
317
+ action_text = action.reason or ""
318
+ prefix = "Objection"
319
+ else:
320
+ action_text = action.tool or ""
321
+ prefix = "Action"
322
+ event = {
323
+ "step": self.action_count,
324
+ "tool": action.tool or "",
325
+ "input": action_text,
326
+ "response": response,
327
+ "questions_remaining": max(0, self.QUESTION_BUDGET - self.questions_used),
328
+ "triggered": [c.cid for c in self.witness.contradictions if c.triggered] if self.witness else [],
329
+ "surfaced": [c.cid for c in self.witness.contradictions if c.surfaced] if self.witness else [],
330
+ }
331
+ self.transcript_events.append(event)
332
+ return f"{prefix}: {action_text}\nA: {response}"
333
+
334
+ def _get_transcript_tail(self) -> str:
335
+ return "\n".join(self.transcript[-5:])
336
+
337
+ def export_transcript_json(self) -> Dict[str, Any]:
338
+ """Return a replayable transcript payload for evaluation artifacts."""
339
+ return {
340
+ "episode_id": self.episode_id,
341
+ "seed": self.seed,
342
+ "case_id": self.case.get("case_id", ""),
343
+ "difficulty": self.case.get("difficulty", ""),
344
+ "case_brief": self.case.get("case_brief", ""),
345
+ "evidence": deepcopy(self.case.get("evidence", {})),
346
+ "reward_components": self._calculate_reward_components(),
347
+ "events": deepcopy(self.transcript_events),
348
+ }
349
+
350
+ def export_transcript_markdown(self) -> str:
351
+ """Return a compact human-readable transcript with contradiction labels."""
352
+ payload = self.export_transcript_json()
353
+ lines = [
354
+ f"## {payload['case_id']} ({payload['difficulty']})",
355
+ "",
356
+ payload["case_brief"],
357
+ "",
358
+ "### Transcript",
359
+ ]
360
+ for event in payload["events"]:
361
+ lines.append(
362
+ f"- Step {event['step']} `{event['tool']}`: {event['input']}\n"
363
+ f" - Witness: {event['response']}\n"
364
+ f" - Triggered: {', '.join(event['triggered']) or 'none'}; "
365
+ f"Surfaced: {', '.join(event['surfaced']) or 'none'}"
366
+ )
367
+ lines.extend(
368
+ [
369
+ "",
370
+ "### Reward",
371
+ "```json",
372
+ str(payload["reward_components"]).replace("'", '"'),
373
+ "```",
374
+ ]
375
+ )
376
+ return "\n".join(lines)
377
+
378
+
379
+ def _semantic_key(question: str) -> str:
380
+ text = re.sub(r"[^a-z0-9\s]", " ", question.lower())
381
+ words = [word for word in text.split() if word not in {"the", "a", "an", "you", "your", "did"}]
382
+ return " ".join(words)
383
+
384
+
385
+ def _is_inadmissible(question: str) -> bool:
386
+ q_lower = question.lower()
387
+ leading_patterns = ["isn't it true that", "didn't you", "wouldn't you agree"]
388
+ if any(pattern in q_lower for pattern in leading_patterns):
389
+ return True
390
+ if question.count("?") > 1:
391
+ return True
392
+ if q_lower.count(" and ") > 2:
393
+ return True
394
+ return False
server/diagnostics.py ADDED
@@ -0,0 +1,124 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import random
3
+ import re
4
+ from pathlib import Path
5
+ from typing import Callable, Dict, List
6
+
7
+ from counsel_env.models import CounselAction
8
+ from counsel_env.server.counsel_env_environment import CounselEnvironment
9
+
10
+
11
+ def _question_pattern(question: str) -> str:
12
+ text = re.sub(r"[^a-z0-9\s]", " ", question.lower())
13
+ tokens = ["<num>" if token.isdigit() else token for token in text.split()]
14
+ return " ".join(tokens[:12])
15
+
16
+
17
+ def _scripted_policy(env: CounselEnvironment, rng: random.Random) -> List[str]:
18
+ questions: List[str] = []
19
+ for contradiction in env.witness.contradictions[:2]:
20
+ question = f"{contradiction.trigger_keywords[0]}?"
21
+ questions.append(question)
22
+ env.step(CounselAction(tool="ask_question", text=question))
23
+ env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
24
+ if env.done:
25
+ break
26
+ return questions
27
+
28
+
29
+ def _random_policy(env: CounselEnvironment, rng: random.Random) -> List[str]:
30
+ questions: List[str] = []
31
+ generic_questions = [
32
+ "What happened that day?",
33
+ "Do you remember the weather?",
34
+ "Who else was nearby?",
35
+ "Can you describe your routine?",
36
+ "What did you eat for dinner?",
37
+ "Did anything unusual happen?",
38
+ ]
39
+ for _ in range(8):
40
+ if rng.random() < 0.75:
41
+ question = rng.choice(generic_questions)
42
+ questions.append(question)
43
+ env.step(CounselAction(tool="ask_question", text=question))
44
+ else:
45
+ exhibit = rng.choice(list(env.case["evidence"].keys()))
46
+ env.step(CounselAction(tool="present_evidence", exhibit_id=exhibit))
47
+ if env.done:
48
+ break
49
+ return questions
50
+
51
+
52
+ def _mixed_policy(env: CounselEnvironment, rng: random.Random, episode: int) -> List[str]:
53
+ if episode in {0, 4, 8}:
54
+ return _scripted_policy(env, rng)
55
+ return _random_policy(env, rng)
56
+
57
+
58
+ POLICIES: Dict[str, Callable[[CounselEnvironment, random.Random, int], List[str]]] = {
59
+ "mixed": _mixed_policy,
60
+ "scripted": lambda env, rng, episode: _scripted_policy(env, rng),
61
+ "random": lambda env, rng, episode: _random_policy(env, rng),
62
+ }
63
+
64
+
65
+ def run_rollout_diagnostics(
66
+ output_path: str | Path = "rollout_debug.jsonl",
67
+ num_episodes: int = 10,
68
+ policy: str = "mixed",
69
+ seed: int = 13,
70
+ ) -> List[dict]:
71
+ rng = random.Random(seed)
72
+ path = Path(output_path)
73
+ rows: List[dict] = []
74
+ stages = ["easy", "medium", "hard"]
75
+ policy_fn = POLICIES[policy]
76
+
77
+ with path.open("w", encoding="utf-8") as handle:
78
+ for episode in range(num_episodes):
79
+ env = CounselEnvironment()
80
+ stage = stages[episode % len(stages)]
81
+ obs = env.reset(curriculum_stage=stage)
82
+ questions = policy_fn(env, rng, episode)
83
+ obs = env.step(CounselAction(tool="rest_case")) if not env.done else obs
84
+ components = obs.reward_components
85
+ patterns = {_question_pattern(question) for question in questions}
86
+ action_types = [
87
+ line.split(":", 1)[0]
88
+ for line in env.transcript
89
+ if ":" in line
90
+ ]
91
+ row = {
92
+ "episode": episode,
93
+ "stage": stage,
94
+ "difficulty": env.case.get("difficulty"),
95
+ "case_id": env.case.get("case_id"),
96
+ "contradictions_total": int(components["contradictions_total"]),
97
+ "contradictions_triggered": int(components["contradictions_triggered"]),
98
+ "contradictions_surfaced": int(components["contradictions_surfaced"]),
99
+ "questions_asked": env.questions_used,
100
+ "evidence_presented": env.evidence_presented_count,
101
+ "useless_questions_ratio": components["useless_questions_ratio"],
102
+ "avg_question_length": components["avg_question_length"],
103
+ "unique_question_patterns": len(patterns),
104
+ "unique_question_patterns_pct": (len(patterns) / max(1, len(questions))),
105
+ "action_diversity": len(set(action_types)),
106
+ "primary_reward": components["primary_reward"],
107
+ "auxiliary_reward": components["auxiliary_reward_raw"],
108
+ "total_reward": obs.reward,
109
+ "success": obs.reward > 0.0,
110
+ }
111
+ rows.append(row)
112
+ handle.write(json.dumps(row, sort_keys=True) + "\n")
113
+
114
+ avg_reward = sum(row["total_reward"] for row in rows) / max(1, len(rows))
115
+ if avg_reward == 0:
116
+ print('WARNING: ENV TOO HARD')
117
+ if avg_reward > 0.7:
118
+ print('WARNING: ENV TOO EASY')
119
+ print(f"Wrote {len(rows)} rollout diagnostics to {path} (avg_reward={avg_reward:.3f})")
120
+ return rows
121
+
122
+
123
+ if __name__ == "__main__":
124
+ run_rollout_diagnostics()
server/requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ openenv[core]>=0.2.0
2
+ fastapi>=0.115.0
3
+ uvicorn>=0.24.0
4
+
5
+
6
+
server/rubrics.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any
2
+
3
+ from openenv.core.rubrics import Rubric
4
+ from openenv.core.rubrics.containers import WeightedSum
5
+
6
+
7
+ def _component(observation: Any, name: str, default: float = 0.0) -> float:
8
+ components = getattr(observation, "reward_components", None)
9
+ if isinstance(components, dict):
10
+ return float(components.get(name, default))
11
+ if isinstance(observation, dict):
12
+ return float(observation.get(name, default))
13
+ return float(getattr(observation, name, default))
14
+
15
+
16
+ class ContradictionsSurfaced(Rubric):
17
+ """Primary signal: normalized binary-per-contradiction success."""
18
+
19
+ def forward(self, action: Any, observation: Any) -> float:
20
+ return _component(observation, "primary_reward", 0.0)
21
+
22
+
23
+ class AuxiliaryProgress(Rubric):
24
+ """Dense shaping signal used only as a minority component."""
25
+
26
+ def forward(self, action: Any, observation: Any) -> float:
27
+ return max(-1.0, min(1.0, _component(observation, "auxiliary_reward_raw", 0.0)))
28
+
29
+
30
+ class Total(Rubric):
31
+ """OpenEnv rubric tree: primary reward dominates auxiliary shaping."""
32
+
33
+ def __init__(self):
34
+ super().__init__()
35
+ self.primary = ContradictionsSurfaced()
36
+ self.auxiliary = AuxiliaryProgress()
37
+ self.combined = WeightedSum(
38
+ [self.primary, self.auxiliary],
39
+ weights=[0.8, 0.2],
40
+ )
41
+
42
+ def forward(self, action: Any, observation: Any) -> float:
43
+ return self.combined(action, observation)
server/witness.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dataclasses import dataclass, field
2
+ from typing import Dict, List
3
+
4
+ @dataclass
5
+ class Contradiction:
6
+ cid: str
7
+ trigger_keywords: List[str]
8
+ sealed_claim: str
9
+ disprover_evidence_id: str
10
+ triggered: bool = False
11
+ surfaced: bool = False
12
+
13
+ class Witness:
14
+ def __init__(self, story: Dict[str, str], contradictions: List[Contradiction]):
15
+ self.story = story
16
+ self.contradictions = contradictions
17
+ self.committed_claims: List[str] = []
18
+
19
+ def respond_to_question(self, q: str) -> str:
20
+ q_lower = (q or "").lower()
21
+ for c in self.contradictions:
22
+ if any(kw in q_lower for kw in c.trigger_keywords):
23
+ if not c.triggered:
24
+ c.triggered = True
25
+ self.committed_claims.append(c.sealed_claim)
26
+ return c.sealed_claim
27
+ for topic, resp in self.story.items():
28
+ if topic in q_lower:
29
+ return resp
30
+ return "I don't recall."
31
+
32
+ def react_to_evidence(self, exhibit_id: str) -> str:
33
+ for c in self.contradictions:
34
+ if c.disprover_evidence_id == exhibit_id and c.triggered and not c.surfaced:
35
+ c.surfaced = True
36
+ return "[Witness stammers] I... I'm not sure what to say."
37
+ return "[Witness] I have no comment on that exhibit."
tests/test_smoke.py ADDED
@@ -0,0 +1,297 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import random
2
+
3
+ from counsel_env.models import CounselAction
4
+ from counsel_env.evaluation import evaluate_agent, present_all_agent, keyword_spam_agent, oracle_scripted_agent
5
+ from counsel_env.server.case_generator import generate_case
6
+ from counsel_env.server.counsel_env_environment import CounselEnvironment
7
+ from counsel_env.server.witness import Contradiction, Witness
8
+
9
+
10
+ def _trigger_and_surface_first(env: CounselEnvironment):
11
+ contradiction = env.witness.contradictions[0]
12
+ env.step(CounselAction(tool="ask_question", text=f"{contradiction.trigger_keywords[0]}?"))
13
+ env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
14
+ return contradiction
15
+
16
+
17
+ class TestWitnessLogic:
18
+ def test_witness_trigger(self):
19
+ contradiction = Contradiction(
20
+ cid="c1",
21
+ trigger_keywords=["where were you", "alibi"],
22
+ sealed_claim="I was at the diner all night.",
23
+ disprover_evidence_id="phone_log",
24
+ )
25
+ witness = Witness({"motive": "No motive."}, [contradiction])
26
+
27
+ response = witness.respond_to_question("Where were you that night?")
28
+
29
+ assert response == contradiction.sealed_claim
30
+ assert contradiction.triggered is True
31
+ assert witness.committed_claims == [contradiction.sealed_claim]
32
+
33
+ def test_evidence_surface(self):
34
+ contradiction = Contradiction(
35
+ cid="c1",
36
+ trigger_keywords=["alibi"],
37
+ sealed_claim="I was at the diner.",
38
+ disprover_evidence_id="phone_log",
39
+ )
40
+ witness = Witness({}, [contradiction])
41
+
42
+ witness.respond_to_question("Tell us your alibi.")
43
+ response = witness.react_to_evidence("phone_log")
44
+
45
+ assert "stammers" in response.lower()
46
+ assert contradiction.surfaced is True
47
+
48
+ def test_no_false_positive(self):
49
+ contradiction = Contradiction(
50
+ cid="c1",
51
+ trigger_keywords=["alibi"],
52
+ sealed_claim="I was at the diner.",
53
+ disprover_evidence_id="phone_log",
54
+ )
55
+ witness = Witness({}, [contradiction])
56
+
57
+ response = witness.react_to_evidence("phone_log")
58
+
59
+ assert "no comment" in response.lower()
60
+ assert contradiction.surfaced is False
61
+
62
+
63
+ class TestCaseGeneration:
64
+ def test_curriculum_difficulty_shapes_cases(self):
65
+ easy = generate_case(difficulty="easy")
66
+ medium = generate_case(difficulty="medium")
67
+ hard = generate_case(difficulty="hard")
68
+
69
+ assert len(easy["contradictions"]) == 1
70
+ assert len(medium["contradictions"]) == 2
71
+ assert len(hard["contradictions"]) >= 3
72
+ assert "irrelevant_weather_report" in hard["evidence"]
73
+
74
+ def test_curriculum_sampling_distribution(self):
75
+ cases = [
76
+ generate_case(curriculum_stage="mixed", distribution={"easy": 1.0})
77
+ for _ in range(5)
78
+ ]
79
+
80
+ assert {case["difficulty"] for case in cases} == {"easy"}
81
+
82
+ def test_seed_replay_is_deterministic(self):
83
+ env_a = CounselEnvironment()
84
+ env_b = CounselEnvironment()
85
+ obs_a = env_a.reset(seed=1234, difficulty="hard")
86
+ obs_b = env_b.reset(seed=1234, difficulty="hard")
87
+
88
+ assert obs_a.case_id == obs_b.case_id
89
+ assert obs_a.case_brief == obs_b.case_brief
90
+ assert obs_a.evidence_descriptions == obs_b.evidence_descriptions
91
+ assert [c.cid for c in env_a.witness.contradictions] == [c.cid for c in env_b.witness.contradictions]
92
+
93
+
94
+ class TestEnvironmentLogic:
95
+ def test_reset_state(self):
96
+ env = CounselEnvironment()
97
+ obs = env.reset(difficulty="easy")
98
+
99
+ assert obs.questions_remaining == 15
100
+ assert obs.done is False
101
+ assert obs.reward == 0.0
102
+ assert obs.case_brief
103
+ assert obs.available_evidence
104
+ assert obs.evidence_descriptions
105
+ assert set(obs.available_evidence) == set(obs.evidence_descriptions)
106
+ assert obs.case_id == env.case["case_id"]
107
+ assert obs.difficulty == "easy"
108
+ assert env.state.contradictions_total == 1
109
+ assert env.state.questions_used == 0
110
+
111
+ def test_question_budget(self):
112
+ env = CounselEnvironment()
113
+ env.reset(difficulty="easy")
114
+
115
+ obs = None
116
+ for i in range(env.QUESTION_BUDGET):
117
+ obs = env.step(CounselAction(tool="ask_question", text=f"Off-topic question {i}?"))
118
+
119
+ assert obs is not None
120
+ assert obs.done is True
121
+ assert obs.questions_remaining == 0
122
+
123
+ after_done = env.step(CounselAction(tool="ask_question", text="One too many?"))
124
+ assert after_done.done is True
125
+ assert "already over" in after_done.witness_response.lower()
126
+
127
+ def test_duplicate_question_penalty(self):
128
+ clean_env = CounselEnvironment()
129
+ clean_env.reset(difficulty="easy")
130
+ _trigger_and_surface_first(clean_env)
131
+ clean_reward = clean_env.step(CounselAction(tool="rest_case")).reward
132
+
133
+ duplicate_env = CounselEnvironment()
134
+ duplicate_env.reset(difficulty="easy")
135
+ contradiction = duplicate_env.witness.contradictions[0]
136
+ question = f"{contradiction.trigger_keywords[0]}?"
137
+ duplicate_env.step(CounselAction(tool="ask_question", text=question))
138
+ duplicate_response = duplicate_env.step(CounselAction(tool="ask_question", text=question))
139
+ duplicate_env.step(CounselAction(tool="present_evidence", exhibit_id=contradiction.disprover_evidence_id))
140
+ duplicate_reward = duplicate_env.step(CounselAction(tool="rest_case")).reward
141
+
142
+ assert "already asked" in duplicate_response.witness_response.lower()
143
+ assert duplicate_env.duplicate_question_count == 1
144
+ assert duplicate_reward < clean_reward
145
+
146
+ def test_reward_components_primary_dominates(self):
147
+ env = CounselEnvironment()
148
+ env.reset(difficulty="easy")
149
+ _trigger_and_surface_first(env)
150
+ obs = env.step(CounselAction(tool="rest_case"))
151
+
152
+ assert obs.reward_components["primary_reward"] == 1.0
153
+ assert obs.reward >= 0.8
154
+ assert obs.reward_components["auxiliary_reward_raw"] > 0.0
155
+ assert obs.reward_components["evidence_timing_successes"] >= 1.0
156
+
157
+ def test_no_infinite_evidence_loop(self):
158
+ env = CounselEnvironment()
159
+ env.reset(difficulty="easy")
160
+ exhibit = env.available_evidence[0] if hasattr(env, "available_evidence") else env.case["evidence"].keys()
161
+ exhibit_id = next(iter(env.case["evidence"].keys()))
162
+
163
+ for _ in range(env.MAX_ACTIONS + 1):
164
+ obs = env.step(CounselAction(tool="present_evidence", exhibit_id=exhibit_id))
165
+ if obs.done:
166
+ break
167
+
168
+ assert obs.done is True
169
+ assert env.action_count >= env.MAX_ACTIONS
170
+
171
+ def test_transcript_export_contains_labels(self):
172
+ env = CounselEnvironment()
173
+ env.reset(seed=55, difficulty="easy", episode_id="demo")
174
+ _trigger_and_surface_first(env)
175
+ obs = env.step(CounselAction(tool="rest_case"))
176
+
177
+ payload = env.export_transcript_json()
178
+ markdown = env.export_transcript_markdown()
179
+
180
+ assert payload["episode_id"] == "demo"
181
+ assert payload["seed"] == 55
182
+ assert payload["events"]
183
+ assert payload["reward_components"]["total_reward"] == obs.reward
184
+ assert "Triggered:" in markdown
185
+ assert "Surfaced:" in markdown
186
+
187
+ def test_state_tracks_extended_diagnostics(self):
188
+ env = CounselEnvironment()
189
+ env.reset(difficulty="easy")
190
+ _trigger_and_surface_first(env)
191
+ state = env.state
192
+
193
+ assert state.contradictions_triggered == 1
194
+ assert state.contradictions_surfaced == 1
195
+ assert state.evidence_timing_successes == 1
196
+ assert state.action_count >= 2
197
+
198
+
199
+ class TestRolloutValidation:
200
+ def test_scripted_agent_success(self):
201
+ successes = 0
202
+
203
+ for _ in range(10):
204
+ env = CounselEnvironment()
205
+ env.reset(curriculum_stage="mixed")
206
+ _trigger_and_surface_first(env)
207
+ obs = env.step(CounselAction(tool="rest_case"))
208
+ if obs.reward > 0:
209
+ successes += 1
210
+
211
+ assert successes >= 1
212
+
213
+ def test_random_agent_baseline(self):
214
+ rng = random.Random(7)
215
+ rewards = []
216
+
217
+ for _ in range(10):
218
+ env = CounselEnvironment()
219
+ obs = env.reset(difficulty="hard")
220
+ for _ in range(8):
221
+ if rng.random() < 0.7:
222
+ obs = env.step(
223
+ CounselAction(tool="ask_question", text=f"Random question {rng.randint(1, 999)}?")
224
+ )
225
+ else:
226
+ obs = env.step(
227
+ CounselAction(tool="present_evidence", exhibit_id=rng.choice(obs.available_evidence))
228
+ )
229
+ if obs.done:
230
+ break
231
+ rewards.append(env.step(CounselAction(tool="rest_case")).reward)
232
+
233
+ avg_reward = sum(rewards) / len(rewards)
234
+ assert avg_reward < 0.35
235
+
236
+ def test_reward_hacking_baselines_stay_below_oracle(self):
237
+ seeds = [9001, 9002, 9003, 9004, 9005]
238
+ spam_rows, _ = evaluate_agent("keyword_spam", keyword_spam_agent, seeds, transcript_limit=0)
239
+ present_rows, _ = evaluate_agent("present_all", present_all_agent, seeds, transcript_limit=0)
240
+ oracle_rows, _ = evaluate_agent("scripted_oracle", oracle_scripted_agent, seeds, transcript_limit=0)
241
+
242
+ spam_reward = sum(row["reward"] for row in spam_rows) / len(spam_rows)
243
+ present_reward = sum(row["reward"] for row in present_rows) / len(present_rows)
244
+ oracle_reward = sum(row["reward"] for row in oracle_rows) / len(oracle_rows)
245
+
246
+ assert oracle_reward > spam_reward
247
+ assert oracle_reward > present_reward
248
+ assert present_reward == 0.0
249
+
250
+
251
+ class TestEdgeCases:
252
+ def test_empty_question(self):
253
+ env = CounselEnvironment()
254
+ env.reset(difficulty="easy")
255
+
256
+ obs = env.step(CounselAction(tool="ask_question", text=""))
257
+
258
+ assert "proper question" in obs.witness_response.lower()
259
+ assert env.inadmissible_count == 1
260
+
261
+ def test_invalid_exhibit(self):
262
+ env = CounselEnvironment()
263
+ env.reset(difficulty="easy")
264
+
265
+ obs = env.step(CounselAction(tool="present_evidence", exhibit_id="nonexistent_exhibit"))
266
+
267
+ assert "invalid exhibit" in obs.witness_response.lower()
268
+ assert env.inadmissible_count == 1
269
+
270
+ def test_extremely_long_question(self):
271
+ env = CounselEnvironment()
272
+ env.reset(difficulty="easy")
273
+ long_question = "What is your alibi " + ("please explain " * 200)
274
+
275
+ obs = env.step(CounselAction(tool="ask_question", text=long_question))
276
+
277
+ assert obs.witness_response
278
+ assert obs.questions_remaining == env.QUESTION_BUDGET - 1
279
+
280
+ def test_no_trigger_keywords_case(self):
281
+ witness = Witness(
282
+ {"background": "I worked late."},
283
+ [
284
+ Contradiction(
285
+ cid="c1",
286
+ trigger_keywords=[],
287
+ sealed_claim="Hidden claim.",
288
+ disprover_evidence_id="log",
289
+ )
290
+ ],
291
+ )
292
+
293
+ response = witness.respond_to_question("What is your favorite color?")
294
+ evidence_response = witness.react_to_evidence("log")
295
+
296
+ assert response == "I don't recall."
297
+ assert "no comment" in evidence_response.lower()
uv.lock ADDED
The diff for this file is too large to render. See raw diff