New info API for vectorized environments #2657

gym/vector/async_vector_env.py

@@ -95,6 +95,9 @@ class AsyncVectorEnv(VectorEnv):
         space in Gym, such as :class:`~gym.spaces.Box`, :class:`~gym.spaces.Discrete`,
         or :class:`~gym.spaces.Dict`) and :obj:`shared_memory` is ``True``.
 
+    InvalidInfoFormat
+        If the info format does not matches any of the available.
+
     Example
     -------
 
@@ -308,8 +311,10 @@ def reset_wait(
         self._state = AsyncState.DEFAULT
 
         if return_info:
-            results, infos = zip(*results)
-            infos = list(infos)
+            infos = {}
+            results, info_data = zip(*results)
+            for i, info in enumerate(info_data):
+                infos = self._add_info(infos, info, i)
 
             if not self.shared_memory:
                 self.observations = concatenate(
@@ -406,10 +411,20 @@ def step_wait(self, timeout=None):
                 f"The call to `step_wait` has timed out after {timeout} second(s)."
             )
 
-        results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes])
+        observations_list, rewards, dones, infos = [], [], [], {}
+        successes = []
+        for i, pipe in enumerate(self.parent_pipes):
+            result, success = pipe.recv()
+            obs, rew, done, info = result
+            infos = self._add_info(infos, info, i)
+
+            successes.append(success)
+            observations_list.append(obs)
+            rewards.append(rew)
+            dones.append(done)
+
         self._raise_if_errors(successes)
         self._state = AsyncState.DEFAULT
-        observations_list, rewards, dones, infos = zip(*results)
 
         if not self.shared_memory:
             self.observations = concatenate(

gym/vector/sync_vector_env.py

@@ -36,6 +36,9 @@ class SyncVectorEnv(VectorEnv):
         :obj:`observation_space` (or, by default, the observation space of
         the first sub-environment).
 
+    InvalidInfoFormat
+        If the info format does not matches any of the available.
+
     Example
     -------
 
@@ -50,7 +53,13 @@ class SyncVectorEnv(VectorEnv):
                [-0.85009176,  0.5266346 ,  0.60007906]], dtype=float32)
     """
 
-    def __init__(self, env_fns, observation_space=None, action_space=None, copy=True):
+    def __init__(
+        self,
+        env_fns,
+        observation_space=None,
+        action_space=None,
+        copy=True,
+    ):
         self.env_fns = env_fns
         self.envs = [env_fn() for env_fn in env_fns]
         self.copy = copy
@@ -98,8 +107,8 @@ def reset_wait(
 
         self._dones[:] = False
         observations = []
-        data_list = []
-        for env, single_seed in zip(self.envs, seed):
+        infos = {}
+        for i, (env, single_seed) in enumerate(zip(self.envs, seed)):
 
             kwargs = {}
             if single_seed is not None:
@@ -115,7 +124,7 @@ def reset_wait(
             else:
                 observation, data = env.reset(**kwargs)
                 observations.append(observation)
-                data_list.append(data)
+                infos = self._add_info(infos, data, i)
 
         self.observations = concatenate(
             self.single_observation_space, observations, self.observations
@@ -125,20 +134,20 @@ def reset_wait(
         else:
             return (
                 deepcopy(self.observations) if self.copy else self.observations
-            ), data_list
+            ), infos
 
     def step_async(self, actions):
         self._actions = iterate(self.action_space, actions)
 
     def step_wait(self):
-        observations, infos = [], []
+        observations, infos = [], {}
         for i, (env, action) in enumerate(zip(self.envs, self._actions)):
             observation, self._rewards[i], self._dones[i], info = env.step(action)
             if self._dones[i]:
                 info["terminal_observation"] = observation
                 observation = env.reset()
             observations.append(observation)
-            infos.append(info)
+            infos = self._add_info(infos, info, i)
         self.observations = concatenate(
             self.single_observation_space, observations, self.observations
         )

gym/vector/vector_env.py

@@ -1,5 +1,7 @@
 from typing import List, Optional, Union
 
+import numpy as np
+
 import gym
 from gym.logger import deprecation
 from gym.vector.utils.spaces import batch_space
@@ -211,6 +213,27 @@ def seed(self, seed=None):
             "Please use `env.reset(seed=seed) instead in VectorEnvs."
         )
 
+    def _add_info(self, infos: dict, info: dict, env_num: int):
+        for k in info.keys():
+            if k not in infos:
+                info_array, array_mask = self._init_info_array(type(info[k]))
+            else:
+                info_array, array_mask = infos[k], infos[f"_{k}"]
+
+            info_array[env_num], array_mask[env_num] = info[k], True
+            infos[k], infos[f"_{k}"] = info_array, array_mask
+        return infos
+
+    def _init_info_array(self, dtype: type) -> np.ndarray:
+        if dtype not in [int, float, bool]:
+            dtype = object
+            array = np.zeros(self.num_envs, dtype=dtype)
+            array[:] = None
+        else:
+            array = np.zeros(self.num_envs, dtype=dtype)
+        array_mask = np.zeros(self.num_envs, dtype=bool)
+        return array, array_mask
+
     def __del__(self):
         if not getattr(self, "closed", True):
             self.close()

gym/wrappers/__init__.py

@@ -1,4 +1,5 @@
 """Module of wrapper classes."""
+from gym import error
 from gym.wrappers.atari_preprocessing import AtariPreprocessing
 from gym.wrappers.autoreset import AutoResetWrapper
 from gym.wrappers.clip_action import ClipAction
@@ -8,7 +9,6 @@
 from gym.wrappers.gray_scale_observation import GrayScaleObservation
 from gym.wrappers.normalize import NormalizeObservation, NormalizeReward
 from gym.wrappers.order_enforcing import OrderEnforcing
-from gym.wrappers.pixel_observation import PixelObservationWrapper
 from gym.wrappers.record_episode_statistics import RecordEpisodeStatistics
 from gym.wrappers.record_video import RecordVideo, capped_cubic_video_schedule
 from gym.wrappers.rescale_action import RescaleAction
@@ -17,3 +17,4 @@
 from gym.wrappers.time_limit import TimeLimit
 from gym.wrappers.transform_observation import TransformObservation
 from gym.wrappers.transform_reward import TransformReward
+from gym.wrappers.vec_info_to_classic import ClassicVectorInfo

gym/wrappers/record_episode_statistics.py

@@ -7,6 +7,89 @@
 import gym
 
 
+class ClassicStatsInfo:
+    """Manage episode statistics."""
+
+    def __init__(self, num_envs: int):
+        """Classic EpisodeStatics info.
+
+        Args:
+            num_envs (int): number of environments.
+        """
+        self.info = {}
+
+    def add_info(self, infos: dict, env_num: int):
+        """Add info.
+
+        Args:
+            infos (dict): info dict of the environment.
+            env_num (int): environment number.
+        """
+        self.info = {**self.info, **infos}
+
+    def add_episode_statistics(self, infos: dict, env_num: int):
+        """Add episode statistics.
+
+        Args:
+            infos (dict): info dict of the environment.
+            env_num (int): env number.
+        """
+        self.info = {**self.info, **infos}
+
+    def get_info(self):
+        """Return info."""
+        return self.info
+
+
+class BraxVecEnvStatsInfo:
+    """Manage episode statistics in the Brax format for vectorized envs."""
+
+    def __init__(self, num_envs: int):
+        """Brax-style episode statistics.
+
+        Args:
+            num_envs (int): number of environments.
+        """
+        self.num_envs = num_envs
+        self.info = {}
+
+    def add_info(self, info: dict, env_num: int):
+        """Add info.
+
+        Args:
+            info (dict): info dict of the environment.
+            env_num (int): environment number.
+        """
+        self.info = {**self.info, **info}
+
+    def add_episode_statistics(self, info: dict, env_num: int):
+        """Add episode statistics.
+
+        Add statistics coming from the vectorized environment.
+
+        Args:
+            info (dict): info dict of the environment.
+            env_num (int): env number of the vectorized environments.
+        """
+        episode_info = info["episode"]
+
+        self.info["episode"] = self.info.get("episode", {})
+
+        self.info["_episode"] = self.info.get(
+            "_episode", np.zeros(self.num_envs, dtype=bool)
+        )
+        self.info["_episode"][env_num] = True
+
+        for k in episode_info.keys():
+            info_array = self.info["episode"].get(k, np.zeros(self.num_envs))
+            info_array[env_num] = episode_info[k]
+            self.info["episode"][k] = info_array
+
+    def get_info(self):
+        """Returns info."""
+        return self.info
+
+
 class RecordEpisodeStatistics(gym.Wrapper):
     """This wrapper will keep track of cumulative rewards and episode lengths.
 
