def write_observations(vec_env, env_start_idx, shared_obs, obs):
obs = list(iterate(vec_env.observation_space, obs))
for i in range(vec_env.num_envs):
write_to_shared_memory(
vec_env.observation_space,
env_start_idx + i,
obs[i],
shared_obs,
)
def _regym_worker_shared_memory(index: int, env_fn: Callable[[], gym.Env],
pipe: Connection, parent_pipe: Connection,
shared_memory: Tuple[SynchronizedArray],
error_queue: Queue):
'''
Based on function `gym.vector.async_vector_env._worker_shared_memory`
See that function's documentation
Custom additions:
- 'environment' command: To return underlying environment
- 'step' command returns:
Note: succ_obs dimensions:
[num_agents, num_environments, environment_observations]
'''
assert shared_memory is not None
env = env_fn()
observation_space = env.observation_space
parent_pipe.close()
try:
while True:
command, data = pipe.recv()
if command == 'reset':
observation = env.reset()
write_to_shared_memory(index, observation, shared_memory,
observation_space)
pipe.send((None, True))
elif command == 'step':
observation, reward, done, info = env.step(data)
if done:
observation = env.reset()
write_to_shared_memory(index, observation, shared_memory,
observation_space)
pipe.send(((None, reward, done, info), True))
elif command == 'seed':
env.seed(data)
pipe.send((None, True))
elif command == 'close':
pipe.send((None, True))
break
elif command == 'environment':
pipe.send((env, True))
elif command == '_check_observation_space':
pipe.send((data == observation_space, True))
else:
raise RuntimeError(
'Received unknown command `{0}`. Must '
'be one of {`reset`, `step`, `seed`, `close`, `env`'
'`_check_observation_space`}.'.format(command))
except (KeyboardInterrupt, Exception):
error_queue.put((index, ) + sys.exc_info()[:2])
pipe.send((None, False))
finally:
env.close()
def _worker_shared_memory(index, env_fn, pipe, parent_pipe, shared_memory,
error_queue):
assert shared_memory is not None
env = env_fn()
observation_space = env.observation_space
parent_pipe.close()
try:
while True:
command, data = pipe.recv()
if command == 'reset':
observation = env.reset()
write_to_shared_memory(index, observation, shared_memory,
observation_space)
pipe.send((None, True))
elif command == 'step':
observation, reward, done, info = env.step(data)
if type(done) == bool:
cond = done
else:
cond = all(done)
if cond:
observation = env.reset()
write_to_shared_memory(index, observation, shared_memory,
observation_space)
pipe.send(((None, reward, done, info), True))
elif command == 'seed':
env.seed(data)
pipe.send((None, True))
elif command == 'close':
pipe.send((None, True))
break
elif command == 'render':
img = env.render('rgb_array')
pipe.send((img, True))
elif command == '_check_observation_space':
pipe.send((data == observation_space, True))
else:
raise RuntimeError(
'Received unknown command `{0}`. Must '
'be one of {`reset`, `step`, `seed`, `close`, '
'`_check_observation_space`}.'.format(command))
except (KeyboardInterrupt, Exception):
error_queue.put((index,) + sys.exc_info()[:2])
pipe.send((None, False))
finally:
env.close()
def _worker_shared_memory(index, env_fn, pipe, parent_pipe, shared_memory, error_queue):
assert shared_memory is not None
env = env_fn()
observation_space = env.observation_space
parent_pipe.close()
try:
while True:
command, data = pipe.recv()
if command == "reset":
observation = env.reset()
write_to_shared_memory(
index, observation, shared_memory, observation_space
)
pipe.send((None, True))
elif command == "step":
observation, reward, done, info = env.step(data)
if done:
info["terminal_observation"] = observation
observation = env.reset()
write_to_shared_memory(
index, observation, shared_memory, observation_space
)
pipe.send(((None, reward, done, info), True))
elif command == "seed":
env.seed(data)
pipe.send((None, True))
elif command == "close":
pipe.send((None, True))
break
elif command == "_check_observation_space":
pipe.send((data == observation_space, True))
else:
raise RuntimeError(
"Received unknown command `{0}`. Must "
"be one of {`reset`, `step`, `seed`, `close`, "
"`_check_observation_space`}.".format(command)
)
except (KeyboardInterrupt, Exception):
error_queue.put((index,) + sys.exc_info()[:2])
pipe.send((None, False))
finally:
env.close()
def _worker_shared_memory(index, env_fn, pipe, parent_pipe, shared_memory,
error_queue):
assert shared_memory is not None
env = env_fn()
observation_space = env.observation_space
parent_pipe.close()
try:
while True:
command, data = pipe.recv()
if command == "reset":
if "return_info" in data and data["return_info"] == True:
observation, info = env.reset(**data)
write_to_shared_memory(observation_space, index,
observation, shared_memory)
pipe.send(((None, info), True))
else:
observation = env.reset(**data)
write_to_shared_memory(observation_space, index,
observation, shared_memory)
pipe.send((None, True))
elif command == "step":
observation, reward, done, info = env.step(data)
if done:
info["terminal_observation"] = observation
observation = env.reset()
write_to_shared_memory(observation_space, index, observation,
shared_memory)
pipe.send(((None, reward, done, info), True))
elif command == "seed":
env.seed(data)
pipe.send((None, True))
elif command == "close":
pipe.send((None, True))
break
elif command == "_call":
name, args, kwargs = data
if name in ["reset", "step", "seed", "close"]:
raise ValueError(
f"Trying to call function `{name}` with "
f"`_call`. Use `{name}` directly instead.")
function = getattr(env, name)
if callable(function):
pipe.send((function(*args, **kwargs), True))
else:
pipe.send((function, True))
elif command == "_setattr":
name, value = data
setattr(env, name, value)
pipe.send((None, True))
elif command == "_check_spaces":
pipe.send(((data[0] == observation_space,
data[1] == env.action_space), True))
else:
raise RuntimeError(
f"Received unknown command `{command}`. Must "
"be one of {`reset`, `step`, `seed`, `close`, `_call`, "
"`_setattr`, `_check_spaces`}.")
except (KeyboardInterrupt, Exception):
error_queue.put((index, ) + sys.exc_info()[:2])
pipe.send((None, False))
finally:
env.close()
