Add method to GymWrapper to return info from gym.Env.
OpenAI's Gym environments return diagnostic information
for every environment step.
Add functionality to track and obtain the information returned
from stepping the environment.
"""Wraps an OpenAI Gym environment to be used as a dm_env environment."""
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
from acme import specs
from acme import types
@@ -37,6 +37,7 @@ class GymWrapper(dm_env.Environment):
self._environment = environment
self._reset_next_step = True
+ self._last_info = None
# Convert action and observation specs.
obs_space = self._environment.observation_space
@@ -48,6 +49,8 @@ class GymWrapper(dm_env.Environment):
"""Resets the episode."""
self._reset_next_step = False
observation = self._environment.reset()
+ # Reset the diagnostic information.
+ self._last_info = None
return dm_env.restart(observation)
def step(self, action: types.NestedArray) -> dm_env.TimeStep:
@@ -57,6 +60,7 @@ class GymWrapper(dm_env.Environment):
observation, reward, done, info = self._environment.step(action)
self._reset_next_step = done
+ self._last_info = info
reward = tree.map_structure(lambda x, t: np.asarray(x, dtype=t.dtype),
reward, self.reward_spec())
@@ -74,6 +78,14 @@ class GymWrapper(dm_env.Environment):
def action_spec(self) -> types.NestedSpec:
return self._action_spec
+ def get_info(self) -> Optional[Dict[str, Any]]:
+ """Returns the last info returned from env.step(action).
+
+ Returns:
+ info: dictionary of diagnostic information from the last environment step
+ """
+ return self._last_info
+
@property
def environment(self) -> gym.Env:
"""Returns the wrapped environment."""
