fix: surface top-level metadata on StepResult in typed env clients

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:

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:

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(

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:

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:

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:

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:

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:

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:

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:

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:

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:

envs/reasoning_gym_env/client.py

@@ -93,14 +93,15 @@ 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"),
         )
 
         return StepResult(
             observation=observation,
             reward=payload.get("reward"),
             done=payload.get("done", False),
+            metadata=payload.get("metadata"),
         )
 
     def _parse_state(self, payload: Dict) -> State:

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:

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:

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:

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:

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:

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:

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: