def _worker(
device_ids: List[int],
interval: float,
conn: mp.connection.Connection,
) -> None:
gpu_timeline: List[List[int]] = []
conn.send(None)
while not conn.poll(timeout=interval):
gpu_utils = collect_gpu_utils(device_ids)
gpu_timeline.append(gpu_utils)
conn.send(gpu_timeline)
def _poll_pipe(
self,
actor_idx: int,
pipe: mp.connection.Connection,
replay_buffer_lock: mp.synchronize.Lock,
exception_event: mp.synchronize.Event,
) -> None:
if pipe.closed:
return
try:
while pipe.poll() and not exception_event.is_set():
cmd, data = pipe.recv()
if cmd == "get_statistics":
assert data is None
with replay_buffer_lock:
stats = self.get_statistics()
pipe.send(stats)
elif cmd == "load":
self.load(data)
pipe.send(None)
elif cmd == "save":
self.save(data)
pipe.send(None)
elif cmd == "transition":
with replay_buffer_lock:
if "env_id" not in data:
data["env_id"] = actor_idx
self.replay_buffer.append(**data)
self._cumulative_steps += 1
elif cmd == "stop_episode":
idx = actor_idx if data is None else data
with replay_buffer_lock:
self.replay_buffer.stop_current_episode(env_id=idx)
stats = self.get_statistics()
pipe.send(stats)
else:
raise RuntimeError(
"Unknown command from actor: {}".format(cmd))
except EOFError:
pipe.close()
except Exception:
self.logger.exception("Poller loop failed. Exiting")
exception_event.set()
def _worker(
env_constructor: bytes,
auto_reset: bool,
pipe: mp.connection.Connection,
polling_period: float = 0.1,
):
"""Process to build and run an environment. Using a pipe to
communicate with parent, the process receives action, steps
the environment, and returns the observations.
Args:
env_constructor (bytes): Cloudpickled callable which constructs the environment.
auto_reset (bool): If True, auto resets environment when episode ends.
pipe (mp.connection.Connection): Child's end of the pipe.
polling_period (float, optional): Time to wait for keyboard interrupts. Defaults to 0.1.
Raises:
KeyError: If unknown message type is received.
"""
env = cloudpickle.loads(env_constructor)()
pipe.send((_Message.RESULT, None))
try:
while True:
if not pipe.poll(polling_period):
continue
message, payload = pipe.recv()
if message == _Message.SEED:
env_seed = env.seed(payload)
pipe.send((_Message.RESULT, env_seed))
elif message == _Message.ACCESS:
result = getattr(env, payload, None)
pipe.send((_Message.RESULT, result))
elif message == _Message.RESET:
observation = env.reset()
pipe.send((_Message.RESULT, observation))
elif message == _Message.STEP:
observation, reward, done, info = env.step(payload)
if done["__all__"] and auto_reset:
# Final observation can be obtained from `info` as follows:
# `final_obs = info[agent_id]["env_obs"]`
observation = env.reset()
pipe.send((_Message.RESULT, (observation, reward, done, info)))
elif message == _Message.CLOSE:
break
else:
raise KeyError(
f"Expected message from {_Message.__members__}, but got unknown message `{message}`."
)
except (Exception, KeyboardInterrupt):
etype, evalue, tb = sys.exc_info()
if etype == KeyboardInterrupt:
stacktrace = "".join(
traceback.format_exception(etype, evalue, None))
else:
stacktrace = "".join(traceback.format_exception(etype, evalue, tb))
payload = (mp.current_process().name, stacktrace)
pipe.send((_Message.EXCEPTION, payload))
finally:
env.close()
pipe.close()