def _worker_env(
connection_read_fn: Callable,
connection_write_fn: Callable,
env_fn: Callable,
env_fn_args: Tuple[Any],
auto_reset_done: bool,
mask_signals: bool = False,
child_pipe: Optional[Connection] = None,
parent_pipe: Optional[Connection] = None,
) -> None:
r"""process worker for creating and interacting with the environment."""
if mask_signals:
signal.signal(signal.SIGINT, signal.SIG_IGN)
signal.signal(signal.SIGTERM, signal.SIG_IGN)
signal.signal(signal.SIGUSR1, signal.SIG_IGN)
signal.signal(signal.SIGUSR2, signal.SIG_IGN)
env = env_fn(*env_fn_args)
if parent_pipe is not None:
parent_pipe.close()
try:
command, data = connection_read_fn()
while command != CLOSE_COMMAND:
if command == STEP_COMMAND:
# different step methods for habitat.RLEnv and habitat.Env
if isinstance(env, (habitat.RLEnv, gym.Env)):
# habitat.RLEnv
observations, reward, done, info = env.step(**data)
if auto_reset_done and done:
observations = env.reset()
with profiling_wrapper.RangeContext(
"worker write after step"
):
connection_write_fn(
(observations, reward, done, info)
)
elif isinstance(env, habitat.Env): # type: ignore
# habitat.Env
observations = env.step(**data)
if auto_reset_done and env.episode_over:
observations = env.reset()
connection_write_fn(observations)
else:
raise NotImplementedError
elif command == RESET_COMMAND:
observations = env.reset()
connection_write_fn(observations)
elif command == RENDER_COMMAND:
connection_write_fn(env.render(*data[0], **data[1]))
elif command == CALL_COMMAND:
function_name, function_args = data
if function_args is None:
function_args = {}
result_or_fn = getattr(env, function_name)
if len(function_args) > 0 or callable(result_or_fn):
result = result_or_fn(**function_args)
else:
result = result_or_fn
connection_write_fn(result)
elif command == COUNT_EPISODES_COMMAND:
connection_write_fn(len(env.episodes))
else:
raise NotImplementedError(f"Unknown command {command}")
with profiling_wrapper.RangeContext("worker wait for command"):
command, data = connection_read_fn()
except KeyboardInterrupt:
logger.info("Worker KeyboardInterrupt")
finally:
if child_pipe is not None:
child_pipe.close()
env.close()
def _worker_env(
connection_read_fn: Callable,
connection_write_fn: Callable,
env_fn: Callable,
env_fn_args: Tuple[Any],
auto_reset_done: bool,
mask_signals: bool = False,
child_pipe: Optional[Connection] = None,
parent_pipe: Optional[Connection] = None,
) -> None:
r"""process worker for creating and interacting with the environment."""
if mask_signals:
signal.signal(signal.SIGINT, signal.SIG_IGN)
signal.signal(signal.SIGTERM, signal.SIG_IGN)
signal.signal(signal.SIGUSR1, signal.SIG_IGN)
signal.signal(signal.SIGUSR2, signal.SIG_IGN)
env = env_fn(*env_fn_args)
if parent_pipe is not None:
parent_pipe.close()
try:
command, data = connection_read_fn()
while command != CLOSE_COMMAND:
if command == STEP_COMMAND:
# different step methods for habitat.RLEnv and habitat.Env
if isinstance(env, habitat.RLEnv) or isinstance(
env, gym.Env
):
# habitat.RLEnv
observations, reward, done, info = env.step(**data)
if auto_reset_done and done:
observations = env.reset()
with profiling_wrapper.RangeContext(
"worker write after step"
):
connection_write_fn(
(observations, reward, done, info)
)
elif isinstance(env, habitat.Env):
# habitat.Env
observations = env.step(**data)
if auto_reset_done and env.episode_over:
observations = env.reset()
connection_write_fn(observations)
else:
raise NotImplementedError
elif command == RESET_COMMAND:
observations = env.reset()
connection_write_fn(observations)
elif command == RENDER_COMMAND:
connection_write_fn(env.render(*data[0], **data[1]))
elif command in {
OBSERVATION_SPACE_COMMAND,
ACTION_SPACE_COMMAND,
NUMBER_OF_EPISODES_COMMAND,
}:
connection_write_fn(getattr(env, command))
elif command == CALL_COMMAND:
function_name, function_args = data
if function_args is None or len(function_args) == 0:
result = getattr(env, function_name)()
else:
result = getattr(env, function_name)(**function_args)
connection_write_fn(result)
# TODO: update CALL_COMMAND for getting attribute like this
elif command == EPISODE_COMMAND:
connection_write_fn(env.current_episode)
elif command == COUNT_EPISODES_COMMAND:
connection_write_fn(len(env.episodes))
elif command == EPISODE_OVER:
connection_write_fn(env.episode_over)
elif command == GET_METRICS:
result = env.get_metrics()
connection_write_fn(result)
else:
raise NotImplementedError
with profiling_wrapper.RangeContext("worker wait for command"):
command, data = connection_read_fn()
if child_pipe is not None:
child_pipe.close()
except KeyboardInterrupt:
logger.info("Worker KeyboardInterrupt")
finally:
env.close()