| import numpy as np |
| from copy import deepcopy |
|
|
| from gym import logger |
| from gym.vector.vector_env import VectorEnv |
| from gym.vector.utils import concatenate, create_empty_array |
|
|
| __all__ = ["SyncVectorEnv"] |
|
|
|
|
| class SyncVectorEnv(VectorEnv): |
| """Vectorized environment that serially runs multiple environments. |
| Parameters |
| ---------- |
| env_fns : iterable of callable |
| Functions that create the environments. |
| observation_space : `gym.spaces.Space` instance, optional |
| Observation space of a single environment. If `None`, then the |
| observation space of the first environment is taken. |
| action_space : `gym.spaces.Space` instance, optional |
| Action space of a single environment. If `None`, then the action space |
| of the first environment is taken. |
| copy : bool (default: `True`) |
| If `True`, then the `reset` and `step` methods return a copy of the |
| observations. |
| """ |
|
|
| 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 |
| self.metadata = self.envs[0].metadata |
|
|
| if (observation_space is None) or (action_space is None): |
| observation_space = observation_space or self.envs[0].observation_space |
| action_space = action_space or self.envs[0].action_space |
| super(SyncVectorEnv, self).__init__( |
| num_envs=len(env_fns), |
| observation_space=observation_space, |
| action_space=action_space, |
| ) |
|
|
| self._check_observation_spaces() |
| self.observations = create_empty_array( |
| self.single_observation_space, n=self.num_envs, fn=np.zeros |
| ) |
| self._rewards = np.zeros((self.num_envs,), dtype=np.float64) |
| self._dones = np.zeros((self.num_envs,), dtype=np.bool_) |
| |
| |
| self._actions = None |
|
|
| def seed(self, seeds=None): |
| if seeds is None: |
| seeds = [None for _ in range(self.num_envs)] |
| if isinstance(seeds, int): |
| seeds = [seeds + i for i in range(self.num_envs)] |
| assert len(seeds) == self.num_envs |
|
|
| for env, seed in zip(self.envs, seeds): |
| env.seed(seed) |
|
|
| def reset_wait(self): |
| self._dones[:] = False |
| observations = [] |
| for env in self.envs: |
| observation = env.reset() |
| observations.append(observation) |
| self.observations = concatenate( |
| observations, self.observations, self.single_observation_space |
| ) |
|
|
| return deepcopy(self.observations) if self.copy else self.observations |
|
|
| def step_async(self, actions): |
| self._actions = actions |
|
|
| def step_wait(self): |
| observations, infos = [], [] |
| for i, (env, action) in enumerate(zip(self.envs, self._actions)): |
| observation, self._rewards[i], self._dones[i], info = env.step(action) |
| |
| |
| observations.append(observation) |
| infos.append(info) |
| self.observations = concatenate( |
| observations, self.observations, self.single_observation_space |
| ) |
|
|
| return ( |
| deepcopy(self.observations) if self.copy else self.observations, |
| np.copy(self._rewards), |
| np.copy(self._dones), |
| infos, |
| ) |
|
|
| def close_extras(self, **kwargs): |
| [env.close() for env in self.envs] |
|
|
| def _check_observation_spaces(self): |
| for env in self.envs: |
| if not (env.observation_space == self.single_observation_space): |
| break |
| else: |
| return True |
| raise RuntimeError( |
| "Some environments have an observation space " |
| "different from `{0}`. In order to batch observations, the " |
| "observation spaces from all environments must be " |
| "equal.".format(self.single_observation_space) |
| ) |
| |
| def call(self, name, *args, **kwargs) -> tuple: |
| """Calls the method with name and applies args and kwargs. |
| |
| Args: |
| name: The method name |
| *args: The method args |
| **kwargs: The method kwargs |
| |
| Returns: |
| Tuple of results |
| """ |
| results = [] |
| for env in self.envs: |
| function = getattr(env, name) |
| if callable(function): |
| results.append(function(*args, **kwargs)) |
| else: |
| results.append(function) |
|
|
| return tuple(results) |
|
|
| def call_each(self, name: str, |
| args_list: list=None, |
| kwargs_list: list=None): |
| n_envs = len(self.envs) |
| if args_list is None: |
| args_list = [[]] * n_envs |
| assert len(args_list) == n_envs |
|
|
| if kwargs_list is None: |
| kwargs_list = [dict()] * n_envs |
| assert len(kwargs_list) == n_envs |
|
|
| results = [] |
| for i, env in enumerate(self.envs): |
| function = getattr(env, name) |
| if callable(function): |
| results.append(function(*args_list[i], **kwargs_list[i])) |
| else: |
| results.append(function) |
|
|
| return tuple(results) |
|
|
|
|
| def render(self, *args, **kwargs): |
| return self.call('render', *args, **kwargs) |
| |
| def set_attr(self, name: str, values): |
| """Sets an attribute of the sub-environments. |
| |
| Args: |
| name: The property name to change |
| values: Values of the property to be set to. If ``values`` is a list or |
| tuple, then it corresponds to the values for each individual |
| environment, otherwise, a single value is set for all environments. |
| |
| Raises: |
| ValueError: Values must be a list or tuple with length equal to the number of environments. |
| """ |
| if not isinstance(values, (list, tuple)): |
| values = [values for _ in range(self.num_envs)] |
| if len(values) != self.num_envs: |
| raise ValueError( |
| "Values must be a list or tuple with length equal to the " |
| f"number of environments. Got `{len(values)}` values for " |
| f"{self.num_envs} environments." |
| ) |
|
|
| for env, value in zip(self.envs, values): |
| setattr(env, name, value) |
