QED Math Environment

Mathematical proof generation and evaluation environment for OpenEnv, ported from QED-Nano. Agents receive math problems, submit proofs, and receive LLM-based rubric grading on a 0–7 scale with normalized rewards.

Features

• LLM-based rubric grading (0–7 scale) via any OpenAI-compatible endpoint
• Process-based answer verification service (math_verify in worker processes)
• Backpressure + retries + worker restart for robust concurrent rollout operation
• Gold-answer cache keyed by problem_id and verifier normalization settings
• Flexible dataset loading: local JSONL/JSON, Hugging Face Hub, or built-in bootstrap problems
• Reward shaping: discount factor, length penalty, and optional score thresholding
• Reasoning stripping: configurable delimiters (e.g. <think>...</think>) removed before grading
• Multi-step problems: configurable max attempts with per-attempt feedback
• Verifier metrics: rollout/staging counters and health signals surfaced in observation metadata, ready for TrackIO / WandB
• MCP tool interface: get_problem, submit_proof, get_grading_guidelines

Quick Start

Async (default)

import asyncio
from qed_math_env import QEDMathEnv

async def main():
    async with QEDMathEnv(base_url="http://localhost:8000") as env:
        # Reset to load a problem
        result = await env.reset()
        obs = result.observation
        print(f"Problem: {obs.problem[:100]}...")

        # Submit a proof
        submission = await env.submit_proof(proof="By induction on n...")
        print(f"Score: {submission.score}/7, Reward: {submission.reward:.2f}")

asyncio.run(main())

Sync

from qed_math_env import QEDMathEnv

with QEDMathEnv(base_url="http://localhost:8000").sync() as env:
    result = env.reset()
    submission = env.call_tool("submit_proof", proof="By induction on n...")

MCP tool-calling

async with QEDMathEnv(base_url="http://localhost:8000") as env:
    await env.reset()

    # Discover tools
    tools = await env.list_tools()
    print([t.name for t in tools])
    # ['get_problem', 'submit_proof', 'get_grading_guidelines']

    # Call tools by name
    problem = await env.call_tool("get_problem")
    guidelines = await env.call_tool("get_grading_guidelines")
    result = await env.call_tool("submit_proof", proof="...")

Building & Running

# Build Docker image (from project root)
docker build -t qed-math-env:latest -f envs/qed_math_env/server/Dockerfile .

# Run the server
docker run -p 8000:8000 -e OPENAI_API_KEY=$OPENAI_API_KEY qed-math-env:latest

# Or run locally with uvicorn
PYTHONPATH=src:envs uvicorn qed_math_env.server.app:app --port 8000

# Or install and run via uv
cd envs/qed_math_env
uv sync
uv run server

Project Structure

qed_math_env/
├── __init__.py              # Module exports (QEDMathEnv, models)
├── models.py                # ProblemObservation, ProofSubmissionObservation
├── client.py                # QEDMathEnv client (MCPToolClient subclass)
├── openenv.yaml             # OpenEnv manifest with metrics declarations
├── pyproject.toml           # Dependencies
├── uv.lock                  # Locked dependencies
├── README.md
├── prompts/
│   └── evaluator_prompts/
│       ├── v0.md            # Evaluator prompt (QED-Nano v0, uses reference solution)
│       ├── v1.md            # Evaluator prompt (QED-Nano v1, default, full 0–7 range)
│       └── v2.md            # Evaluator prompt (QED-Nano v2, scores constrained to {0,1,6,7})
└── server/
    ├── __init__.py
    ├── app.py               # FastAPI server (create_app factory)
    ├── qed_math_environment.py  # QEDMathEnvironment (MCPEnvironment)
    ├── math_verify_service.py   # Process-pool verifier service + health/metrics
    ├── mcp_server.py        # MCP tool registration
    ├── rubric.py            # MathProofRubric + GradingResult
    └── Dockerfile

Configuration

The environment is configured via QEDMathConfig:

Environment variables:

• OPENAI_API_KEY — API key for the grader LLM
• OPENAI_BASE_URL — Base URL override (for non-OpenAI providers)

Dataset Format

Local JSONL/JSON

{
  "problem": "Prove that the sum of two even integers is even.",
  "solution": "Let a=2m and b=2n. Then a+b=2(m+n), which is even.",
  "rubrics": [
    {"title": "Definitions", "points": 2, "desc": "Correctly defines even integers."},
    {"title": "Algebra", "points": 3, "desc": "Valid algebraic manipulation."},
    {"title": "Conclusion", "points": 2, "desc": "Correctly concludes evenness."}
  ],
  "dataset": "FineProofs-RL",
  "problem_id": "fp_001"
}

Hugging Face Hub

QEDMathConfig(dataset_path="meta-math/MetaMathQA")
# or with config
QEDMathConfig(dataset_path={"hub_id": "meta-math/MetaMathQA", "split": "train", "config": "default"})

Field Aliases

The environment normalizes many dataset formats automatically:

Observation Space

ProblemObservation (from reset / get_problem)

ProofSubmissionObservation (from submit_proof)

MCP Tools

Note: output_length_tokens is not an agent-supplied parameter. Token counts are injected by the training harness via the HTTP step request body (see Reward Shaping) to preserve reward integrity — the agent cannot influence its own discount factor.

Reward Shaping

The reward pipeline follows QED-Nano conventions:

1. LLM grading: Score 0–7 via evaluator prompt with <score>N</score> parsing
2. Optional thresholding: Collapses 1–5 → 1 (when custom_reward_threshold=True)
3. Normalization: reward = score / 7.0
4. Discount factor: reward *= discount_factor ** output_length_tokens
5. Length penalty: Linear penalty when output approaches max_tokens

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

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.

Harness-injected token count

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 pattern from PrimeIntellect/verifiers:

# Training harness (pseudocode)
completion_tokens = llm_call.usage.completion_tokens  # from inference API

step_response = await openenv_client.step(
    action=CallToolAction(tool_name="submit_proof", arguments={"proof": proof_text}),
    output_length_tokens=completion_tokens,  # injected here, not via MCP tool
)

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.

Verifier Metrics

Every submit_proof call emits verifier metrics in metadata["verifier_metrics"], compatible with TrackIO and WandB:

TrackIO Integration

import trackio

run = trackio.init(project="qed-math-training")

# After each submit_proof call:
verifier_metrics = result["metadata"]["verifier_metrics"]
numeric = {k: v for k, v in verifier_metrics.items() if isinstance(v, (int, float))}
run.log(numeric, step=global_step)

Or with TRL's GRPOTrainer:

from trl import GRPOConfig

config = GRPOConfig(
    report_to="trackio",
    trackio_space_id="your-org/qed-math-grpo",
    # ...
)

Deployment

# Optional: run rollout/staging verifier validation first
PYTHONPATH=src:envs uv run python scripts/qed_math_verifier_staging_validation.py \
  --workers 4 --queue-size 128 --concurrency 64 --requests 2000 \
  --max-timeout-rate 0.05 --max-error-rate 0.02

openenv push