Skip to content

Commit d9ed15d

Browse files
rycerzesburtenshaw
andauthored
feat: QED Math Environment with MCP tools and LLM-judge rubric (#865)
* initial qed-math OpenEnv scaffold with models and client - add QEDMathAction and QEDMathObservation dataclasses in models.py - implement QEDMathEnv client with reset, step, get_problem, submit_proof * implement MCP server tools and QEDMath environment - main env class - mcp server tools - step & reset logic * rubric implementation - impl MathProofRubric - LLM Grading Logic - rubric config in env * data pipeline integration - map problem data structure - support multiple types of problems * grading prompts * API & Client - wss handling - client methods * fix deps * proof submission with output length tokens and reasoning stripping - `submit_proof` method accepts an optional `output_length_tokens` parameter for reward shaping. - Introduced `remove_reasoning` function to strip reasoning traces from model output. - Added `length_penalty` function to compute penalties for overlong sequences. - Adjusted grading logic to apply discount factors and penalties based on token count. * fix eval prompt * original_problem field for grading in QED-Nano RC stream * add verifier metrics to grading results - implement metrics aggregation * docs dockerfile * tests for QED Math Env * add QED Math inference example - fix dockerfile * improve async step handling - refer to #456 * increase timeout for LLM judge calls and adjust WebSocket ping settings * add feedback output for incorrect answers in QED Math inference * refactor get_problem method to include metadata and adjust reference solution handling - update tests for answer mode * refactor math expression parsing and simplify async calls in tests * impl process-based math verification service - integrate with QED Math environment * gold answer caching and request admission control - tests * worker pool restart and health probe reporting * preserve custom observation fields during parsing - Fix MCP client parsing bug where reset/step observations were coerced into base Observation, dropping env-specific fields - tests * metrics tracking for MathVerifierService and QEDMathEnvironment * update env docs * refactor: remove output_length_tokens param from submit_proof and related methods * chore(lint): usort + ruff format * feat(qed_math): add upstream v0/v1 evaluator prompts Ship the QED-Nano v0 and v1 evaluator prompt templates verbatim alongside the existing v2. v1 is the QED-Nano default (full 0-7 score range); v2 constrains scores to {0,1,6,7}; v0 additionally uses the reference solution. * feat(qed_math): align defaults, grading, and reward routing with QED-Nano Bring the per-rollout reward signal in line with QED-Nano: - Default prompt_name v2 -> v1 (matches QED-Nano base.yaml; full 0-7 range). - Default reasoning_delimiters None -> ["</think>"] so only the post-think answer is graded, as upstream does. - Forward grader sampling params (grader_temperature=1.0, optional grader_max_output_tokens) to the Responses/Chat calls, and capture real token usage from the provider response (heuristic fallback only when absent). - Auto-detect proof vs answer mode from grading-schema presence, mirroring upstream's `if "schema" in problem` routing, when no explicit type is set. - Add answer_reward_preset (pure_success default, base adds wrong/no_answer/ unparsable penalties) keyed on verifier status; infra errors stay neutral. - Align is_correct success threshold to upstream score==7 (default 7). - Fix stale meta-pytorch/OpenEnv issue link in a docstring. * fix(qed_math): point build refs at canonical huggingface/OpenEnv The repo moved from meta-pytorch/OpenEnv to huggingface/OpenEnv. Update the Docker base image (ghcr.io/huggingface/openenv-base), the openenv-core git dependency in pyproject.toml, and the matching uv.lock source entries. * test(qed_math): cover QED-Nano fidelity changes Add tests for the new defaults and behavior: v1 default + v0/v1 prompt loading, default </think> reasoning stripping, proof/answer auto-detection from schema presence, answer reward presets (pure_success/base), success threshold == 7, grader sampling-param forwarding, and real token-usage capture. * docs(qed_math): refresh README, fix QED-Nano link, add docs stub Update the config tables and prose for the new defaults (v1 prompt, </think> stripping, grader sampling params, answer_reward_preset, success threshold 7, proof/answer auto-detection) and clarify token-usage metrics. Fix the broken QED-Nano source link (meta-pytorch -> CMU-AIRe) and add the generated docs/source/environments/qed_math.md stub. * fix: add pytest-asyncio, forward precision params, collapse verify timeout * fix: propagate reward/done in step_async, raise on unknown problem_id * fix: release verifier pool on close, collapse no-op imports, fix docstrings * fix: nested boxed regex, stale token count, queue metrics, dict grading_guidelines * fix: repair qed math client session helpers --------- Co-authored-by: burtenshaw <ben.burtenshaw@gmail.com>
1 parent cd37686 commit d9ed15d

23 files changed

Lines changed: 10308 additions & 3 deletions

File tree

Lines changed: 320 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,320 @@
1+
<!-- openenv-source: qed_math_env -->
2+
# QED Math Environment
3+
4+
Mathematical proof generation and evaluation environment for OpenEnv, ported from [QED-Nano](https://github.com/CMU-AIRe/QED-Nano). Agents receive math problems, submit proofs, and receive LLM-based rubric grading on a 0–7 scale with normalized rewards.
5+
6+
## Features
7+
8+
- **LLM-based rubric grading** (0–7 scale) via any OpenAI-compatible endpoint
9+
- **Process-based answer verification service** (`math_verify` in worker processes)
10+
- **Backpressure + retries + worker restart** for robust concurrent rollout operation
11+
- **Gold-answer cache** keyed by `problem_id` and verifier normalization settings
12+
- **Flexible dataset loading**: local JSONL/JSON, Hugging Face Hub, or built-in bootstrap problems
13+
- **Reward shaping**: discount factor, length penalty, and optional score thresholding
14+
- **Reasoning stripping**: configurable delimiters (e.g. `<think>...</think>`) removed before grading
15+
- **Multi-step problems**: configurable max attempts with per-attempt feedback
16+
- **Verifier metrics**: rollout/staging counters and health signals surfaced in observation metadata, ready for TrackIO / WandB
17+
- **MCP tool interface**: `get_problem`, `submit_proof`, `get_grading_guidelines`
18+
19+
## Quick Start
20+
21+
### Async (default)
22+
23+
```python
24+
import asyncio
25+
from qed_math_env import QEDMathEnv
26+
27+
async def main():
28+
async with QEDMathEnv(base_url="http://localhost:8000") as env:
29+
# Reset to load a problem
30+
result = await env.reset()
31+
obs = result.observation
32+
print(f"Problem: {obs.problem[:100]}...")
33+
34+
# Submit a proof
35+
submission = await env.submit_proof(proof="By induction on n...")
36+
print(f"Score: {submission.score}/7, Reward: {submission.reward:.2f}")
37+
38+
asyncio.run(main())
39+
```
40+
41+
### Sync
42+
43+
```python
44+
from qed_math_env import QEDMathEnv
45+
46+
with QEDMathEnv(base_url="http://localhost:8000").sync() as env:
47+
result = env.reset()
48+
submission = env.call_tool("submit_proof", proof="By induction on n...")
49+
```
50+
51+
### MCP tool-calling
52+
53+
```python
54+
async with QEDMathEnv(base_url="http://localhost:8000") as env:
55+
await env.reset()
56+
57+
# Discover tools
58+
tools = await env.list_tools()
59+
print([t.name for t in tools])
60+
# ['get_problem', 'submit_proof', 'get_grading_guidelines']
61+
62+
# Call tools by name
63+
problem = await env.call_tool("get_problem")
64+
guidelines = await env.call_tool("get_grading_guidelines")
65+
result = await env.call_tool("submit_proof", proof="...")
66+
```
67+
68+
## Building & Running
69+
70+
```bash
71+
# Build Docker image (from project root)
72+
docker build -t qed-math-env:latest -f envs/qed_math_env/server/Dockerfile .
73+
74+
# Run the server
75+
docker run -p 8000:8000 -e OPENAI_API_KEY=$OPENAI_API_KEY qed-math-env:latest
76+
77+
# Or run locally with uvicorn
78+
PYTHONPATH=src:envs uvicorn qed_math_env.server.app:app --port 8000
79+
80+
# Or install and run via uv
81+
cd envs/qed_math_env
82+
uv sync
83+
uv run server
84+
```
85+
86+
## Project Structure
87+
88+
```
89+
qed_math_env/
90+
├── __init__.py # Module exports (QEDMathEnv, models)
91+
├── models.py # ProblemObservation, ProofSubmissionObservation
92+
├── client.py # QEDMathEnv client (MCPToolClient subclass)
93+
├── openenv.yaml # OpenEnv manifest with metrics declarations
94+
├── pyproject.toml # Dependencies
95+
├── uv.lock # Locked dependencies
96+
├── README.md
97+
├── prompts/
98+
│ └── evaluator_prompts/
99+
│ ├── v0.md # Evaluator prompt (QED-Nano v0, uses reference solution)
100+
│ ├── v1.md # Evaluator prompt (QED-Nano v1, default, full 0–7 range)
101+
│ └── v2.md # Evaluator prompt (QED-Nano v2, scores constrained to {0,1,6,7})
102+
└── server/
103+
├── __init__.py
104+
├── app.py # FastAPI server (create_app factory)
105+
├── qed_math_environment.py # QEDMathEnvironment (MCPEnvironment)
106+
├── math_verify_service.py # Process-pool verifier service + health/metrics
107+
├── mcp_server.py # MCP tool registration
108+
├── rubric.py # MathProofRubric + GradingResult
109+
└── Dockerfile
110+
```
111+
112+
## Configuration
113+
114+
The environment is configured via `QEDMathConfig`:
115+
116+
| Parameter | Default | Description |
117+
|-----------|---------|-------------|
118+
| `dataset_path` | `None` | Dataset source: local path, Hub ID, or list of specs. `None` uses bootstrap problems. |
119+
| `grader_model` | `"gemini-3-pro"` | Model identifier for the LLM grader (any OpenAI-compatible endpoint) |
120+
| `prompt_name` | `"v1"` | Evaluator prompt template name (`v0`, `v1`, or `v2` in `prompts/evaluator_prompts/`). `v1` matches the QED-Nano default (full 0–7 range); `v2` constrains scores to `{0,1,6,7}` |
121+
| `grader_temperature` | `1.0` | Sampling temperature forwarded to the grader (matches QED-Nano) |
122+
| `grader_max_output_tokens` | `None` | Optional output-token cap for the grader. `None` uses the provider default to avoid truncating the trailing `<score>` tag |
123+
| `custom_reward_threshold` | `False` | When `True`, collapses partial-credit scores 1–5 → 1 |
124+
| `answer_reward_preset` | `"pure_success"` | Answer-mode reward table: `pure_success` (correct→1, else 0) or `base` (adds penalties: wrong −0.5, no_answer/unparsable −1) |
125+
| `max_attempts` | `1` | Max proof attempts per problem (>1 for multi-step) |
126+
| `discount_factor` | `1.0` | Exponential discount: `reward *= discount_factor ** output_length_tokens` |
127+
| `buffer_tokens` | `0` | Length penalty zone width. `0` disables the penalty. |
128+
| `max_tokens` | `0` | Max token limit for length penalty computation |
129+
| `reasoning_delimiters` | `["</think>"]` | Delimiter strings to strip reasoning before grading (matches QED-Nano). Set to `None` to grade the full completion. |
130+
| `verifier_workers` | `max(2, min(8, cpu_count//2))` | Number of process workers used for answer-mode verification |
131+
| `verifier_queue_size` | `verifier_workers * 32` | Max in-flight verifier requests before backpressure |
132+
| `verifier_request_timeout_seconds` | `5.0` | Per-request client-side timeout when awaiting worker response |
133+
| `verifier_max_retries` | `1` | Retry budget for transient verifier infra failures |
134+
| `verifier_strict` | `True` | Strict `math_verify` equivalence mode |
135+
| `verifier_numeric_precision` | `5` | Numeric precision setting used in verifier request contract |
136+
| `verifier_float_rounding` | `10` | Float rounding setting used in verifier request contract |
137+
138+
Environment variables:
139+
- `OPENAI_API_KEY` — API key for the grader LLM
140+
- `OPENAI_BASE_URL` — Base URL override (for non-OpenAI providers)
141+
142+
## Dataset Format
143+
144+
### Local JSONL/JSON
145+
146+
```json
147+
{
148+
"problem": "Prove that the sum of two even integers is even.",
149+
"solution": "Let a=2m and b=2n. Then a+b=2(m+n), which is even.",
150+
"rubrics": [
151+
{"title": "Definitions", "points": 2, "desc": "Correctly defines even integers."},
152+
{"title": "Algebra", "points": 3, "desc": "Valid algebraic manipulation."},
153+
{"title": "Conclusion", "points": 2, "desc": "Correctly concludes evenness."}
154+
],
155+
"dataset": "FineProofs-RL",
156+
"problem_id": "fp_001"
157+
}
158+
```
159+
160+
### Hugging Face Hub
161+
162+
```python
163+
QEDMathConfig(dataset_path="meta-math/MetaMathQA")
164+
# or with config
165+
QEDMathConfig(dataset_path={"hub_id": "meta-math/MetaMathQA", "split": "train", "config": "default"})
166+
```
167+
168+
### Field Aliases
169+
170+
The environment normalizes many dataset formats automatically:
171+
172+
| Canonical Field | Accepted Aliases |
173+
|----------------|------------------|
174+
| `problem` | `task`, `Problem` |
175+
| `reference_solution` | `solution`, `answer`, `Solution` |
176+
| `grading_guidelines` | `rubrics`, `schema`, `schema_0`, `Grading guidelines`, `details` |
177+
| `problem_id` | `id` |
178+
| `original_problem` | Used for RC-stream problems where the actor prompt differs from grading prompt |
179+
180+
## Observation Space
181+
182+
### ProblemObservation (from `reset` / `get_problem`)
183+
184+
| Field | Type | Description |
185+
|-------|------|-------------|
186+
| `problem` | `str` | Math problem statement |
187+
| `reference_solution` | `str` | Ground-truth solution |
188+
| `grading_guidelines` | `str` | Rubric / marking scheme |
189+
| `problem_id` | `str` | Unique identifier |
190+
| `problem_type` | `str` | `"proof"`, `"answer"`, or `"multi_step"` |
191+
| `dataset_source` | `str` | Source dataset name |
192+
| `metadata` | `dict` | Additional context (e.g. `original_problem`) |
193+
194+
### ProofSubmissionObservation (from `submit_proof`)
195+
196+
| Field | Type | Description |
197+
|-------|------|-------------|
198+
| `proof` | `str` | Submitted proof text |
199+
| `score` | `int` | Raw grade (0–7) |
200+
| `feedback` | `str` | Full grader response |
201+
| `reward` | `float` | Shaped reward in [0, 1] |
202+
| `done` | `bool` | Whether the episode is over |
203+
| `is_correct` | `bool` | Whether score >= success threshold (default 7, matching QED-Nano's `score == 7`) |
204+
| `attempt_number` | `int` | Current attempt count |
205+
| `attempts_remaining` | `int` | Remaining attempts |
206+
| `problem_type` | `str` | Problem type |
207+
| `metadata` | `dict` | Contains `verifier_metrics`, `base_reward`, `shaped_reward` |
208+
209+
## MCP Tools
210+
211+
| Tool | Description | Parameters |
212+
|------|-------------|------------|
213+
| `get_problem` | Return current problem statement and metadata ||
214+
| `submit_proof` | Submit a proof for LLM-based rubric grading | `proof` (str, required) |
215+
| `get_grading_guidelines` | Return the rubric/marking scheme ||
216+
217+
> **Note:** `output_length_tokens` is **not** an agent-supplied parameter. Token counts are
218+
> injected by the training harness via the HTTP step request body (see [Reward Shaping](#reward-shaping))
219+
> to preserve reward integrity — the agent cannot influence its own discount factor.
220+
221+
## Reward Shaping
222+
223+
The reward pipeline follows QED-Nano conventions:
224+
225+
1. **LLM grading**: Score 0–7 via evaluator prompt with `<score>N</score>` parsing
226+
2. **Optional thresholding**: Collapses 1–5 → 1 (when `custom_reward_threshold=True`)
227+
3. **Normalization**: `reward = score / 7.0`
228+
4. **Discount factor**: `reward *= discount_factor ** output_length_tokens`
229+
5. **Length penalty**: Linear penalty when output approaches `max_tokens`
230+
231+
For answer-mode problems (`evaluation_mode: "answer"`), grading is routed through the process-based verifier service: `\boxed{}` answers are extracted and verified against cached gold answers, with timeout/retry/backpressure handling for concurrent rollouts. The answer-mode reward is selected from `answer_reward_preset` keyed on the verifier status (`correct`, `wrong`, `no_answer`, `unparsable`); transient `timeout`/`internal_error` statuses stay neutral (0).
232+
233+
**Proof vs. answer routing:** when a dataset row does not set an explicit `problem_type`/`evaluation_mode`, the mode is auto-detected like QED-Nano's `if "schema" in problem` — a row carrying a grading rubric/schema is graded as a **proof** (LLM judge), while a row with no rubric is treated as an **answer** problem (boxed gold + `math_verify`). Set `problem_type`/`evaluation_mode` explicitly to override.
234+
235+
### Harness-injected token count
236+
237+
Steps 4 and 5 require the full generation length (including any reasoning trace that is stripped before grading). This value cannot come from the agent — it is supplied by the training harness as an out-of-band field in the HTTP step request body, mirroring the [`StateUsageTracker`](https://github.com/PrimeIntellect-ai/verifiers/blob/main/verifiers/utils/usage_utils.py) pattern from PrimeIntellect/verifiers:
238+
239+
```python
240+
# Training harness (pseudocode)
241+
completion_tokens = llm_call.usage.completion_tokens # from inference API
242+
243+
step_response = await openenv_client.step(
244+
action=CallToolAction(tool_name="submit_proof", arguments={"proof": proof_text}),
245+
output_length_tokens=completion_tokens, # injected here, not via MCP tool
246+
)
247+
```
248+
249+
When `output_length_tokens` is absent (local testing, eval without a training loop) shaping is skipped entirely — no estimation is attempted, consistent with verifiers' behaviour of returning `None` from `StateUsageTracker.snapshot()` when no usage was recorded.
250+
251+
## Verifier Metrics
252+
253+
Every `submit_proof` call emits verifier metrics in `metadata["verifier_metrics"]`, compatible with TrackIO and WandB:
254+
255+
| Metric | Description |
256+
|--------|-------------|
257+
| `verifier/rollouts/success` | 1 if grading succeeded |
258+
| `verifier/rollouts/failure` | 1 if grading failed |
259+
| `verifier/failures/timeout` | Count of timeout errors |
260+
| `verifier/failures/rate_limit` | Count of rate-limit errors |
261+
| `verifier/failures/no_input` | 1 if proof was empty |
262+
| `verifier/failures/no_score_tag` | 1 if LLM response had no `<score>` tag |
263+
| `verifier/failures/all_attempts_failed` | 1 if all retries exhausted |
264+
| `verifier/failures/num_retries` | Number of retries used |
265+
| `verifier/runtime/latency_per_request` | Grading wall-clock time (seconds) |
266+
| `verifier/requests/count` | Total verifier requests processed by the service |
267+
| `verifier/requests/latency_ms` | Service-level average request latency |
268+
| `verifier/requests/timeout_count` | Service-level timeout counter |
269+
| `verifier/requests/error_count` | Service-level internal error counter |
270+
| `verifier/queue/depth` | Current in-flight verifier queue depth |
271+
| `verifier/cache/hit_rate` | Gold-answer cache hit rate |
272+
| `verifier/workers/restart_count` | Worker-pool restart count |
273+
| `verifier/workers/worker_restarted` | 1 if current request required worker restart |
274+
| `verifier/workers/heartbeat_lag_ms` | Time since last verifier activity |
275+
| `verifier/runtime/input_tokens` | Grader input tokens (real provider usage when reported, else ~chars/4 estimate) |
276+
| `verifier/runtime/output_tokens` | Grader output tokens (real provider usage when reported, else ~chars/4 estimate) |
277+
| `reward/base` | Pre-shaping reward |
278+
| `reward/shaped` | Post-shaping reward |
279+
| `reward/score_raw` | Raw integer score (0–7) |
280+
| `reward/overlong_penalty` | Length penalty applied |
281+
| `episode/attempt_number` | Current attempt |
282+
| `episode/is_correct` | 1 if correct |
283+
| `episode/problem_type` | proof / answer / multi_step |
284+
| `episode/dataset_source` | Source dataset name |
285+
286+
### TrackIO Integration
287+
288+
```python
289+
import trackio
290+
291+
run = trackio.init(project="qed-math-training")
292+
293+
# After each submit_proof call:
294+
verifier_metrics = result["metadata"]["verifier_metrics"]
295+
numeric = {k: v for k, v in verifier_metrics.items() if isinstance(v, (int, float))}
296+
run.log(numeric, step=global_step)
297+
```
298+
299+
Or with TRL's GRPOTrainer:
300+
301+
```python
302+
from trl import GRPOConfig
303+
304+
config = GRPOConfig(
305+
report_to="trackio",
306+
trackio_space_id="your-org/qed-math-grpo",
307+
# ...
308+
)
309+
```
310+
311+
## Deployment
312+
313+
```bash
314+
# Optional: run rollout/staging verifier validation first
315+
PYTHONPATH=src:envs uv run python scripts/qed_math_verifier_staging_validation.py \
316+
--workers 4 --queue-size 128 --concurrency 64 --requests 2000 \
317+
--max-timeout-rate 0.05 --max-error-rate 0.02
318+
319+
openenv push
320+
```

envs/qed_math_env/.dockerignore

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
.venv
2+
.git
3+
.gitignore
4+
.env
5+
__pycache__/
6+
*.pyc
7+
*.pyo
8+
*.pyd
9+
*.pyw
10+
*.pyz
11+
*.pywz
12+
*.pyzw
13+
*.pyzwz

0 commit comments

Comments
 (0)