|
| 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 | +``` |
0 commit comments