diff --git a/envs/atari_env/client.py b/envs/atari_env/client.py index a4808692c..4e3c7157f 100644 --- a/envs/atari_env/client.py +++ b/envs/atari_env/client.py @@ -88,13 +88,14 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[AtariObservation] frame_number=obs_data.get("frame_number", 0), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> AtariState: diff --git a/envs/browsergym_env/client.py b/envs/browsergym_env/client.py index 1c1196ff1..23709c2dd 100644 --- a/envs/browsergym_env/client.py +++ b/envs/browsergym_env/client.py @@ -99,13 +99,14 @@ def _parse_result( last_action_error=obs_data.get("last_action_error", False), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> BrowserGymState: diff --git a/envs/calendar_env/client.py b/envs/calendar_env/client.py index d66199293..fc18c2105 100644 --- a/envs/calendar_env/client.py +++ b/envs/calendar_env/client.py @@ -140,12 +140,13 @@ def _parse_step_result( tool_result=obs_data.get("tool_result"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def reset( diff --git a/envs/chat_env/client.py b/envs/chat_env/client.py index eb4a3de9d..c9c73459a 100644 --- a/envs/chat_env/client.py +++ b/envs/chat_env/client.py @@ -89,13 +89,14 @@ def _parse_result(self, payload: Dict) -> StepResult[ChatObservation]: tokens=tokens, done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> ChatState: diff --git a/envs/chess_env/client.py b/envs/chess_env/client.py index 0bc577e71..e2fd6f4e6 100644 --- a/envs/chess_env/client.py +++ b/envs/chess_env/client.py @@ -73,13 +73,14 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[ChessObservation] done=obs_data.get("done", False), reward=obs_data.get("reward", 0.0), result=obs_data.get("result"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=observation.reward, done=observation.done, + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> ChessState: diff --git a/envs/connect4_env/client.py b/envs/connect4_env/client.py index 46b614e19..e20cd688e 100644 --- a/envs/connect4_env/client.py +++ b/envs/connect4_env/client.py @@ -71,13 +71,14 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Connect4Observati legal_actions=obs_data.get("legal_actions", []), done=payload.get("done", False), reward=payload.get("reward", 0.0), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward", 0.0), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> Connect4State: diff --git a/envs/dm_control_env/client.py b/envs/dm_control_env/client.py index cc77d9fe0..d0c927a92 100644 --- a/envs/dm_control_env/client.py +++ b/envs/dm_control_env/client.py @@ -144,13 +144,14 @@ def _parse_result(self, payload: Dict) -> StepResult[DMControlObservation]: pixels=obs_data.get("pixels"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> DMControlState: diff --git a/envs/finrl_env/client.py b/envs/finrl_env/client.py index 6b0e27075..8c9abc1ae 100644 --- a/envs/finrl_env/client.py +++ b/envs/finrl_env/client.py @@ -109,13 +109,14 @@ def _parse_result(self, payload: Dict) -> StepResult[FinRLObservation]: date=obs_data.get("date", ""), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> State: diff --git a/envs/kernrl/client.py b/envs/kernrl/client.py index 7eceaf755..884811442 100644 --- a/envs/kernrl/client.py +++ b/envs/kernrl/client.py @@ -121,13 +121,14 @@ def _parse_result(self, payload: Dict) -> StepResult[KernelObservation]: speedup=obs_data.get("speedup"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> KernelState: diff --git a/envs/maze_env/client.py b/envs/maze_env/client.py index 914d1a67a..ceacad648 100644 --- a/envs/maze_env/client.py +++ b/envs/maze_env/client.py @@ -70,13 +70,14 @@ def _parse_result(self, payload: Dict) -> StepResult[MazeObservation]: previous_position=obs_data.get("previous_position", []), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> MazeState: diff --git a/envs/openapp_env/client.py b/envs/openapp_env/client.py index eb4acdaeb..eef996268 100644 --- a/envs/openapp_env/client.py +++ b/envs/openapp_env/client.py @@ -123,13 +123,14 @@ def _parse_result(self, payload: Dict) -> StepResult[OpenAppObservation]: last_action_error=obs_data.get("last_action_error"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> State: diff --git a/envs/openspiel_env/client.py b/envs/openspiel_env/client.py index acf190187..f9fdc26e2 100644 --- a/envs/openspiel_env/client.py +++ b/envs/openspiel_env/client.py @@ -88,13 +88,14 @@ def _parse_result( opponent_last_action=obs_data.get("opponent_last_action"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> OpenSpielState: diff --git a/envs/reasoning_gym_env/client.py b/envs/reasoning_gym_env/client.py index 5ab845969..abf7ed6ea 100644 --- a/envs/reasoning_gym_env/client.py +++ b/envs/reasoning_gym_env/client.py @@ -93,7 +93,7 @@ def _parse_result(self, payload: Dict) -> StepResult[ReasoningGymObservation]: correct_answer=obs_data.get("correct_answer"), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), dataset_metadata=obs_data.get("dataset_metadata"), ) @@ -101,6 +101,7 @@ def _parse_result(self, payload: Dict) -> StepResult[ReasoningGymObservation]: observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> State: diff --git a/envs/repl_env/client.py b/envs/repl_env/client.py index a94792613..a24e6b589 100644 --- a/envs/repl_env/client.py +++ b/envs/repl_env/client.py @@ -76,13 +76,14 @@ def _parse_result(self, payload: dict[str, Any]) -> StepResult[REPLObservation]: max_iterations=obs_data.get("max_iterations", 30), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: dict[str, Any]) -> REPLState: diff --git a/envs/snake_env/client.py b/envs/snake_env/client.py index 8d02314b8..63334f69c 100644 --- a/envs/snake_env/client.py +++ b/envs/snake_env/client.py @@ -90,13 +90,14 @@ def _parse_result(self, payload: Dict) -> StepResult[SnakeObservation]: alive=obs_data.get("alive", True), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> State: diff --git a/envs/sumo_rl_env/client.py b/envs/sumo_rl_env/client.py index cf4e4345c..1ebbcd7ef 100644 --- a/envs/sumo_rl_env/client.py +++ b/envs/sumo_rl_env/client.py @@ -106,13 +106,14 @@ def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SumoObservation]: sim_time=obs_data.get("sim_time", 0.0), done=obs_data.get("done", False), reward=obs_data.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> SumoState: diff --git a/envs/tbench2_env/client.py b/envs/tbench2_env/client.py index 01290ee28..2392405a5 100644 --- a/envs/tbench2_env/client.py +++ b/envs/tbench2_env/client.py @@ -54,12 +54,13 @@ def _parse_result(self, payload: dict[str, Any]) -> StepResult[Tbench2Observatio info=obs_data.get("info", {}), reward=payload.get("reward"), done=payload.get("done", False), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: dict[str, Any]) -> Tbench2State: diff --git a/envs/textarena_env/client.py b/envs/textarena_env/client.py index 3d18743fc..d943871a6 100644 --- a/envs/textarena_env/client.py +++ b/envs/textarena_env/client.py @@ -95,12 +95,13 @@ def _parse_result(self, payload: Dict) -> StepResult[TextArenaObservation]: info=obs_data.get("info", {}), reward=payload.get("reward"), done=payload.get("done", False), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict[str, Any]) -> TextArenaState: diff --git a/envs/unity_env/client.py b/envs/unity_env/client.py index df66fe7fa..6216acdf8 100644 --- a/envs/unity_env/client.py +++ b/envs/unity_env/client.py @@ -188,13 +188,14 @@ def _parse_result(self, payload: Dict) -> StepResult[UnityObservation]: observation_spec_info=obs_data.get("observation_spec_info", {}), done=payload.get("done", False), reward=payload.get("reward"), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> UnityState: diff --git a/envs/websearch_env/client.py b/envs/websearch_env/client.py index 15c89533c..635a7edc1 100644 --- a/envs/websearch_env/client.py +++ b/envs/websearch_env/client.py @@ -83,13 +83,14 @@ def _parse_result(self, payload: Dict) -> StepResult[WebSearchObservation]: observation = WebSearchObservation( content=obs_data.get("content", ""), web_contents=obs_data.get("web_contents", []), - metadata=obs_data.get("metadata", {}), + metadata=payload.get("metadata", obs_data.get("metadata", {})), ) return StepResult( observation=observation, reward=payload.get("reward"), done=payload.get("done", False), + metadata=payload.get("metadata"), ) def _parse_state(self, payload: Dict) -> State: