ParallelRolloutRunner

class maze.core.rollout.parallel_rollout_runner.ParallelRolloutRunner(n_episodes: int, max_episode_steps: int, deterministic: bool, n_processes: int, record_trajectory: bool, record_event_logs: bool)

Runs rollout in multiple processes in parallel.

Both agent and environment are run in multiple instances across multiple processes. While this greatly speeds up the rollout, the memory consumption might be high for large environments and agents.

Trajectory recording, event logging, as well as stats logging are supported. Trajectory logging happens in the child processes. Event logs and stats are shipped back to the main process so that they can be handled together there. This allows monitoring of progress and calculation of summary stats across all the processes.

(Note that the relevant wrappers need to be present in the config for the trajectory/event/stats logging to work. Data are logged into the working directory managed by hydra.)

In case of early rollout termination using a keyboard interrupt, data for all episodes completed till that point will be preserved (= written out). Graceful shutdown will be attempted, including calculation of statistics across the episodes completed before the rollout was terminated.

Parameters
  • n_episodes – Count of episodes to run.

  • max_episode_steps – Count of steps to run in each episode (if environment returns done, the episode will be finished earlier though).

  • deterministic – Deterministic or stochastic action sampling.

  • n_processes – Count of processes to spread the rollout across.

  • record_trajectory – Whether to record trajectory data.

  • record_event_logs – Whether to record event logs.

run_with(env: Union[None, Mapping[str, Any], Any], wrappers: Union[List[Union[None, Mapping[str, Any], Any]], Mapping[Union[str, Type], Union[None, Mapping[str, Any], Any]]], agent: Union[None, Mapping[str, Any], Any])

(overrides RolloutRunner)

Run the parallel rollout in multiple worker processes.