def __init__(
self,
config: ExperimentConfig,
output_dir: str,
loaded_config_src_files: Optional[Dict[str, str]],
seed: Optional[int] = None,
mode: str = "train",
deterministic_cudnn: bool = False,
deterministic_agents: bool = False,
mp_ctx: Optional[BaseContext] = None,
multiprocessing_start_method: str = "forkserver",
extra_tag: str = "",
disable_tensorboard: bool = False,
disable_config_saving: bool = False,
):
self.config = config
self.output_dir = output_dir
self.loaded_config_src_files = loaded_config_src_files
self.seed = seed
self.deterministic_cudnn = deterministic_cudnn
self.mp_ctx = self.init_context(mp_ctx, multiprocessing_start_method)
self.extra_tag = extra_tag
self.mode = mode
self.visualizer: Optional[VizSuite] = None
self.deterministic_agents = deterministic_agents
self.disable_tensorboard = disable_tensorboard
self.disable_config_saving = disable_config_saving
assert self.mode in [
"train",
"test",
], "Only 'train' and 'test' modes supported in runner"
if self.deterministic_cudnn:
set_deterministic_cudnn()
if self.seed is not None:
set_seed(self.seed)
self.queues = {
"results": self.mp_ctx.Queue(),
"checkpoints": self.mp_ctx.Queue(),
}
self.processes: Dict[str, List[Union[BaseProcess,
mp.Process]]] = defaultdict(list)
self.current_checkpoint = None
self.local_start_time_str = time.strftime("%Y-%m-%d_%H-%M-%S",
time.localtime(time.time()))
self._is_closed: bool = False
def set_seed(self, seed: int) -> None:
set_seed(seed)
self.np_seeded_random_gen, _ = seeding.np_random(seed)
self.seed = seed
def set_seed(self, seed: int):
self.seed = seed
if seed is not None:
set_seed(seed)
