Spaces:
Sleeping
Sleeping
Upload folder using huggingface_hub
Browse files- Dockerfile +81 -0
- README.md +225 -4
- __init__.py +17 -0
- client.py +119 -0
- evaluation.py +337 -0
- models.py +38 -0
- notebooks/train_counsel.ipynb +327 -0
- openenv.yaml +10 -0
- pyproject.toml +35 -0
- server/__init__.py +14 -0
- server/app.py +89 -0
- server/case_generator.py +705 -0
- server/counsel_env_environment.py +394 -0
- server/diagnostics.py +124 -0
- server/requirements.txt +6 -0
- server/rubrics.py +43 -0
- server/witness.py +37 -0
- tests/test_smoke.py +297 -0
- uv.lock +0 -0
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:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: docker
|
| 7 |
pinned: false
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
---
|
| 9 |
|
| 10 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+

|
| 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 |
+

|
| 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
|
|
|