def main(args, unknown_args): args, config = parse_args_uargs(args, unknown_args) set_global_seed(args.seed) if args.logdir is not None: os.makedirs(args.logdir, exist_ok=True) dump_config(config, args.logdir, args.configs) if args.expdir is not None: module = import_module(expdir=args.expdir) # noqa: F841 env = ENVIRONMENTS.get_from_params(**config["environment"]) algorithm_name = config["algorithm"].pop("algorithm") if algorithm_name in OFFPOLICY_ALGORITHMS_NAMES: ALGORITHMS = OFFPOLICY_ALGORITHMS trainer_fn = OffpolicyTrainer sync_epoch = False weights_sync_mode = "critic" if env.discrete_actions else "actor" elif algorithm_name in ONPOLICY_ALGORITHMS_NAMES: ALGORITHMS = ONPOLICY_ALGORITHMS trainer_fn = OnpolicyTrainer sync_epoch = True weights_sync_mode = "actor" else: # @TODO: add registry for algorithms, trainers, samplers raise NotImplementedError() db_server = DATABASES.get_from_params( **config.get("db", {}), sync_epoch=sync_epoch ) algorithm_fn = ALGORITHMS.get(algorithm_name) algorithm = algorithm_fn.prepare_for_trainer(env_spec=env, config=config) if args.resume is not None: algorithm.load_checkpoint(filepath=args.resume) trainer = trainer_fn( algorithm=algorithm, env_spec=env, db_server=db_server, logdir=args.logdir, weights_sync_mode=weights_sync_mode, **config["trainer"], ) trainer.run()
def main(args, unknown_args): args, config = parse_args_uargs(args, unknown_args) set_global_seed(args.seed) if args.logdir is not None: os.makedirs(args.logdir, exist_ok=True) dump_config(args.configs, args.logdir) if args.expdir is not None: module = import_module(expdir=args.expdir) # noqa: F841 algorithm_name = config["algorithm"].pop("algorithm") if algorithm_name in OFFPOLICY_ALGORITHMS_NAMES: ALGORITHMS = OFFPOLICY_ALGORITHMS trainer_fn = OffpolicyTrainer sync_epoch = False else: ALGORITHMS = ONPOLICY_ALGORITHMS trainer_fn = OnpolicyTrainer sync_epoch = True db_server = DATABASES.get_from_params(**config.get("db", {}), sync_epoch=sync_epoch) env = ENVIRONMENTS.get_from_params(**config["environment"]) algorithm_fn = ALGORITHMS.get(algorithm_name) algorithm = algorithm_fn.prepare_for_trainer(env_spec=env, config=config) if args.resume is not None: algorithm.load_checkpoint(filepath=args.resume) trainer = trainer_fn( algorithm=algorithm, env_spec=env, db_server=db_server, **config["trainer"], logdir=args.logdir, ) def on_exit(): for p in trainer.get_processes(): p.terminate() atexit.register(on_exit) trainer.run()
def main(args, unknown_args): args, config = parse_args_uargs(args, unknown_args) if args.expdir is not None: module = import_module(expdir=args.expdir) # noqa: F841 environment_name = config["environment"].pop("environment") environment_fn = ENVIRONMENTS.get(environment_name) algorithm_name = config["algorithm"].pop("algorithm") algorithm_fn = ALGORITHMS.get(algorithm_name) processes = [] sampler_id = 0 def on_exit(): for p in processes: p.terminate() atexit.register(on_exit) params = dict( seed=args.seed, logdir=args.logdir, algorithm_fn=algorithm_fn, environment_fn=environment_fn, config=config, resume=args.resume, db=args.db ) if args.check: params_ = dict( vis=False, infer=False, id=sampler_id ) run_sampler(**params, **params_) for i in range(args.vis): params_ = dict( vis=True, infer=True, id=sampler_id, exploration_power=0.0 ) p = mp.Process(target=run_sampler, kwargs=dict(**params, **params_)) p.start() processes.append(p) sampler_id += 1 time.sleep(STEP_DELAY) for i in range(args.infer): params_ = dict( vis=False, infer=True, id=sampler_id, exploration_power=0.0 ) p = mp.Process(target=run_sampler, kwargs=dict(**params, **params_)) p.start() processes.append(p) sampler_id += 1 time.sleep(STEP_DELAY) for i in range(1, args.train + 1): exploration_power = i / args.train params_ = dict( vis=False, infer=False, id=sampler_id, exploration_power=exploration_power ) p = mp.Process(target=run_sampler, kwargs=dict(**params, **params_)) p.start() processes.append(p) sampler_id += 1 time.sleep(STEP_DELAY) for p in processes: p.join()