@@ -46,6 +129,10 @@ def __init__(self, env: gym.Env, deque_size: int = 100):
         self.return_queue = deque(maxlen=deque_size)
         self.length_queue = deque(maxlen=deque_size)
         self.is_vector_env = getattr(env, "is_vector_env", False)
+        if self.is_vector_env:
+            self.stats_info_processor = BraxVecEnvStatsInfo
+        else:
+            self.stats_info_processor = ClassicStatsInfo
 
     def reset(self, **kwargs):
         """Resets the environment using kwargs and resets the episode returns and lengths."""
@@ -56,35 +143,35 @@ def reset(self, **kwargs):
 
     def step(self, action):
         """Steps through the environment, recording the episode statistics."""
+        infos_processor = self.stats_info_processor(self.num_envs)
         observations, rewards, dones, infos = super().step(action)
         self.episode_returns += rewards
         self.episode_lengths += 1
         if not self.is_vector_env:
-            infos = [infos]
             dones = [dones]
-        else:
-            infos = list(infos)  # Convert infos to mutable type
+        dones = list(dones)
+
         for i in range(len(dones)):
             if dones[i]:
-                infos[i] = infos[i].copy()
+                infos_processor.add_info(infos, i)
                 episode_return = self.episode_returns[i]
                 episode_length = self.episode_lengths[i]
                 episode_info = {
-                    "r": episode_return,
-                    "l": episode_length,
-                    "t": round(time.perf_counter() - self.t0, 6),
+                    "episode": {
+                        "r": episode_return,
+                        "l": episode_length,
+                        "t": round(time.perf_counter() - self.t0, 6),
+                    }
                 }
-                infos[i]["episode"] = episode_info
+                infos_processor.add_episode_statistics(episode_info, i)
                 self.return_queue.append(episode_return)
                 self.length_queue.append(episode_length)
                 self.episode_count += 1
                 self.episode_returns[i] = 0
                 self.episode_lengths[i] = 0
-        if self.is_vector_env:
-            infos = tuple(infos)
         return (
             observations,
             rewards,
             dones if self.is_vector_env else dones[0],
-            infos if self.is_vector_env else infos[0],
+            infos_processor.get_info(),
         )