| """ |
| An interface for asynchronous vectorized environments. |
| """ |
|
|
| import multiprocessing as mp |
| import numpy as np |
| from .vec_env import VecEnv, CloudpickleWrapper, clear_mpi_env_vars |
| import ctypes |
| from baselines import logger |
|
|
| from .util import dict_to_obs, obs_space_info, obs_to_dict |
|
|
| _NP_TO_CT = {np.float32: ctypes.c_float, |
| np.int32: ctypes.c_int32, |
| np.int8: ctypes.c_int8, |
| np.uint8: ctypes.c_char, |
| np.bool: ctypes.c_bool} |
|
|
|
|
| class ShmemVecEnv(VecEnv): |
| """ |
| Optimized version of SubprocVecEnv that uses shared variables to communicate observations. |
| """ |
|
|
| def __init__(self, env_fns, spaces=None, context='spawn'): |
| """ |
| If you don't specify observation_space, we'll have to create a dummy |
| environment to get it. |
| """ |
| ctx = mp.get_context(context) |
| if spaces: |
| observation_space, action_space = spaces |
| else: |
| logger.log('Creating dummy env object to get spaces') |
| with logger.scoped_configure(format_strs=[]): |
| dummy = env_fns[0]() |
| observation_space, action_space = dummy.observation_space, dummy.action_space |
| dummy.close() |
| del dummy |
| VecEnv.__init__(self, len(env_fns), observation_space, action_space) |
| self.obs_keys, self.obs_shapes, self.obs_dtypes = obs_space_info(observation_space) |
| self.obs_bufs = [ |
| {k: ctx.Array(_NP_TO_CT[self.obs_dtypes[k].type], int(np.prod(self.obs_shapes[k]))) for k in self.obs_keys} |
| for _ in env_fns] |
| self.parent_pipes = [] |
| self.procs = [] |
| with clear_mpi_env_vars(): |
| for env_fn, obs_buf in zip(env_fns, self.obs_bufs): |
| wrapped_fn = CloudpickleWrapper(env_fn) |
| parent_pipe, child_pipe = ctx.Pipe() |
| proc = ctx.Process(target=_subproc_worker, |
| args=(child_pipe, parent_pipe, wrapped_fn, obs_buf, self.obs_shapes, self.obs_dtypes, self.obs_keys)) |
| proc.daemon = True |
| self.procs.append(proc) |
| self.parent_pipes.append(parent_pipe) |
| proc.start() |
| child_pipe.close() |
| self.waiting_step = False |
| self.viewer = None |
|
|
| def reset(self): |
| if self.waiting_step: |
| logger.warn('Called reset() while waiting for the step to complete') |
| self.step_wait() |
| for pipe in self.parent_pipes: |
| pipe.send(('reset', None)) |
| return self._decode_obses([pipe.recv() for pipe in self.parent_pipes]) |
|
|
| def step_async(self, actions): |
| assert len(actions) == len(self.parent_pipes) |
| for pipe, act in zip(self.parent_pipes, actions): |
| pipe.send(('step', act)) |
| self.waiting_step = True |
|
|
| def step_wait(self): |
| outs = [pipe.recv() for pipe in self.parent_pipes] |
| self.waiting_step = False |
| obs, rews, dones, infos = zip(*outs) |
| return self._decode_obses(obs), np.array(rews), np.array(dones), infos |
|
|
| def close_extras(self): |
| if self.waiting_step: |
| self.step_wait() |
| for pipe in self.parent_pipes: |
| pipe.send(('close', None)) |
| for pipe in self.parent_pipes: |
| pipe.recv() |
| pipe.close() |
| for proc in self.procs: |
| proc.join() |
|
|
| def get_images(self, mode='human'): |
| for pipe in self.parent_pipes: |
| pipe.send(('render', None)) |
| return [pipe.recv() for pipe in self.parent_pipes] |
|
|
| def _decode_obses(self, obs): |
| result = {} |
| for k in self.obs_keys: |
|
|
| bufs = [b[k] for b in self.obs_bufs] |
| o = [np.frombuffer(b.get_obj(), dtype=self.obs_dtypes[k]).reshape(self.obs_shapes[k]) for b in bufs] |
| result[k] = np.array(o) |
| return dict_to_obs(result) |
|
|
|
|
| def _subproc_worker(pipe, parent_pipe, env_fn_wrapper, obs_bufs, obs_shapes, obs_dtypes, keys): |
| """ |
| Control a single environment instance using IPC and |
| shared memory. |
| """ |
| def _write_obs(maybe_dict_obs): |
| flatdict = obs_to_dict(maybe_dict_obs) |
| for k in keys: |
| dst = obs_bufs[k].get_obj() |
| dst_np = np.frombuffer(dst, dtype=obs_dtypes[k]).reshape(obs_shapes[k]) |
| np.copyto(dst_np, flatdict[k]) |
|
|
| env = env_fn_wrapper.x() |
| parent_pipe.close() |
| try: |
| while True: |
| cmd, data = pipe.recv() |
| if cmd == 'reset': |
| pipe.send(_write_obs(env.reset())) |
| elif cmd == 'step': |
| obs, reward, done, info = env.step(data) |
| if done: |
| obs = env.reset() |
| pipe.send((_write_obs(obs), reward, done, info)) |
| elif cmd == 'render': |
| pipe.send(env.render(mode='rgb_array')) |
| elif cmd == 'close': |
| pipe.send(None) |
| break |
| else: |
| raise RuntimeError('Got unrecognized cmd %s' % cmd) |
| except KeyboardInterrupt: |
| print('ShmemVecEnv worker: got KeyboardInterrupt') |
| finally: |
| env.close() |
|
|
