mirror of
https://github.com/microsoft/qlib.git
synced 2026-06-30 01:21:18 +08:00
696 lines
24 KiB
Python
696 lines
24 KiB
Python
import gym
|
||
import time
|
||
import ctypes
|
||
import numpy as np
|
||
from collections import OrderedDict
|
||
from multiprocessing.context import Process
|
||
from multiprocessing import Array, Pipe, connection, Queue
|
||
from typing import Any, List, Tuple, Union, Callable, Optional
|
||
|
||
from tianshou.env.worker import EnvWorker
|
||
from tianshou.env.utils import CloudpickleWrapper
|
||
|
||
|
||
_NP_TO_CT = {
|
||
np.bool: ctypes.c_bool,
|
||
np.bool_: ctypes.c_bool,
|
||
np.uint8: ctypes.c_uint8,
|
||
np.uint16: ctypes.c_uint16,
|
||
np.uint32: ctypes.c_uint32,
|
||
np.uint64: ctypes.c_uint64,
|
||
np.int8: ctypes.c_int8,
|
||
np.int16: ctypes.c_int16,
|
||
np.int32: ctypes.c_int32,
|
||
np.int64: ctypes.c_int64,
|
||
np.float32: ctypes.c_float,
|
||
np.float64: ctypes.c_double,
|
||
}
|
||
|
||
|
||
class ShArray:
|
||
"""Wrapper of multiprocessing Array."""
|
||
|
||
def __init__(self, dtype: np.generic, shape: Tuple[int]) -> None:
|
||
self.arr = Array(
|
||
_NP_TO_CT[dtype.type], # type: ignore
|
||
int(np.prod(shape)),
|
||
)
|
||
self.dtype = dtype
|
||
self.shape = shape
|
||
|
||
def save(self, ndarray: np.ndarray) -> None:
|
||
"""
|
||
|
||
:param ndarray: np.ndarray:
|
||
:param ndarray: np.ndarray:
|
||
:param ndarray: np.ndarray:
|
||
|
||
"""
|
||
assert isinstance(ndarray, np.ndarray)
|
||
dst = self.arr.get_obj()
|
||
dst_np = np.frombuffer(dst, dtype=self.dtype).reshape(self.shape)
|
||
np.copyto(dst_np, ndarray)
|
||
|
||
def get(self) -> np.ndarray:
|
||
""" """
|
||
obj = self.arr.get_obj()
|
||
return np.frombuffer(obj, dtype=self.dtype).reshape(self.shape)
|
||
|
||
|
||
def _setup_buf(space: gym.Space) -> Union[dict, tuple, ShArray]:
|
||
"""
|
||
|
||
:param space: gym.Space:
|
||
:param space: gym.Space:
|
||
:param space: gym.Space:
|
||
|
||
"""
|
||
if isinstance(space, gym.spaces.Dict):
|
||
assert isinstance(space.spaces, OrderedDict)
|
||
return {k: _setup_buf(v) for k, v in space.spaces.items()}
|
||
elif isinstance(space, gym.spaces.Tuple):
|
||
assert isinstance(space.spaces, tuple)
|
||
return tuple([_setup_buf(t) for t in space.spaces])
|
||
else:
|
||
return ShArray(space.dtype, space.shape)
|
||
|
||
|
||
def _worker(
|
||
parent: connection.Connection,
|
||
p: connection.Connection,
|
||
env_fn_wrapper: CloudpickleWrapper,
|
||
obs_bufs: Optional[Union[dict, tuple, ShArray]] = None,
|
||
) -> None:
|
||
"""
|
||
|
||
:param parent: connection.Connection:
|
||
:param p: connection.Connection:
|
||
:param env_fn_wrapper: CloudpickleWrapper:
|
||
:param obs_bufs: Optional[Union[dict:
|
||
:param tuple: param ShArray]]: (Default value = None)
|
||
:param parent: connection.Connection:
|
||
:param p: connection.Connection:
|
||
:param env_fn_wrapper: CloudpickleWrapper:
|
||
:param obs_bufs: Optional[Union[dict:
|
||
:param ShArray]]: (Default value = None)
|
||
:param parent: connection.Connection:
|
||
:param p: connection.Connection:
|
||
:param env_fn_wrapper: CloudpickleWrapper:
|
||
:param obs_bufs: Optional[Union[dict:
|
||
|
||
"""
|
||
|
||
def _encode_obs(obs: Union[dict, tuple, np.ndarray], buffer: Union[dict, tuple, ShArray],) -> None:
|
||
"""
|
||
|
||
:param obs: Union[dict:
|
||
:param tuple: param np.ndarray]:
|
||
:param buffer: Union[dict:
|
||
:param ShArray:
|
||
:param obs: Union[dict:
|
||
:param np.ndarray]:
|
||
:param buffer: Union[dict:
|
||
:param ShArray]:
|
||
:param obs: Union[dict:
|
||
:param buffer: Union[dict:
|
||
|
||
"""
|
||
if isinstance(obs, np.ndarray) and isinstance(buffer, ShArray):
|
||
buffer.save(obs)
|
||
elif isinstance(obs, tuple) and isinstance(buffer, tuple):
|
||
for o, b in zip(obs, buffer):
|
||
_encode_obs(o, b)
|
||
elif isinstance(obs, dict) and isinstance(buffer, dict):
|
||
for k in obs.keys():
|
||
_encode_obs(obs[k], buffer[k])
|
||
return None
|
||
|
||
parent.close()
|
||
env = env_fn_wrapper.data()
|
||
try:
|
||
while True:
|
||
try:
|
||
cmd, data = p.recv()
|
||
except EOFError: # the pipe has been closed
|
||
p.close()
|
||
break
|
||
if cmd == "step":
|
||
obs, reward, done, info = env.step(data)
|
||
if obs_bufs is not None:
|
||
_encode_obs(obs, obs_bufs)
|
||
obs = None
|
||
p.send((obs, reward, done, info))
|
||
elif cmd == "reset":
|
||
obs = env.reset(data)
|
||
if obs_bufs is not None:
|
||
_encode_obs(obs, obs_bufs)
|
||
obs = None
|
||
p.send(obs)
|
||
elif cmd == "close":
|
||
p.send(env.close())
|
||
p.close()
|
||
break
|
||
elif cmd == "render":
|
||
p.send(env.render(**data) if hasattr(env, "render") else None)
|
||
elif cmd == "seed":
|
||
p.send(env.seed(data) if hasattr(env, "seed") else None)
|
||
elif cmd == "getattr":
|
||
p.send(getattr(env, data) if hasattr(env, data) else None)
|
||
elif cmd == "toggle_log":
|
||
env.toggle_log(data)
|
||
else:
|
||
p.close()
|
||
raise NotImplementedError
|
||
except KeyboardInterrupt:
|
||
p.close()
|
||
|
||
|
||
class SubprocEnvWorker(EnvWorker):
|
||
"""Subprocess worker used in SubprocVectorEnv and ShmemVectorEnv."""
|
||
|
||
def __init__(self, env_fn: Callable[[], gym.Env], share_memory: bool = False) -> None:
|
||
super().__init__(env_fn)
|
||
self.parent_remote, self.child_remote = Pipe()
|
||
self.share_memory = share_memory
|
||
self.buffer: Optional[Union[dict, tuple, ShArray]] = None
|
||
if self.share_memory:
|
||
dummy = env_fn()
|
||
obs_space = dummy.observation_space
|
||
dummy.close()
|
||
del dummy
|
||
self.buffer = _setup_buf(obs_space)
|
||
args = (
|
||
self.parent_remote,
|
||
self.child_remote,
|
||
CloudpickleWrapper(env_fn),
|
||
self.buffer,
|
||
)
|
||
self.process = Process(target=_worker, args=args, daemon=True)
|
||
self.process.start()
|
||
self.child_remote.close()
|
||
|
||
def __getattr__(self, key: str) -> Any:
|
||
self.parent_remote.send(["getattr", key])
|
||
return self.parent_remote.recv()
|
||
|
||
def _decode_obs(self) -> Union[dict, tuple, np.ndarray]:
|
||
""" """
|
||
|
||
def decode_obs(buffer: Optional[Union[dict, tuple, ShArray]]) -> Union[dict, tuple, np.ndarray]:
|
||
"""
|
||
|
||
:param buffer: Optional[Union[dict:
|
||
:param tuple: param ShArray]]:
|
||
:param buffer: Optional[Union[dict:
|
||
:param ShArray]]:
|
||
:param buffer: Optional[Union[dict:
|
||
|
||
"""
|
||
if isinstance(buffer, ShArray):
|
||
return buffer.get()
|
||
elif isinstance(buffer, tuple):
|
||
return tuple([decode_obs(b) for b in buffer])
|
||
elif isinstance(buffer, dict):
|
||
return {k: decode_obs(v) for k, v in buffer.items()}
|
||
else:
|
||
raise NotImplementedError
|
||
|
||
return decode_obs(self.buffer)
|
||
|
||
def reset(self, sample) -> Any:
|
||
"""
|
||
|
||
:param sample:
|
||
|
||
"""
|
||
self.parent_remote.send(["reset", sample])
|
||
# obs = self.parent_remote.recv()
|
||
# if self.share_memory:
|
||
# obs = self._decode_obs()
|
||
# return obs
|
||
|
||
def get_reset_result(self):
|
||
""" """
|
||
obs = self.parent_remote.recv()
|
||
if self.share_memory:
|
||
obs = self._decode_obs()
|
||
return obs
|
||
|
||
@staticmethod
|
||
def wait( # type: ignore
|
||
workers: List["SubprocEnvWorker"], wait_num: int, timeout: Optional[float] = None,
|
||
) -> List["SubprocEnvWorker"]:
|
||
"""
|
||
|
||
:param # type: ignoreworkers: List["SubprocEnvWorker"]:
|
||
:param wait_num: int:
|
||
:param timeout: Optional[float]: (Default value = None)
|
||
:param # type: ignoreworkers: List["SubprocEnvWorker"]:
|
||
:param wait_num: int:
|
||
:param timeout: Optional[float]: (Default value = None)
|
||
|
||
"""
|
||
remain_conns = conns = [x.parent_remote for x in workers]
|
||
ready_conns: List[connection.Connection] = []
|
||
remain_time, t1 = timeout, time.time()
|
||
while len(remain_conns) > 0 and len(ready_conns) < wait_num:
|
||
if timeout:
|
||
remain_time = timeout - (time.time() - t1)
|
||
if remain_time <= 0:
|
||
break
|
||
# connection.wait hangs if the list is empty
|
||
new_ready_conns = connection.wait(remain_conns, timeout=remain_time)
|
||
ready_conns.extend(new_ready_conns) # type: ignore
|
||
remain_conns = [conn for conn in remain_conns if conn not in ready_conns]
|
||
return [workers[conns.index(con)] for con in ready_conns]
|
||
|
||
def send_action(self, action: np.ndarray) -> None:
|
||
"""
|
||
|
||
:param action: np.ndarray:
|
||
:param action: np.ndarray:
|
||
:param action: np.ndarray:
|
||
|
||
"""
|
||
self.parent_remote.send(["step", action])
|
||
|
||
def toggle_log(self, log):
|
||
self.parent_remote.send(["toggle_log", log])
|
||
|
||
def get_result(self,) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||
""" """
|
||
obs, rew, done, info = self.parent_remote.recv()
|
||
if self.share_memory:
|
||
obs = self._decode_obs()
|
||
return obs, rew, done, info
|
||
|
||
def seed(self, seed: Optional[int] = None) -> Optional[List[int]]:
|
||
"""
|
||
|
||
:param seed: Optional[int]: (Default value = None)
|
||
:param seed: Optional[int]: (Default value = None)
|
||
:param seed: Optional[int]: (Default value = None)
|
||
|
||
"""
|
||
self.parent_remote.send(["seed", seed])
|
||
return self.parent_remote.recv()
|
||
|
||
def render(self, **kwargs: Any) -> Any:
|
||
"""
|
||
|
||
:param **kwargs: Any:
|
||
:param **kwargs: Any:
|
||
|
||
"""
|
||
self.parent_remote.send(["render", kwargs])
|
||
return self.parent_remote.recv()
|
||
|
||
def close_env(self) -> None:
|
||
""" """
|
||
try:
|
||
self.parent_remote.send(["close", None])
|
||
# mp may be deleted so it may raise AttributeError
|
||
self.parent_remote.recv()
|
||
self.process.join()
|
||
except (BrokenPipeError, EOFError, AttributeError):
|
||
pass
|
||
# ensure the subproc is terminated
|
||
self.process.terminate()
|
||
|
||
|
||
class BaseVectorEnv(gym.Env):
|
||
"""Base class for vectorized environments wrapper.
|
||
Usage:
|
||
::
|
||
env_num = 8
|
||
envs = DummyVectorEnv([lambda: gym.make(task) for _ in range(env_num)])
|
||
assert len(envs) == env_num
|
||
It accepts a list of environment generators. In other words, an environment
|
||
generator ``efn`` of a specific task means that ``efn()`` returns the
|
||
environment of the given task, for example, ``gym.make(task)``.
|
||
All of the VectorEnv must inherit :class:`~tianshou.env.BaseVectorEnv`.
|
||
Here are some other usages:
|
||
::
|
||
envs.seed(2) # which is equal to the next line
|
||
envs.seed([2, 3, 4, 5, 6, 7, 8, 9]) # set specific seed for each env
|
||
obs = envs.reset() # reset all environments
|
||
obs = envs.reset([0, 5, 7]) # reset 3 specific environments
|
||
obs, rew, done, info = envs.step([1] * 8) # step synchronously
|
||
envs.render() # render all environments
|
||
envs.close() # close all environments
|
||
.. warning::
|
||
If you use your own environment, please make sure the ``seed`` method
|
||
is set up properly, e.g.,
|
||
::
|
||
def seed(self, seed):
|
||
np.random.seed(seed)
|
||
Otherwise, the outputs of these envs may be the same with each other.
|
||
|
||
:param env_fns: a list of callable envs
|
||
:param env:
|
||
:param worker_fn: a callable worker
|
||
:param worker: which contains the i
|
||
:param int: wait_num
|
||
:param env: step
|
||
:param environments: to finish a step is time
|
||
:param return: when
|
||
:param simulation: in these environments
|
||
:param is: disabled
|
||
:param float: timeout
|
||
:param vectorized: step it only deal with those environments spending time
|
||
:param within: timeout
|
||
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
env_fns: List[Callable[[], gym.Env]],
|
||
worker_fn: Callable[[Callable[[], gym.Env]], EnvWorker],
|
||
sampler=None,
|
||
testing: Optional[bool] = False,
|
||
wait_num: Optional[int] = None,
|
||
timeout: Optional[float] = None,
|
||
) -> None:
|
||
self._env_fns = env_fns
|
||
# A VectorEnv contains a pool of EnvWorkers, which corresponds to
|
||
# interact with the given envs (one worker <-> one env).
|
||
self.workers = [worker_fn(fn) for fn in env_fns]
|
||
self.worker_class = type(self.workers[0])
|
||
assert issubclass(self.worker_class, EnvWorker)
|
||
assert all([isinstance(w, self.worker_class) for w in self.workers])
|
||
|
||
self.env_num = len(env_fns)
|
||
self.wait_num = wait_num or len(env_fns)
|
||
assert 1 <= self.wait_num <= len(env_fns), f"wait_num should be in [1, {len(env_fns)}], but got {wait_num}"
|
||
self.timeout = timeout
|
||
assert self.timeout is None or self.timeout > 0, f"timeout is {timeout}, it should be positive if provided!"
|
||
self.is_async = self.wait_num != len(env_fns) or timeout is not None or testing
|
||
self.waiting_conn: List[EnvWorker] = []
|
||
# environments in self.ready_id is actually ready
|
||
# but environments in self.waiting_id are just waiting when checked,
|
||
# and they may be ready now, but this is not known until we check it
|
||
# in the step() function
|
||
self.waiting_id: List[int] = []
|
||
# all environments are ready in the beginning
|
||
self.ready_id = list(range(self.env_num))
|
||
self.is_closed = False
|
||
self.sampler = sampler
|
||
self.sample_obs = None
|
||
|
||
def _assert_is_not_closed(self) -> None:
|
||
""" """
|
||
assert not self.is_closed, f"Methods of {self.__class__.__name__} cannot be called after " "close."
|
||
|
||
def __len__(self) -> int:
|
||
"""Return len(self), which is the number of environments."""
|
||
return self.env_num
|
||
|
||
def __getattribute__(self, key: str) -> Any:
|
||
"""Switch the attribute getter depending on the key.
|
||
Any class who inherits ``gym.Env`` will inherit some attributes, like
|
||
``action_space``. However, we would like the attribute lookup to go
|
||
straight into the worker (in fact, this vector env's action_space is
|
||
always None).
|
||
"""
|
||
if key in [
|
||
"metadata",
|
||
"reward_range",
|
||
"spec",
|
||
"action_space",
|
||
"observation_space",
|
||
]: # reserved keys in gym.Env
|
||
return self.__getattr__(key)
|
||
else:
|
||
return super().__getattribute__(key)
|
||
|
||
def __getattr__(self, key: str) -> List[Any]:
|
||
"""Fetch a list of env attributes.
|
||
This function tries to retrieve an attribute from each individual
|
||
wrapped environment, if it does not belong to the wrapping vector
|
||
environment class.
|
||
"""
|
||
return [getattr(worker, key) for worker in self.workers]
|
||
|
||
def _wrap_id(self, id: Optional[Union[int, List[int], np.ndarray]] = None) -> Union[List[int], np.ndarray]:
|
||
"""
|
||
|
||
:param id: Optional[Union[int:
|
||
:param List: int]:
|
||
:param np: ndarray]]: (Default value = None)
|
||
:param id: Optional[Union[int:
|
||
:param List[int]:
|
||
:param np.ndarray]]: (Default value = None)
|
||
:param id: Optional[Union[int:
|
||
|
||
"""
|
||
if id is None:
|
||
id = list(range(self.env_num))
|
||
elif np.isscalar(id):
|
||
id = [id]
|
||
return id
|
||
|
||
def _assert_id(self, id: List[int]) -> None:
|
||
"""
|
||
|
||
:param id: List[int]:
|
||
:param id: List[int]:
|
||
:param id: List[int]:
|
||
|
||
"""
|
||
for i in id:
|
||
assert i not in self.waiting_id, f"Cannot interact with environment {i} which is stepping now."
|
||
assert i in self.ready_id, f"Can only interact with ready environments {self.ready_id}."
|
||
|
||
def reset(self, id: Optional[Union[int, List[int], np.ndarray]] = None) -> np.ndarray:
|
||
"""Reset the state of some envs and return initial observations.
|
||
If id is None, reset the state of all the environments and return
|
||
initial observations, otherwise reset the specific environments with
|
||
the given id, either an int or a list.
|
||
|
||
:param id: Optional[Union[int:
|
||
:param List: int]:
|
||
:param np: ndarray]]: (Default value = None)
|
||
:param id: Optional[Union[int:
|
||
:param List[int]:
|
||
:param np.ndarray]]: (Default value = None)
|
||
:param id: Optional[Union[int:
|
||
|
||
"""
|
||
start_time = time.time()
|
||
self._assert_is_not_closed()
|
||
id = self._wrap_id(id)
|
||
if self.is_async:
|
||
self._assert_id(id)
|
||
obs = []
|
||
stop_id = []
|
||
for i in id:
|
||
sample = self.sampler.sample()
|
||
if sample is None:
|
||
stop_id.append(i)
|
||
else:
|
||
self.workers[i].reset(sample)
|
||
for i in id:
|
||
if i in stop_id:
|
||
obs.append(self.sample_obs)
|
||
else:
|
||
this_obs = self.workers[i].get_reset_result()
|
||
if self.sample_obs is None:
|
||
self.sample_obs = this_obs
|
||
for j in range(len(obs)):
|
||
if obs[j] is None:
|
||
obs[j] = self.sample_obs
|
||
obs.append(this_obs)
|
||
|
||
if len(obs) > 0:
|
||
obs = np.stack(obs)
|
||
# if len(stop_id)> 0:
|
||
# obs_zero =
|
||
# print(time.time() - start_timed)
|
||
|
||
return obs, stop_id
|
||
|
||
def toggle_log(self, log):
|
||
for worker in self.workers:
|
||
worker.toggle_log(log)
|
||
|
||
def reset_sampler(self):
|
||
""" """
|
||
self.sampler.reset()
|
||
|
||
def step(self, action: np.ndarray, id: Optional[Union[int, List[int], np.ndarray]] = None) -> List[np.ndarray]:
|
||
"""Run one timestep of some environments' dynamics.
|
||
If id is None, run one timestep of all the environments’ dynamics;
|
||
otherwise run one timestep for some environments with given id, either
|
||
an int or a list. When the end of episode is reached, you are
|
||
responsible for calling reset(id) to reset this environment’s state.
|
||
Accept a batch of action and return a tuple (batch_obs, batch_rew,
|
||
batch_done, batch_info) in numpy format.
|
||
|
||
:param numpy: ndarray action: a batch of action provided by the agent.
|
||
:param action: np.ndarray:
|
||
:param id: Optional[Union[int:
|
||
:param List: int]:
|
||
:param np: ndarray]]: (Default value = None)
|
||
:param action: np.ndarray:
|
||
:param id: Optional[Union[int:
|
||
:param List[int]:
|
||
:param np.ndarray]]: (Default value = None)
|
||
:param action: np.ndarray:
|
||
:param id: Optional[Union[int:
|
||
:rtype: A tuple including four items
|
||
|
||
"""
|
||
self._assert_is_not_closed()
|
||
id = self._wrap_id(id)
|
||
if not self.is_async:
|
||
assert len(action) == len(id)
|
||
for i, j in enumerate(id):
|
||
self.workers[j].send_action(action[i])
|
||
result = []
|
||
for j in id:
|
||
obs, rew, done, info = self.workers[j].get_result()
|
||
info["env_id"] = j
|
||
result.append((obs, rew, done, info))
|
||
else:
|
||
if action is not None:
|
||
self._assert_id(id)
|
||
assert len(action) == len(id)
|
||
for i, (act, env_id) in enumerate(zip(action, id)):
|
||
self.workers[env_id].send_action(act)
|
||
self.waiting_conn.append(self.workers[env_id])
|
||
self.waiting_id.append(env_id)
|
||
self.ready_id = [x for x in self.ready_id if x not in id]
|
||
ready_conns: List[EnvWorker] = []
|
||
while not ready_conns:
|
||
ready_conns = self.worker_class.wait(self.waiting_conn, self.wait_num, self.timeout)
|
||
result = []
|
||
for conn in ready_conns:
|
||
waiting_index = self.waiting_conn.index(conn)
|
||
self.waiting_conn.pop(waiting_index)
|
||
env_id = self.waiting_id.pop(waiting_index)
|
||
obs, rew, done, info = conn.get_result()
|
||
info["env_id"] = env_id
|
||
result.append((obs, rew, done, info))
|
||
self.ready_id.append(env_id)
|
||
return list(map(np.stack, zip(*result)))
|
||
|
||
def seed(self, seed: Optional[Union[int, List[int]]] = None) -> List[Optional[List[int]]]:
|
||
"""Set the seed for all environments.
|
||
Accept ``None``, an int (which will extend ``i`` to
|
||
``[i, i + 1, i + 2, ...]``) or a list.
|
||
|
||
:param seed: Optional[Union[int:
|
||
:param List: int]]]: (Default value = None)
|
||
:param seed: Optional[Union[int:
|
||
:param List[int]]]: (Default value = None)
|
||
:param seed: Optional[Union[int:
|
||
:returns: The list of seeds used in this env's random number generators.
|
||
The first value in the list should be the "main" seed, or the value
|
||
which a reproducer pass to "seed".
|
||
|
||
"""
|
||
self._assert_is_not_closed()
|
||
seed_list: Union[List[None], List[int]]
|
||
if seed is None:
|
||
seed_list = [seed] * self.env_num
|
||
elif isinstance(seed, int):
|
||
seed_list = [seed + i for i in range(self.env_num)]
|
||
else:
|
||
seed_list = seed
|
||
return [w.seed(s) for w, s in zip(self.workers, seed_list)]
|
||
|
||
def render(self, **kwargs: Any) -> List[Any]:
|
||
"""Render all of the environments.
|
||
|
||
:param **kwargs: Any:
|
||
:param **kwargs: Any:
|
||
|
||
"""
|
||
self._assert_is_not_closed()
|
||
if self.is_async and len(self.waiting_id) > 0:
|
||
raise RuntimeError(f"Environments {self.waiting_id} are still stepping, cannot " "render them now.")
|
||
return [w.render(**kwargs) for w in self.workers]
|
||
|
||
def close(self) -> None:
|
||
"""Close all of the environments.
|
||
This function will be called only once (if not, it will be called
|
||
during garbage collected). This way, ``close`` of all workers can be
|
||
assured.
|
||
|
||
|
||
"""
|
||
self._assert_is_not_closed()
|
||
for w in self.workers:
|
||
w.close()
|
||
self.is_closed = True
|
||
|
||
def __del__(self) -> None:
|
||
"""Redirect to self.close()."""
|
||
if not self.is_closed:
|
||
self.close()
|
||
|
||
|
||
class SubprocVectorEnv(BaseVectorEnv):
|
||
"""Vectorized environment wrapper based on subprocess.
|
||
.. seealso::
|
||
Please refer to :class:`~tianshou.env.BaseVectorEnv` for more detailed
|
||
explanation.
|
||
|
||
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
env_fns: List[Callable[[], gym.Env]],
|
||
sampler=None,
|
||
testing=False,
|
||
wait_num: Optional[int] = None,
|
||
timeout: Optional[float] = None,
|
||
) -> None:
|
||
def worker_fn(fn: Callable[[], gym.Env]) -> SubprocEnvWorker:
|
||
"""
|
||
|
||
:param fn: Callable[[]:
|
||
:param gym: Env]:
|
||
:param fn: Callable[[]:
|
||
:param gym.Env]:
|
||
:param fn: Callable[[]:
|
||
|
||
"""
|
||
return SubprocEnvWorker(fn, share_memory=False)
|
||
|
||
super().__init__(env_fns, worker_fn, sampler, testing, wait_num=wait_num, timeout=timeout)
|
||
|
||
|
||
class ShmemVectorEnv(BaseVectorEnv):
|
||
"""Optimized SubprocVectorEnv with shared buffers to exchange observations.
|
||
ShmemVectorEnv has exactly the same API as SubprocVectorEnv.
|
||
.. seealso::
|
||
Please refer to :class:`~tianshou.env.SubprocVectorEnv` for more
|
||
detailed explanation.
|
||
|
||
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
env_fns: List[Callable[[], gym.Env]],
|
||
sampler=None,
|
||
testing=False,
|
||
wait_num: Optional[int] = None,
|
||
timeout: Optional[float] = None,
|
||
) -> None:
|
||
def worker_fn(fn: Callable[[], gym.Env]) -> SubprocEnvWorker:
|
||
"""
|
||
|
||
:param fn: Callable[[]:
|
||
:param gym: Env]:
|
||
:param fn: Callable[[]:
|
||
:param gym.Env]:
|
||
:param fn: Callable[[]:
|
||
|
||
"""
|
||
return SubprocEnvWorker(fn, share_memory=True)
|
||
|
||
super().__init__(env_fns, worker_fn, sampler, testing, wait_num=wait_num, timeout=timeout)
|