Keep the qed_math_env answer-mode reference server-side and stop scoring after done in reasoning_gym_env and finqa_env

docs/source/environments/finqa.md

@@ -99,6 +99,8 @@ Tools are auto-discovered via MCP. Use `await env.list_tools()` to see all avail
 | `sql_query` | Execute SQL query on a table (requires filters) | `company_name: str, table_name: str, query: str` |
 | `submit_answer` | Submit final answer (ends episode) | `answer: str` |
 
+Once the episode has ended (an answer was submitted or `FINQA_MAX_STEPS` was reached), further `step()` calls (the WebSocket session) do not run tools or score answers. They return the final observation again (`done=True`, same reward) until the next `reset()`.
+
 ### Tool Constraints
 
 - **sql_query**: Must include filters (`WHERE`, `HAVING`, etc.). `SELECT *` is not allowed.

docs/source/environments/qed_math.md

@@ -184,7 +184,7 @@ The environment normalizes many dataset formats automatically:
 | Field | Type | Description |
 |-------|------|-------------|
 | `problem` | `str` | Math problem statement |
-| `reference_solution` | `str` | Ground-truth solution |
+| `reference_solution` | `str` | Always empty. The ground-truth solution stays on the server and is only used for grading. |
 | `grading_guidelines` | `str` | Rubric / marking scheme |
 | `problem_id` | `str` | Unique identifier |
 | `problem_type` | `str` | `"proof"`, `"answer"`, or `"multi_step"` |

docs/source/environments/reasoning_gym.md

@@ -76,6 +76,8 @@ Each episode is **single-step**:
 2. `step(answer)` returns score and marks episode as done
 3. `reset()` without params gets next question from same dataset
 
+Only the first answer of an episode is scored. Further `step()` calls before the next `reset()` return the same final observation (`done=True`, same reward) and are not scored.
+
 ### Action
 
 **ReasoningGymAction**: Contains the agent's answer

envs/finqa_env/README.md

@@ -113,6 +113,8 @@ Tools are auto-discovered via MCP. Use `await env.list_tools()` to see all avail
 | `sql_query` | Execute SQL query on a table (requires filters) | `company_name: str, table_name: str, query: str` |
 | `submit_answer` | Submit final answer (ends episode) | `answer: str` |
 
+Once the episode has ended (an answer was submitted or `FINQA_MAX_STEPS` was reached), further `step()` calls (the WebSocket session) do not run tools or score answers. They return the final observation again (`done=True`, same reward) until the next `reset()`.
+
 ### Tool Constraints
 
 - **sql_query**: Must include filters (`WHERE`, `HAVING`, etc.). `SELECT *` is not allowed.

envs/finqa_env/server/finqa_environment.py

@@ -131,6 +131,7 @@ def submit_answer(answer: str) -> str:
 
         self._state = FinQAState()
         self._history: List[Dict[str, Any]] = []
+        self._final_observation: Optional[Observation] = None
 
     def _load_questions(self) -> List[Dict[str, Any]]:
         """Load questions from the benchmark CSV."""
@@ -206,6 +207,7 @@ def reset(
             question_id=question["id"],
         )
         self._history = []
+        self._final_observation = None
 
         logger.info(
             f"Reset episode {self._state.episode_id} with question: {question['question'][:200]}..."
@@ -242,6 +244,11 @@ def _step_impl(
             },
         )
 
+    def _end_episode(self, observation: Observation) -> Observation:
+        """Keep the terminal observation; later steps return it again."""
+        self._final_observation = observation.model_copy(deep=True)
+        return observation
+
     def step(
         self,
         action: Action,
@@ -254,6 +261,11 @@ def step(
         Delegates to base class for MCP actions. Handles submit_answer
         reward computation and max-step termination.
         """
+        # After the episode ended (answer submitted or max steps reached), no
+        # tool runs and nothing is scored: return the final observation again.
+        if self._final_observation is not None:
+            return self._final_observation.model_copy(deep=True)
+
         self._state.step_count += 1
 
         # Let the base class handle MCP actions
@@ -267,28 +279,32 @@ def step(
                 f"Episode {self._state.episode_id} ended: "
                 f"submitted='{submitted_answer}', truth='{self._state.ground_truth}', reward={reward}"
             )
-            return Observation(
-                done=True,
-                reward=reward,
-                metadata={
-                    **obs.metadata,
-                    "ground_truth": self._state.ground_truth,
-                    "submitted_answer": submitted_answer,
-                },
+            return self._end_episode(
+                Observation(
+                    done=True,
+                    reward=reward,
+                    metadata={
+                        **obs.metadata,
+                        "ground_truth": self._state.ground_truth,
+                        "submitted_answer": submitted_answer,
+                    },
+                )
             )
 
         # Check for max steps
         if self._state.step_count >= self.max_steps:
             logger.info(
                 f"Episode {self._state.episode_id} terminated: max steps reached"
             )
-            return Observation(
-                done=True,
-                reward=0.0,
-                metadata={
-                    **obs.metadata,
-                    "error": f"Max steps ({self.max_steps}) reached without submitting answer.",
-                },
+            return self._end_episode(
+                Observation(
+                    done=True,
+                    reward=0.0,
+                    metadata={
+                        **obs.metadata,
+                        "error": f"Max steps ({self.max_steps}) reached without submitting answer.",
+                    },
+                )
             )
 
         return obs
@@ -306,6 +322,11 @@ async def step_async(
         termination) but delegates to MCPEnvironment.step_async to avoid
         the run_async_safely event-loop deadlock on the WebSocket path.
         """
+        # After the episode ended (answer submitted or max steps reached), no
+        # tool runs and nothing is scored: return the final observation again.
+        if self._final_observation is not None:
+            return self._final_observation.model_copy(deep=True)
+
         self._state.step_count += 1
 
         obs = await super().step_async(action, timeout_s=timeout_s, **kwargs)
@@ -317,27 +338,31 @@ async def step_async(
                 f"Episode {self._state.episode_id} ended: "
                 f"submitted='{submitted_answer}', truth='{self._state.ground_truth}', reward={reward}"
             )
-            return Observation(
-                done=True,
-                reward=reward,
-                metadata={
-                    **obs.metadata,
-                    "ground_truth": self._state.ground_truth,
-                    "submitted_answer": submitted_answer,
-                },
+            return self._end_episode(
+                Observation(
+                    done=True,
+                    reward=reward,
+                    metadata={
+                        **obs.metadata,
+                        "ground_truth": self._state.ground_truth,
+                        "submitted_answer": submitted_answer,
+                    },
+                )
             )
 
         if self._state.step_count >= self.max_steps:
             logger.info(
                 f"Episode {self._state.episode_id} terminated: max steps reached"
             )
-            return Observation(
-                done=True,
-                reward=0.0,
-                metadata={
-                    **obs.metadata,
-                    "error": f"Max steps ({self.max_steps}) reached without submitting answer.",
-                },
+            return self._end_episode(
+                Observation(
+                    done=True,
+                    reward=0.0,
+                    metadata={
+                        **obs.metadata,
+                        "error": f"Max steps ({self.max_steps}) reached without submitting answer.",
+                    },
+                )
             )
 
         return obs

envs/qed_math_env/README.md

@@ -199,7 +199,7 @@ The environment normalizes many dataset formats automatically:
 | Field | Type | Description |
 |-------|------|-------------|
 | `problem` | `str` | Math problem statement |
-| `reference_solution` | `str` | Ground-truth solution |
+| `reference_solution` | `str` | Always empty. The ground-truth solution stays on the server and is only used for grading. |
 | `grading_guidelines` | `str` | Rubric / marking scheme |
 | `problem_id` | `str` | Unique identifier |
 | `problem_type` | `str` | `"proof"`, `"answer"`, or `"multi_step"` |

envs/qed_math_env/models.py

@@ -45,7 +45,10 @@ class ProblemObservation(QEDMathObservation):
     """Observation containing the problem statement."""
 
     problem: str = Field(default="", description="The mathematical problem")
-    reference_solution: str = Field(default="", description="Ground truth solution")
+    reference_solution: str = Field(
+        default="",
+        description="Always empty. The ground truth stays server-side for grading.",
+    )
     grading_guidelines: str = Field(
         default="", description="Rubric for grading (0-7 scale)"
     )

envs/qed_math_env/server/mcp_server.py

@@ -9,8 +9,8 @@
 
 Tools are registered on the environment's FastMCP instance inside
 QEDMathEnvironment.__init__:
-- get_problem(): Return current problem and metadata; reference solution is
-    only included for answer-mode evaluation.
+- get_problem(): Return current problem and metadata. The reference solution
+    is not included; it stays server-side and is only used for grading.
 - submit_proof(proof): Grade proof via MathProofRubric and return score/reward.
 - get_grading_guidelines(): Return the rubric for the current problem.
 

envs/qed_math_env/server/qed_math_environment.py

@@ -520,8 +520,8 @@ class QEDMathEnvironment(MCPEnvironment):
     Math proof environment with MCP tools and rubric-based grading.
 
     This environment provides MCP tools for:
-    - get_problem(): Return current problem statement, reference solution,
-      and grading guidelines.
+    - get_problem(): Return current problem statement and grading guidelines.
+      The reference solution stays server-side and is only used for grading.
     - submit_proof(proof): Grade a proof via MathProofRubric and return
       score (0-7) and normalized reward.
     - get_grading_guidelines(): Return the rubric for the current problem.
@@ -829,8 +829,9 @@ def reset(
             ProblemObservation with problem details.
 
         Loads a problem from the dataset and returns a ProblemObservation
-        with problem, reference_solution, grading_guidelines, problem_id,
-        and dataset_source.
+        with problem, grading_guidelines, problem_id, and dataset_source.
+        reference_solution is left empty: the reference stays server-side
+        and is only used for grading.
         """
         selected_problem_id = kwargs.pop("problem_id", None)
         self._refresh_gold_cache_if_needed()
@@ -887,11 +888,9 @@ def reset(
 
         return ProblemObservation(
             problem=self._current_problem.get("problem", ""),
-            reference_solution=(
-                self._current_problem.get("reference_solution", "")
-                if self._current_problem.get("evaluation_mode", "proof") == "answer"
-                else ""
-            ),
+            # The reference solution is the grading key in both modes, so it is
+            # never sent to the agent. Grading reads it from the server-side problem.
+            reference_solution="",
             grading_guidelines=parse_schema(
                 self._current_problem.get("grading_guidelines", "") or ""
             ),
@@ -1084,15 +1083,10 @@ def get_problem_payload(self) -> dict:
             }
 
         problem_type = self._current_problem.get("problem_type", "proof")
-        evaluation_mode = self._current_problem.get("evaluation_mode")
-        if isinstance(evaluation_mode, str):
-            evaluation_mode = evaluation_mode.strip().lower()
-        else:
-            evaluation_mode = "answer" if problem_type == "answer" else "proof"
-        if evaluation_mode not in {"proof", "answer"}:
-            evaluation_mode = "proof"
 
-        payload = {
+        # The reference solution is the grading key in both modes, so it is
+        # not part of the payload the agent sees.
+        return {
             "problem": self._current_problem.get("problem", ""),
             "grading_guidelines": self._current_grading_guidelines_text(),
             "problem_id": self._current_problem.get("problem_id", ""),
@@ -1104,15 +1098,6 @@ def get_problem_payload(self) -> dict:
             "reward": 0.0,
         }
 
-        # Keep answer-mode parity with QED-Nano while preventing proof-mode
-        # agents from seeing the grading key via MCP.
-        if evaluation_mode == "answer":
-            payload["reference_solution"] = self._current_problem.get(
-                "reference_solution", ""
-            )
-
-        return payload
-
     @staticmethod
     def _verify_math(
         prediction: str,

envs/reasoning_gym_env/README.md

@@ -90,6 +90,8 @@ Each episode is **single-step**:
 2. `step(answer)` returns score and marks episode as done
 3. `reset()` without params gets next question from same dataset
 
+Only the first answer of an episode is scored. Further `step()` calls before the next `reset()` return the same final observation (`done=True`, same reward) and are not scored.
+
 ### Action
 
 **ReasoningGymAction**: Contains the agent's answer

envs/reasoning_gym_env/server/reasoning_gym_environment.py

@@ -70,6 +70,7 @@ def __init__(self):
         self._dataset = None
         self._dataset_iterator = None
         self._current_entry = None
+        self._final_observation: Optional[ReasoningGymObservation] = None
 
         # Dataset configuration tracking
         self._dataset_name: Optional[str] = None
@@ -192,6 +193,7 @@ def reset(
             # Iterator exhausted - restart it
             self._dataset_iterator = iter(self._dataset)
             self._current_entry = next(self._dataset_iterator)
+        self._final_observation = None
 
         question = self._current_entry["question"]
 
@@ -220,6 +222,12 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation:  # type:
         Returns:
             ReasoningGymObservation with score and correct answer, always done=True
         """
+        # The episode is single-step: once an answer has been scored (and the
+        # correct answer revealed), later steps return the same final
+        # observation and are not scored again.
+        if self._final_observation is not None:
+            return self._final_observation.model_copy(deep=True)
+
         self._state.step_count += 1
 
         # Validate current state
@@ -242,14 +250,16 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation:  # type:
         correct_answer = self._current_entry["answer"]
         metadata = self._current_entry.get("metadata", {})
 
-        return ReasoningGymObservation(
+        observation = ReasoningGymObservation(
             question=None,
             score=score,
             correct_answer=correct_answer,
             done=True,
             reward=reward,
             dataset_metadata=metadata if metadata else None,
         )
+        self._final_observation = observation.model_copy(deep=True)
+        return observation
 
     @property
     def state(self) -> State: