diff --git a/rllib/env/utils/infinite_lookback_buffer.py b/rllib/env/utils/infinite_lookback_buffer.py index 269e3827ca20..cd84a8518097 100644 --- a/rllib/env/utils/infinite_lookback_buffer.py +++ b/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,16 @@ 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