diff --git a/docs/source/environments/finqa.md b/docs/source/environments/finqa.md index ba06ac728..e5e4fa14f 100644 --- a/docs/source/environments/finqa.md +++ b/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. diff --git a/docs/source/environments/qed_math.md b/docs/source/environments/qed_math.md index b3de606df..36d66f91a 100644 --- a/docs/source/environments/qed_math.md +++ b/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"` | diff --git a/docs/source/environments/reasoning_gym.md b/docs/source/environments/reasoning_gym.md index 58e38932f..53c9b6eea 100644 --- a/docs/source/environments/reasoning_gym.md +++ b/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 diff --git a/envs/finqa_env/README.md b/envs/finqa_env/README.md index 9f0681f9d..809f73806 100644 --- a/envs/finqa_env/README.md +++ b/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. diff --git a/envs/finqa_env/server/finqa_environment.py b/envs/finqa_env/server/finqa_environment.py index 513d54614..4fefeba91 100644 --- a/envs/finqa_env/server/finqa_environment.py +++ b/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,14 +279,16 @@ 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 @@ -282,13 +296,15 @@ def step( 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 diff --git a/envs/qed_math_env/README.md b/envs/qed_math_env/README.md index a9209d437..0f9d97b39 100644 --- a/envs/qed_math_env/README.md +++ b/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"` | diff --git a/envs/qed_math_env/models.py b/envs/qed_math_env/models.py index 56703a5b5..01344ec89 100644 --- a/envs/qed_math_env/models.py +++ b/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)" ) diff --git a/envs/qed_math_env/server/mcp_server.py b/envs/qed_math_env/server/mcp_server.py index 8bce1c997..4ac75e544 100644 --- a/envs/qed_math_env/server/mcp_server.py +++ b/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. diff --git a/envs/qed_math_env/server/qed_math_environment.py b/envs/qed_math_env/server/qed_math_environment.py index 8bf5d357f..4edf2429a 100644 --- a/envs/qed_math_env/server/qed_math_environment.py +++ b/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, diff --git a/envs/reasoning_gym_env/README.md b/envs/reasoning_gym_env/README.md index d39c1bee4..6720c5242 100644 --- a/envs/reasoning_gym_env/README.md +++ b/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 diff --git a/envs/reasoning_gym_env/server/reasoning_gym_environment.py b/envs/reasoning_gym_env/server/reasoning_gym_environment.py index d8f0ffdf1..3428cc550 100644 --- a/envs/reasoning_gym_env/server/reasoning_gym_environment.py +++ b/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,7 +250,7 @@ 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, @@ -250,6 +258,8 @@ def step(self, action: ReasoningGymAction) -> ReasoningGymObservation: # type: reward=reward, dataset_metadata=metadata if metadata else None, ) + self._final_observation = observation.model_copy(deep=True) + return observation @property def state(self) -> State: diff --git a/tests/envs/test_finqa_environment.py b/tests/envs/test_finqa_environment.py index 45fad019c..ab6f79b5e 100644 --- a/tests/envs/test_finqa_environment.py +++ b/tests/envs/test_finqa_environment.py @@ -653,6 +653,144 @@ def test_different_seeds_give_different_sequences(self, data_path): ) +class TestEpisodeEnd: + """Nothing is scored after the episode ends (answer submitted or max steps). + + Self-contained (synthetic data, no downloaded dataset needed) so it runs in CI. + """ + + @pytest.fixture + def data_path(self, tmp_path): + pytest.importorskip("pandas") + questions = tmp_path / "benchmark_questions" + questions.mkdir() + rows = ["id,user_query,company,question,answer"] + rows += [f"{i},query {i},acme,question {i},\\boxed{{{i}}}" for i in range(5)] + (questions / "finqa.csv").write_text("\n".join(rows) + "\n") + (tmp_path / "input_companies").mkdir() + return str(tmp_path) + + @staticmethod + def _make_env(data_path, max_steps=10): + from envs.finqa_env.server.finqa_environment import FinQAEnvironment + + env = FinQAEnvironment(data_path=data_path, max_steps=max_steps) + env.reset(seed=1) + return env + + @staticmethod + def _submit(answer): + from openenv.core.env_server.mcp_types import CallToolAction + + return CallToolAction(tool_name="submit_answer", arguments={"answer": answer}) + + @staticmethod + def _tool(): + from openenv.core.env_server.mcp_types import CallToolAction + + return CallToolAction( + tool_name="get_descriptions", arguments={"company_name": "acme"} + ) + + def test_submit_after_submit_is_not_scored(self, data_path): + env = self._make_env(data_path) + gold = env.state.ground_truth + + first = env.step(self._submit("-1")) + assert first.done is True + assert first.reward == 0.0 + + second = env.step(self._submit(gold)) + assert second.done is True + assert second.reward == 0.0 + assert second.metadata == first.metadata + + def test_submit_after_max_steps_is_not_scored(self, data_path): + env = self._make_env(data_path, max_steps=2) + gold = env.state.ground_truth + + assert env.step(self._tool()).done is False + last = env.step(self._tool()) + assert last.done is True + assert last.reward == 0.0 + + late = env.step(self._submit(gold)) + assert late.done is True + assert late.reward == 0.0 + assert env.state.step_count == 2 + + def test_submit_on_last_step_is_scored(self, data_path): + env = self._make_env(data_path, max_steps=2) + gold = env.state.ground_truth + + assert env.step(self._tool()).done is False + obs = env.step(self._submit(gold)) + assert obs.done is True + assert obs.reward == 1.0 + + def test_any_step_after_done_keeps_the_final_reward(self, data_path): + from openenv.core.env_server.mcp_types import ListToolsAction + + env = self._make_env(data_path, max_steps=2) + assert env.step(self._tool()).done is False + assert env.step(self._submit(env.state.ground_truth)).reward == 1.0 + + obs = env.step(ListToolsAction()) + assert obs.done is True + assert obs.reward == 1.0 + assert env.state.step_count == 2 + + def test_reset_starts_a_new_episode(self, data_path): + env = self._make_env(data_path) + env.step(self._submit("-1")) + + env.reset() + obs = env.step(self._submit(env.state.ground_truth)) + assert obs.done is True + assert obs.reward == 1.0 + + @pytest.mark.asyncio + async def test_step_async_submit_after_submit_is_not_scored(self, data_path): + env = self._make_env(data_path) + gold = env.state.ground_truth + + first = await env.step_async(self._submit("-1")) + assert first.done is True + assert first.reward == 0.0 + + second = await env.step_async(self._submit(gold)) + assert second.done is True + assert second.reward == 0.0 + + @pytest.mark.asyncio + async def test_step_async_submit_after_max_steps_is_not_scored(self, data_path): + env = self._make_env(data_path, max_steps=1) + gold = env.state.ground_truth + + last = await env.step_async(self._tool()) + assert last.done is True + + late = await env.step_async(self._submit(gold)) + assert late.done is True + assert late.reward == 0.0 + + @pytest.mark.asyncio + async def test_step_async_any_step_after_done_keeps_the_final_reward( + self, data_path + ): + from openenv.core.env_server.mcp_types import ListToolsAction + + env = self._make_env(data_path, max_steps=1) + assert ( + await env.step_async(self._submit(env.state.ground_truth)) + ).reward == 1.0 + + obs = await env.step_async(ListToolsAction()) + assert obs.done is True + assert obs.reward == 1.0 + assert env.state.step_count == 1 + + @pytest.mark.skipif(_integration_skip, reason=_integration_reason) class TestEnvironment: """Test environment logic using MCP actions.""" diff --git a/tests/envs/test_qed_math_environment.py b/tests/envs/test_qed_math_environment.py index 785fcb2e8..79db4ccb5 100644 --- a/tests/envs/test_qed_math_environment.py +++ b/tests/envs/test_qed_math_environment.py @@ -564,11 +564,32 @@ def test_after_reset_returns_problem(self): assert "problem_id" in payload assert "reference_solution" not in payload - def test_answer_mode_includes_reference_solution(self): + def test_answer_mode_hides_reference_solution(self): env = _make_env_with_problem(ANSWER_PROBLEM) - env.reset() + obs = env.reset() + assert obs.reference_solution == "" payload = env.get_problem_payload() - assert payload["reference_solution"] == r"\boxed{4}" + assert "reference_solution" not in payload + assert r"\boxed{4}" not in str(payload) + + def test_auto_detected_answer_mode_hides_reference_solution(self): + # A row without a rubric is routed to answer mode automatically. + env = _make_env_with_problem({"problem": "What is 17*23?", "answer": "391"}) + obs = env.reset() + assert obs.problem_type == "answer" + assert obs.reference_solution == "" + assert "391" not in str(env.get_problem_payload()) + + @pytest.mark.asyncio + async def test_answer_mode_correct_answer_still_scores(self): + env = _make_env_with_problem(ANSWER_PROBLEM) + try: + env.reset() + payload = await env.submit_proof_payload(r"The answer is \boxed{4}.") + assert payload["reward"] == pytest.approx(1.0) + assert payload["is_correct"] is True + finally: + await env._verifier_service.stop() def test_grading_guidelines_payload_without_reset(self): env = _make_env() diff --git a/tests/envs/test_reasoning_gym_environment.py b/tests/envs/test_reasoning_gym_environment.py index e0fe70f7a..4ca39c47d 100644 --- a/tests/envs/test_reasoning_gym_environment.py +++ b/tests/envs/test_reasoning_gym_environment.py @@ -196,6 +196,47 @@ def test_step_without_current_entry(self): assert obs.correct_answer is None assert obs.reward == 0.0 + def test_answer_after_done_is_not_scored(self): + """A second answer, e.g. the revealed correct answer, is not scored.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=1) + + first = env.step(ReasoningGymAction(answer="x")) + assert first.done is True + assert first.reward == 0.0 + assert first.correct_answer is not None + + second = env.step(ReasoningGymAction(answer=str(first.correct_answer))) + assert second.done is True + assert second.reward == 0.0 + assert second.score == first.score + + def test_extra_step_after_correct_answer_does_not_rescore(self): + """The correct answer scores 1.0, and an extra step keeps that reward.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=1) + answer = str(env._current_entry["answer"]) + + first = env.step(ReasoningGymAction(answer=answer)) + assert first.reward == 1.0 + + second = env.step(ReasoningGymAction(answer="x")) + assert second.done is True + assert second.reward == 1.0 + assert env.state.step_count == 1 + + def test_reset_allows_answering_the_next_question(self): + """reset() starts a new episode that is scored normally.""" + env = ReasoningGymEnvironment() + env.reset(dataset_name="leg_counting", dataset_config={}, seed=1, size=2) + env.step(ReasoningGymAction(answer="x")) + + env.reset() + answer = str(env._current_entry["answer"]) + obs = env.step(ReasoningGymAction(answer=answer)) + assert obs.done is True + assert obs.reward == 1.0 + def test_dataset_iterator_wraps_around(self): """Test that dataset iterator restarts when exhausted.""" env = ReasoningGymEnvironment()