[RLlib] Fix small bug in 'InfiniteLookBackBuffer.get_state/from_state'.

rllib/env/utils/infinite_lookback_buffer.py

@@ -8,8 +8,10 @@
 from ray.rllib.utils.serialization import gym_space_from_dict, gym_space_to_dict
 from ray.rllib.utils.spaces.space_utils import (
     batch,
+    from_jsonable_if_needed,
     get_dummy_batch_for_space,
     get_base_struct_from_space,
+    to_jsonable_if_needed,
 )
 
 
@@ -71,12 +73,11 @@ def get_state(self) -> Dict[str, Any]:
             A dict containing all the data and metadata from the buffer.
         """
         return {
-            "data": self.data,
+            "data": to_jsonable_if_needed(self.data, self.space)
+            if self.space
+            else self.data,
             "lookback": self.lookback,
             "finalized": self.finalized,
-            "space_struct": gym_space_to_dict(self.space_struct)
-            if self.space_struct
-            else self.space_struct,
             "space": gym_space_to_dict(self.space) if self.space else self.space,
         }
 
@@ -92,16 +93,18 @@ def from_state(state: Dict[str, Any]) -> None:
             from the state dict.
         """
         buffer = InfiniteLookbackBuffer()
-        buffer.data = state["data"]
         buffer.lookback = state["lookback"]
         buffer.finalized = state["finalized"]
+        buffer.space = (
+            gym_space_from_dict(state["space"]) if state["space"] else None
+        )
         buffer.space_struct = (
-            gym_space_from_dict(state["space_struct"])
-            if state["space_struct"]
-            else state["space_struct"]
+            get_base_struct_from_space(buffer.space) if buffer.space else None
         )
-        buffer.space = (
-            gym_space_from_dict(state["space"]) if state["space"] else state["space"]
+        buffer.data = (
+            from_jsonable_if_needed(state["data"], buffer.space)
+            if buffer.space
+            else state["data"]
         )
 
         return buffer