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()