diff --git a/envs/chess_env/client.py b/envs/chess_env/client.py index e2fd6f4e6..c247dd58a 100644 --- a/envs/chess_env/client.py +++ b/envs/chess_env/client.py @@ -65,13 +65,18 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation] StepResult with ChessObservation. """ obs_data = payload.get("observation", {}) + # `serialize_observation` keeps reward and done on the envelope, not in + # the observation dict. Reading them from obs_data left every step + # reporting reward 0.0 and done False. + reward = payload.get("reward", 0.0) + done = payload.get("done", False) observation = ChessObservation( fen=obs_data.get("fen", ""), legal_moves=obs_data.get("legal_moves", []), is_check=obs_data.get("is_check", False), - done=obs_data.get("done", False), - reward=obs_data.get("reward", 0.0), + done=bool(done), + reward=reward, result=obs_data.get("result"), metadata=payload.get("metadata", obs_data.get("metadata", {})), ) diff --git a/envs/sophistry_bench_sprint_env/client.py b/envs/sophistry_bench_sprint_env/client.py index 070a7528a..0b16b1672 100644 --- a/envs/sophistry_bench_sprint_env/client.py +++ b/envs/sophistry_bench_sprint_env/client.py @@ -53,13 +53,17 @@ def _parse_result(self, data: dict) -> StepResult[AdvocacyObservation]: error = obs_data.get("error") or "" if error and "error" not in metadata: metadata["error"] = error + # ``serialize_observation`` keeps reward and done on the envelope, not in + # the observation dict, so set them on the observation too. Without this + # ``result.observation.reward`` stayed None while ``result.reward`` was + # correct, which traps callers that read the observation. + reward = data.get("reward") + done = bool(data.get("done", False)) # Construct once with metadata set, rather than mutating the model after. - observation = AdvocacyObservation(**obs_data, metadata=metadata) - return StepResult( - observation=observation, - reward=data.get("reward"), - done=data.get("done", False), + observation = AdvocacyObservation( + **obs_data, metadata=metadata, reward=reward, done=done ) + return StepResult(observation=observation, reward=reward, done=done) def _parse_state(self, data: dict) -> State: return State(**data) diff --git a/envs/sumo_rl_env/client.py b/envs/sumo_rl_env/client.py index 1ebbcd7ef..1868097e3 100644 --- a/envs/sumo_rl_env/client.py +++ b/envs/sumo_rl_env/client.py @@ -104,8 +104,10 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SumoObservation]: observation_shape=obs_data.get("observation_shape", []), action_mask=obs_data.get("action_mask", []), sim_time=obs_data.get("sim_time", 0.0), - done=obs_data.get("done", False), - reward=obs_data.get("reward"), + # Read from the envelope, where `serialize_observation` puts these, + # so `result.observation.reward` agrees with `result.reward`. + done=bool(payload.get("done", False)), + reward=payload.get("reward"), metadata=payload.get("metadata", obs_data.get("metadata", {})), ) diff --git a/tests/envs/test_client_step_result_contract.py b/tests/envs/test_client_step_result_contract.py new file mode 100644 index 000000000..789c5973d --- /dev/null +++ b/tests/envs/test_client_step_result_contract.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Typed clients must read reward and done from where the server puts them. + +`serialize_observation` deliberately excludes `reward` and `done` from the +nested observation dict and surfaces them on the envelope. A client that reads +them from the observation dict instead sees reward 0.0 and done False on every +step, which silently breaks any training loop or episode-termination check +built on it. + +These tests drive each client's `_parse_result` with the output of the real +serializer, so they fail if a client and the wire format ever drift apart +again. +""" + +from __future__ import annotations + +from openenv.core.env_server.serialization import serialize_observation + +# No `importorskip` here on purpose. These clients and their models depend only +# on openenv and pydantic, never on an engine package, so skipping when the +# engine is absent would quietly skip the very regression being guarded. + + +def test_serializer_keeps_reward_and_done_on_the_envelope(): + """Pin the wire contract the clients are being checked against.""" + from envs.chess_env.models import ChessObservation + + payload = serialize_observation(ChessObservation(done=True, reward=1.0)) + + assert payload["reward"] == 1.0 + assert payload["done"] is True + assert "reward" not in payload["observation"] + assert "done" not in payload["observation"] + + +class TestChessClient: + def parse(self, observation): + from envs.chess_env.client import ChessEnv + + return ChessEnv._parse_result(None, serialize_observation(observation)) + + def test_step_result_carries_reward_and_done(self): + from envs.chess_env.models import ChessObservation + + result = self.parse( + ChessObservation(fen="8/8/8/8/8/8/8/K6k w - - 0 1", done=True, reward=1.0) + ) + + assert result.reward == 1.0 + assert result.done is True + + def test_observation_agrees_with_the_step_result(self): + from envs.chess_env.models import ChessObservation + + result = self.parse(ChessObservation(done=True, reward=-1.0)) + + assert result.observation.reward == result.reward + assert result.observation.done == result.done + + def test_other_observation_fields_still_arrive(self): + from envs.chess_env.models import ChessObservation + + result = self.parse( + ChessObservation( + fen="rn6/8/8/8/8/8/8/K6k b - - 0 1", is_check=True, result="1-0" + ) + ) + + assert result.observation.fen.startswith("rn6") + assert result.observation.is_check is True + assert result.observation.result == "1-0" + + +class TestSumoClient: + def test_observation_agrees_with_the_step_result(self): + from envs.sumo_rl_env.client import SumoRLEnv + from envs.sumo_rl_env.models import SumoObservation + + result = SumoRLEnv._parse_result( + None, serialize_observation(SumoObservation(done=True, reward=2.5)) + ) + + assert result.reward == 2.5 + assert result.done is True + assert result.observation.reward == 2.5 + assert result.observation.done is True + + +class TestSophistryClient: + def test_observation_agrees_with_the_step_result(self): + from envs.sophistry_bench_sprint_env.client import SophistryBenchSprintEnv + from envs.sophistry_bench_sprint_env.models import AdvocacyObservation + + result = SophistryBenchSprintEnv._parse_result( + None, serialize_observation(AdvocacyObservation(done=True, reward=0.75)) + ) + + assert result.reward == 0.75 + assert result.done is True + assert result.observation.reward == 0.75 + assert result.observation.done is True