def __init__(self, env_fns, is_multi_agent=False):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, CloudpickleWrapper(env_fn),
is_multi_agent))
for (work_remote, env_fn) in zip(self.work_remotes, env_fns)
]
for p in self.ps:
p.daemon = True
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
self.action_space, self.observation_space = self.remotes[0].recv()
self.is_multi_agent = is_multi_agent
self.num_agents = None
if is_multi_agent:
try:
n = len(self.action_space)
except:
n = len(self.action_space.spaces)
self.num_agents = n
def __init__(self, env_fns, spaces=None, context='spawn', in_series=1):
"""
Arguments:
env_fns: iterable of callables - functions that create environments to run in subprocesses. Need to be cloud-pickleable
in_series: number of environments to run in series in a single process
(e.g. when len(env_fns) == 12 and in_series == 3, it will run 4 processes, each running 3 envs in series)
"""
self.waiting = False
self.closed = False
self.in_series = in_series
nenvs = len(env_fns)
assert nenvs % in_series == 0, "Number of envs must be divisible by number of envs to run in series"
self.nremotes = nenvs // in_series
env_fns = np.array_split(env_fns, self.nremotes)
ctx = mp.get_context(context)
self.remotes, self.work_remotes = zip(
*[ctx.Pipe() for _ in range(self.nremotes)])
self.ps = [
ctx.Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
with clear_mpi_env_vars():
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces_spec', None))
observation_space, action_space, self.spec, self.n_actors = self.remotes[
0].recv()
self.viewer = None
VecEnv.__init__(self, nenvs, observation_space, action_space)
def __init__(self, env_fns, spaces=None):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
observation_space, action_space = self.remotes[0].recv()
self.remotes[0].send(('get_agent_types', None))
self.agent_types = self.remotes[0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def worker(remote, parent_remote, env_fn_wrappers):
def step_env(env, action):
ob, reward, done, info = env.step(action)
if done:
ob = env.reset()
return ob, reward, done, info
parent_remote.close()
envs = [env_fn_wrapper() for env_fn_wrapper in env_fn_wrappers.x]
try:
while True:
cmd, data = remote.recv()
if cmd == 'step':
remote.send([step_env(env, action) for env, action in zip(envs, data)])
elif cmd == 'reset':
remote.send([env.reset() for env in envs])
elif cmd == 'adr':
remote.send([env.adr(progress) for env, progress in zip(envs, data)])
elif cmd == 'render':
remote.send([env.render(mode='rgb_array') for env in envs])
elif cmd == 'close':
remote.close()
break
elif cmd == 'get_spaces_spec':
remote.send(CloudpickleWrapper((envs[0].observation_space, envs[0].action_space, envs[0].spec)))
else:
raise NotImplementedError
except KeyboardInterrupt:
print('SubprocVecEnv worker: got KeyboardInterrupt')
finally:
for env in envs:
env.close()
def __init__(self, env_fns, spaces=None):
"""
Arguments:
env_fns: iterable of callables - functions that create environments to run in subprocesses. Need to be cloud-pickleable
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
observation_space, action_space, randomization_space = self.remotes[
0].recv()
self.randomization_space = randomization_space
self.viewer = None
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def __init__(self, env_fns, spaces=None):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
env = env_fns[0]()
o = env.reset()
n = len(o)
dim_o = o[0].shape[-1]
if gym.__version__ == "0.9.4":
obs = gym.spaces.Box(low=-np.inf, high=+np.inf, shape=(n, dim_o))
else:
obs = gym.spaces.Box(low=-np.inf,
high=+np.inf,
shape=(n, dim_o),
dtype=np.float32)
act = gym.spaces.Discrete(5)
self.nenvs = nenvs
VecEnv.__init__(self, nenvs, obs, act)
def __init__(self, env_fns):
"""
Creates a multiprocess vectorized wrapper for multiple environments
:param env_fns: ([Gym Environment]) Environments to run in subprocesses
"""
self.waiting = False
self.closed = False
n_envs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(n_envs)])
self.processes = [
Process(target=_worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for process in self.processes:
process.daemon = True # if the main process crashes, we should not cause things to hang
process.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
observation_space, action_space = self.remotes[0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def __init__(self, env_fns):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
self.num_envs = len(env_fns)
self.remotes, self.work_remotes = zip(
*[Pipe() for _ in range(self.num_envs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
self.first = self.remotes[0]
self.first.send(('get', 'action_space'))
self.action_space = self.first.recv()
self.first.send(('get', 'observation_space'))
self.observation_space = self.first.recv()
self.first.send(('get', 'spec'))
self.spec = self.first.recv()
def __init__(self, env_fns, spaces=None):
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=atari_subproc_worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
self.observation_space, self.action_space = self.remotes[0].recv()
self.env_handles = [
AtariSubprocEnvHandle(remote) for remote in self.remotes
]
self.nenvs = nenvs
def __init__(self,
env_fns,
spaces=None,
level_selector=None,
experiment=None,
dataset=None):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn),
level_selector))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
print("start processes")
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
observation_space, action_space = self.remotes[0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
if USE_IMMITATION_ENV:
self.featurizer = TDCFeaturizer(84,
84,
84,
84,
feature_vector_size=1024,
learning_rate=0,
experiment_name=experiment)
self.featurizer.load(dataset)
video_dataset = generate_dataset(experiment,
framerate=60,
width=84,
height=84)[0]
self.featurized_dataset = self.featurizer.featurize(video_dataset)
self.checkpoint_indexes = [0] * nenvs
self.rewards = 0
self.counter = 0
def __init__(self, model_fns, context='spawn'):
self.n_remotes = len(model_fns)
model_fns = np.array_split(model_fns, self.n_remotes)
ctx = mp.get_context(context)
self.remotes, self.work_remotes = zip(
*[ctx.Pipe() for _ in range(self.n_remotes)])
self.ps = [
ctx.Process(target=worker,
args=(work_remote, remote,
CloudpickleWrapper(model_fn)))
for (work_remote, remote,
model_fn) in zip(self.work_remotes, self.remotes, model_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
def __init__(self, env_fns, h5_path=None, spaces=None, context='spawn'):
"""
Arguments:
env_fns: iterable of callables - functions that create environments to run in subprocesses. Need to be cloud-pickleable
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
ctx = mp.get_context(context)
self.remotes, self.work_remotes = zip(
*[ctx.Pipe() for _ in range(nenvs)])
self.ps = [
ctx.Process(target=worker_acc,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
with clear_mpi_env_vars():
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces_spec', None))
observation_space, action_space, self.spec = self.remotes[0].recv()
self.viewer = None
# load HighD data
self.highd_f = h5py.File(h5_path, "r")
self.highd_vels = self.highd_f['initial_velocities']
self.highd_accs = self.highd_f['x_accelerations']
self.staliro_trace_queue = None
# self.obs_pos = None
# self.obs_vel = None
# self.obs_driver = None
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def __init__(self, env_fns, spaces=None):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
# Pipe() function returns a pair of connection objects connected by a pipe which by default is duplex (two-way)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [Process(target=worker, args=(work_remote, remote, CloudpickleWrapper(env_fn))) # Create process objects
for (work_remote, remote, env_fn) in zip(self.work_remotes, self.remotes, env_fns)]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang (Process's daemon flag must be set
# before before process starts)
p.start() # Start the process
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
observation_space, action_space = self.remotes[0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
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 __init__(self, env_fns):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [
Process(target=worker,
args=(work_remote, remote, CloudpickleWrapper(env_fn)))
for (work_remote, remote,
env_fn) in zip(self.work_remotes, self.remotes, env_fns)
]
for p in self.ps:
p.daemon = False # can't since vizdoom envs have their own daemonic subprocesses
p.start()
for remote in self.work_remotes:
remote.close()
self.remotes[0].send(('get_spaces', None))
self.state_space, observation_space, action_space = self.remotes[
0].recv()
VecEnv.__init__(self, len(env_fns), observation_space, action_space)
def __init__(self, env_fns, agent):
"""
envs: list of gym environments to run in subprocesses
"""
self.waiting = False
self.closed = False
nenvs = len(env_fns)
self.remotes, self.work_remotes = zip(*[Pipe() for _ in range(nenvs)])
self.ps = [Process(target=worker, args=(work_remote, remote, CloudpickleWrapper(env_fn), agent))
for (work_remote, remote, env_fn) in zip(self.work_remotes, self.remotes, env_fns)]
for p in self.ps:
p.daemon = True # if the main process crashes, we should not cause things to hang
p.start()
for remote in self.work_remotes:
remote.close()
for r in self.remotes:
r.send(('setup', None))
results = [r.recv() for r in self.remotes]
obs_space, action_space = zip(*results)
super().__init__(len(env_fns), obs_space[0], action_space[0])
self.num_actions = len(action_space[0][1])
