| 48 | |
| 49 | |
| 50 | class SubprocVecEnv(): |
| 51 | def __init__(self, env_fns): |
| 52 | """ |
| 53 | envs: list of gym environments to run in subprocesses |
| 54 | """ |
| 55 | self.closed = False |
| 56 | nenvs = len(env_fns) |
| 57 | self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)]) |
| 58 | self.ps = [Process(target=worker, args=(work_remote, remote, CloudpickleWrapper(env_fn))) |
| 59 | for (work_remote, remote, env_fn) in zip(self.work_remotes, self.remotes, env_fns)] |
| 60 | for p in self.ps: |
| 61 | p.daemon = True # if the main process crashes, we should not cause things to hang |
| 62 | p.start() |
| 63 | for remote in self.work_remotes: |
| 64 | remote.close() |
| 65 | |
| 66 | self.remotes[0].send(('get_spaces', None)) |
| 67 | self.action_space, self.observation_space = self.remotes[0].recv() |
| 68 | |
| 69 | self.remotes[0].send(('get_id', None)) |
| 70 | self.env_id = self.remotes[0].recv() |
| 71 | |
| 72 | def step(self, actions): |
| 73 | for remote, action in zip(self.remotes, actions): |
| 74 | remote.send(('step', action)) |
| 75 | results = [remote.recv() for remote in self.remotes] |
| 76 | obs, rews, dones, infos = zip(*results) |
| 77 | # print("Infos:", infos) |
| 78 | # for done, info in zip(dones, infos): |
| 79 | # if done: |
| 80 | # # print("Total reward:", info['reward'], "Num steps:", info['episode_length']) |
| 81 | # print("Returned info:", info, "Done:", done) |
| 82 | return np.stack(obs), np.stack(rews), np.stack(dones), infos |
| 83 | |
| 84 | def reset(self): |
| 85 | for remote in self.remotes: |
| 86 | remote.send(('reset', None)) |
| 87 | return np.stack([remote.recv() for remote in self.remotes]) |
| 88 | |
| 89 | def reset_task(self): |
| 90 | for remote in self.remotes: |
| 91 | remote.send(('reset_task', None)) |
| 92 | return np.stack([remote.recv() for remote in self.remotes]) |
| 93 | |
| 94 | def close(self): |
| 95 | if self.closed: |
| 96 | return |
| 97 | |
| 98 | for remote in self.remotes: |
| 99 | remote.send(('close', None)) |
| 100 | for p in self.ps: |
| 101 | p.join() |
| 102 | self.closed = True |
| 103 | |
| 104 | @property |
| 105 | def num_envs(self): |
| 106 | return len(self.remotes) |