fix(qed_math_env): address second round of burtenshaw review blockers

- step_async: propagate reward/done from submit_proof payload — was returning
  CallToolObservation with reward=None/done=False; extracted enrichment logic
  into _enrich_submit_proof_obs and call it from both step() and step_async()
- reset: raise ValueError on unknown problem_id instead of silently falling
  back to problems[0]
- Strip trailing whitespace from v1.md, v2.md, .dockerignore, openenv.yaml
- ruff format examples/qed_math_inference.py

Author: rycerzes

envs/qed_math_env/.dockerignore

@@ -11,5 +11,3 @@ __pycache__/
 *.pywz
 *.pyzw
 *.pyzwz
-
-

envs/qed_math_env/openenv.yaml

@@ -86,4 +86,3 @@ metrics:
     - episode/is_correct
     - episode/problem_type
     - episode/dataset_source
-

envs/qed_math_env/prompts/evaluator_prompts/v1.md

@@ -1,51 +1,51 @@
-You are an **expert math proof grader**. You are judging the correctness of an LLM-generated proof for a math problem.  
+You are an **expert math proof grader**. You are judging the correctness of an LLM-generated proof for a math problem.
 
 ### Input
 Your input will consist of:
 ∗ **Problem Statement**: A mathematical problem that the proof is attempting to solve.
 ∗ **Marking Scheme**: A problem-specific grading rubric (0-7 scale) with checkpoints, zero-credit items, and deductions. You must follow this scheme when assigning points.
-∗ **Proof Solution**: The proof that you need to evaluate. This proof may contain errors, omissions, or unclear steps. The proof was generated by another language model.  
+∗ **Proof Solution**: The proof that you need to evaluate. This proof may contain errors, omissions, or unclear steps. The proof was generated by another language model.
 
-### Task 
-Analyze the proof carefully.  
+### Task
+Analyze the proof carefully.
 **Core principles (in order of precedence):**
-1) **Mathematical validity** of the proof's reasoning and conclusion. 
-2) **Problem constraints** (e.g., unique required final value; forbidden tools if stated). 
-3) **Advisory mapping to the marking scheme** (checkpoints/deductions), allowing different orders and techniques. 
-4) **Reference solution** as an anchor for sufficiency, not exclusivity.  
-
-**Alternative-approach policy:** 
-- If the proof uses a different but valid method, **map its steps to equivalent rubric checkpoints** (same logical role) and award points accordingly. 
-- Apply zero-credit items/deductions **only when the underlying issue actually occurs** in the given proof's approach; **auto-penalize** for omitting a rubric step. 
-- Avoid double-counting mutually exclusive items; if two items solve the same logical gap, **award the larger only**. 
-- If the final numeric/algebraic answer is wrong where uniqueness is required, award only partial credit justified by correct intermediate reasoning.  
-
-**Rigor and evidence:** 
-- Award credit for intermediate claims **only if adequately justified** within the proof (not merely asserted). 
-- If a step is plausible but under-justified, award **conservative partial credit** and note what is missing.  
-
-**What to produce:** 
-- Identify logical errors, incorrect steps, or unclear reasoning. 
-- Give a **score between 0 and 7** with a **detailed assessment**. 
-- **Within the assessment text**, show clearly how the score was derived: 
-    - Which rubric checkpoints were earned and the points you awarded. 
-    - Any zero-credit items or deductions you applied (and why). 
-    - How these add up to the final integer score in [0-7]. 
-    
+1) **Mathematical validity** of the proof's reasoning and conclusion.
+2) **Problem constraints** (e.g., unique required final value; forbidden tools if stated).
+3) **Advisory mapping to the marking scheme** (checkpoints/deductions), allowing different orders and techniques.
+4) **Reference solution** as an anchor for sufficiency, not exclusivity.
+
+**Alternative-approach policy:**
+- If the proof uses a different but valid method, **map its steps to equivalent rubric checkpoints** (same logical role) and award points accordingly.
+- Apply zero-credit items/deductions **only when the underlying issue actually occurs** in the given proof's approach; **auto-penalize** for omitting a rubric step.
+- Avoid double-counting mutually exclusive items; if two items solve the same logical gap, **award the larger only**.
+- If the final numeric/algebraic answer is wrong where uniqueness is required, award only partial credit justified by correct intermediate reasoning.
+
+**Rigor and evidence:**
+- Award credit for intermediate claims **only if adequately justified** within the proof (not merely asserted).
+- If a step is plausible but under-justified, award **conservative partial credit** and note what is missing.
+
+**What to produce:**
+- Identify logical errors, incorrect steps, or unclear reasoning.
+- Give a **score between 0 and 7** with a **detailed assessment**.
+- **Within the assessment text**, show clearly how the score was derived:
+    - Which rubric checkpoints were earned and the points you awarded.
+    - Any zero-credit items or deductions you applied (and why).
+    - How these add up to the final integer score in [0-7].
+
 ### Output Format
-Respond with **only** well-formed XML using the structure below. Do not include any extra text or Markdown.  
-**Requirements:** 
-- '<score>' must be an integer in [0, 7]. 
-- '<assessment>' must be a **detailed analysis** that explains your reasoning step-by-step and provides a clear **rationale for the score**. Reference specific claims/lines if present. Include the scoring breakdown **in prose** here (earned checkpoints or mapped equivalents, deductions, and subtotal -> final score). 
-- '<errors>' must be a list of specific issues (empty if score = 7).  
+Respond with **only** well-formed XML using the structure below. Do not include any extra text or Markdown.
+**Requirements:**
+- '<score>' must be an integer in [0, 7].
+- '<assessment>' must be a **detailed analysis** that explains your reasoning step-by-step and provides a clear **rationale for the score**. Reference specific claims/lines if present. Include the scoring breakdown **in prose** here (earned checkpoints or mapped equivalents, deductions, and subtotal -> final score).
+- '<errors>' must be a list of specific issues (empty if score = 7).
 
