mirror of
https://github.com/saymrwulf/stable-baselines3.git
synced 2026-05-18 21:30:19 +00:00
98 lines
4 KiB
Python
98 lines
4 KiB
Python
from collections import OrderedDict
|
|
from copy import deepcopy
|
|
|
|
import numpy as np
|
|
|
|
from torchy_baselines.common.vec_env.base_vec_env import VecEnv
|
|
from torchy_baselines.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
|
|
|
|
|
|
class DummyVecEnv(VecEnv):
|
|
"""
|
|
Creates a simple vectorized wrapper for multiple environments, calling each environment in sequence on the current
|
|
Python process. This is useful for computationally simple environment such as ``cartpole-v1``, as the overhead of
|
|
multiprocess or multithread outweighs the environment computation time. This can also be used for RL methods that
|
|
require a vectorized environment, but that you want a single environments to train with.
|
|
|
|
:param env_fns: ([Gym Environment]) the list of environments to vectorize
|
|
"""
|
|
|
|
def __init__(self, env_fns):
|
|
self.envs = [fn() for fn in env_fns]
|
|
env = self.envs[0]
|
|
VecEnv.__init__(self, len(env_fns), env.observation_space, env.action_space)
|
|
obs_space = env.observation_space
|
|
self.keys, shapes, dtypes = obs_space_info(obs_space)
|
|
|
|
self.buf_obs = OrderedDict([
|
|
(k, np.zeros((self.num_envs,) + tuple(shapes[k]), dtype=dtypes[k]))
|
|
for k in self.keys])
|
|
self.buf_dones = np.zeros((self.num_envs,), dtype=np.bool)
|
|
self.buf_rews = np.zeros((self.num_envs,), dtype=np.float32)
|
|
self.buf_infos = [{} for _ in range(self.num_envs)]
|
|
self.actions = None
|
|
self.metadata = env.metadata
|
|
|
|
def step_async(self, actions):
|
|
self.actions = actions
|
|
|
|
def step_wait(self):
|
|
for env_idx in range(self.num_envs):
|
|
obs, self.buf_rews[env_idx], self.buf_dones[env_idx], self.buf_infos[env_idx] =\
|
|
self.envs[env_idx].step(self.actions[env_idx])
|
|
if self.buf_dones[env_idx]:
|
|
# save final observation where user can get it, then reset
|
|
self.buf_infos[env_idx]['terminal_observation'] = obs
|
|
obs = self.envs[env_idx].reset()
|
|
self._save_obs(env_idx, obs)
|
|
return (self._obs_from_buf(), np.copy(self.buf_rews), np.copy(self.buf_dones),
|
|
deepcopy(self.buf_infos))
|
|
|
|
def reset(self):
|
|
for env_idx in range(self.num_envs):
|
|
obs = self.envs[env_idx].reset()
|
|
self._save_obs(env_idx, obs)
|
|
return self._obs_from_buf()
|
|
|
|
def close(self):
|
|
for env in self.envs:
|
|
env.close()
|
|
|
|
def get_images(self):
|
|
return [env.render(mode='rgb_array') for env in self.envs]
|
|
|
|
def render(self, *args, **kwargs):
|
|
if self.num_envs == 1:
|
|
return self.envs[0].render(*args, **kwargs)
|
|
else:
|
|
return super().render(*args, **kwargs)
|
|
|
|
def _save_obs(self, env_idx, obs):
|
|
for key in self.keys:
|
|
if key is None:
|
|
self.buf_obs[key][env_idx] = obs
|
|
else:
|
|
self.buf_obs[key][env_idx] = obs[key]
|
|
|
|
def _obs_from_buf(self):
|
|
return dict_to_obs(self.observation_space, copy_obs_dict(self.buf_obs))
|
|
|
|
def get_attr(self, attr_name, indices=None):
|
|
"""Return attribute from vectorized environment (see base class)."""
|
|
target_envs = self._get_target_envs(indices)
|
|
return [getattr(env_i, attr_name) for env_i in target_envs]
|
|
|
|
def set_attr(self, attr_name, value, indices=None):
|
|
"""Set attribute inside vectorized environments (see base class)."""
|
|
target_envs = self._get_target_envs(indices)
|
|
for env_i in target_envs:
|
|
setattr(env_i, attr_name, value)
|
|
|
|
def env_method(self, method_name, *method_args, indices=None, **method_kwargs):
|
|
"""Call instance methods of vectorized environments."""
|
|
target_envs = self._get_target_envs(indices)
|
|
return [getattr(env_i, method_name)(*method_args, **method_kwargs) for env_i in target_envs]
|
|
|
|
def _get_target_envs(self, indices):
|
|
indices = self._get_indices(indices)
|
|
return [self.envs[i] for i in indices]
|