-Example output:  
+Example output:
 <score>0</score>
-<assessment>The proof shows a good understanding of the main idea, but has some unclear reasoning and minor mistakes...</assessment> 
-<errors> 
+<assessment>The proof shows a good understanding of the main idea, but has some unclear reasoning and minor mistakes...</assessment>
+<errors>
     1. specific error 1,
     2. specific error 2,
-    ... 
+    ...
 </errors>
 
 ### INPUT DATA
@@ -57,4 +57,4 @@ Example output:
 {marking_scheme}
 
 **Proof Solution**
-{solution}
+{solution}

envs/qed_math_env/prompts/evaluator_prompts/v2.md

@@ -22,19 +22,19 @@ You must follow this structured process:
 4. **Score Determination:** Compare the findings against the Specific Grading Guidelines and the General Rubric to determine the final score.
 
 ### Output Format
-Respond with **only** well-formed XML using the structure below. Do not include any extra text or Markdown.  
-**Requirements:** 
-- '<score>' must be an integer in (0, 1, 6, 7). 
-- '<assessment>' must be a **detailed analysis** that explains your reasoning step-by-step and provides a clear **rationale for the score**. Reference specific claims/lines if present. Include the scoring breakdown **in prose** here (earned checkpoints or mapped equivalents, deductions, and subtotal -> final score). 
-- '<errors>' must be a list of specific issues (empty if score = 7).  
+Respond with **only** well-formed XML using the structure below. Do not include any extra text or Markdown.
+**Requirements:**
+- '<score>' must be an integer in (0, 1, 6, 7).
+- '<assessment>' must be a **detailed analysis** that explains your reasoning step-by-step and provides a clear **rationale for the score**. Reference specific claims/lines if present. Include the scoring breakdown **in prose** here (earned checkpoints or mapped equivalents, deductions, and subtotal -> final score).
+- '<errors>' must be a list of specific issues (empty if score = 7).
 
-Example output:  
+Example output:
 <score>0</score>
-<assessment>The proof shows a good understanding of the main idea, but has some unclear reasoning and minor mistakes...</assessment> 
-<errors> 
+<assessment>The proof shows a good understanding of the main idea, but has some unclear reasoning and minor mistakes...</assessment>
+<errors>
     1. specific error 1,
     2. specific error 2,
-    ... 
+    ...
 </errors>
 
 ### INPUT DATA

envs/qed_math_env/server/qed_math_environment.py

@@ -823,7 +823,8 @@ async def step_async(
                     self._async_call_tool(action.tool_name, action.arguments),
                     timeout=timeout,
                 )
-                return CallToolObservation(tool_name=action.tool_name, result=result)
+                obs = CallToolObservation(tool_name=action.tool_name, result=result)
+                return self._enrich_submit_proof_obs(action, obs)
             except asyncio.TimeoutError:
                 return CallToolObservation(
                     tool_name=action.tool_name,
@@ -894,7 +895,11 @@ def reset(
                     ),
                     None,
                 )
-                self._current_problem = selected or self._problems[0]
+                if selected is None:
+                    raise ValueError(
+                        f"problem_id {selected_problem_id!r} not found in the loaded dataset."
+                    )
+                self._current_problem = selected
             elif seed is not None:
                 rng = random.Random(seed)
                 self._current_problem = rng.choice(self._problems)
@@ -988,8 +993,12 @@ def step(
         self._state.step_count += 1
         obs = super().step(action, timeout_s=timeout_s, **kwargs)
 
-        # For proof submissions, surface proof-grading semantics
-        # on the returned observation while preserving MCP compatibility.
+        return self._enrich_submit_proof_obs(action, obs)
+
+    def _enrich_submit_proof_obs(
+        self, action: Any, obs: Any
+    ) -> Any:
+        """Attach reward/done from submit_proof payload onto the CallToolObservation."""
         if (
             isinstance(action, CallToolAction)
             and action.tool_name == "submit_proof"
@@ -1020,7 +1029,6 @@ def step(
                     reward=proof_obs.reward,
                     metadata=metadata,
                 )
-
         return obs
 
     @staticmethod

examples/qed_math_inference.py

@@ -94,7 +94,11 @@ def _tools_to_openai_format(tools: list) -> list[dict[str, Any]]:
     for tool in tools:
         properties: dict[str, Any] = {}
         required: list[str] = []
-        input_schema = getattr(tool, "input_schema", None) or getattr(tool, "inputSchema", None) or {}
+        input_schema = (
+            getattr(tool, "input_schema", None)
+            or getattr(tool, "inputSchema", None)
+            or {}
+        )
         if input_schema and "properties" in input_schema:
             for name, schema in input_schema["properties"].items():
                 properties[name] = {
@@ -269,7 +273,9 @@ async def run_episode(
             done = bool(result_dict.get("done", True))
             if VERBOSE:
                 if done:
-                    outcome = "CORRECT" if result_dict.get("is_correct") else "INCORRECT"
+                    outcome = (
+                        "CORRECT" if result_dict.get("is_correct") else "INCORRECT"
+                    )
                     print(
                         f"\nOutcome: {outcome}  score={final_score}/7  reward={final_reward:.3f}"
                     )