def from_defaults(args): if args.agent: agent_cls = R.lookup_agent(args.agent) agent_args = agent_cls.args else: h = R.lookup_actor(args.actor_host) w = R.lookup_actor(args.actor_worker) l = R.lookup_learner(args.learner) e = R.lookup_exp(args.exp) agent_args = {**h.args, **w.args, **l.args, **e.args} env_cls = R.lookup_env(args.env) rwdnorm_cls = R.lookup_reward_normalizer(args.rwd_norm) env_args = env_cls.args rwdnorm_args = rwdnorm_cls.args if args.custom_network: net_args = R.lookup_network(args.custom_network).args else: net_args = R.lookup_modular_args(args) args = DotDict({ **args, **agent_args, **env_args, **rwdnorm_args, **net_args }) return args
def from_prompt(args): if args.agent: agent_cls = R.lookup_agent(args.agent) agent_args = agent_cls.prompt(provided=args) else: h = R.lookup_actor(args.actor_host) w = R.lookup_actor(args.actor_worker) l = R.lookup_learner(args.learner) e = R.lookup_exp(args.exp) agent_args = { **h.prompt(args), **w.prompt(args), **l.prompt(args), **e.prompt(args), } env_cls = R.lookup_env(args.env) rwdnorm_cls = R.lookup_reward_normalizer(args.rwd_norm) env_args = env_cls.prompt(provided=args) rwdnorm_args = rwdnorm_cls.prompt(provided=args) if args.custom_network: net_args = R.lookup_network(args.custom_network).prompt() else: net_args = R.prompt_modular_args(args) args = DotDict({ **args, **agent_args, **env_args, **rwdnorm_args, **net_args }) return args
def __init__( self, args, logger, log_id_dir, initial_step_count, local_rank, global_rank, world_size, ): seed = ( args.seed if global_rank == 0 else args.seed + args.nb_env * global_rank ) logger.info("Using {} for rank {} seed.".format(seed, global_rank)) # ENV engine = REGISTRY.lookup_engine(args.env) env_cls = REGISTRY.lookup_env(args.env) mgr_cls = REGISTRY.lookup_manager(args.manager) env_mgr = mgr_cls.from_args(args, engine, env_cls, seed=seed) # NETWORK torch.manual_seed(args.seed) device = torch.device("cuda:{}".format(local_rank)) output_space = REGISTRY.lookup_output_space( args.agent, env_mgr.action_space ) if args.custom_network: net_cls = REGISTRY.lookup_network(args.custom_network) else: net_cls = ModularNetwork net = net_cls.from_args( args, env_mgr.observation_space, output_space, env_mgr.gpu_preprocessor, REGISTRY, ) logger.info("Network parameters: " + str(self.count_parameters(net))) def optim_fn(x): return torch.optim.RMSprop(x, lr=args.lr, eps=1e-5, alpha=0.99) # AGENT rwd_norm = REGISTRY.lookup_reward_normalizer(args.rwd_norm).from_args( args ) agent_cls = REGISTRY.lookup_agent(args.agent) builder = agent_cls.exp_spec_builder( env_mgr.observation_space, env_mgr.action_space, net.internal_space(), env_mgr.nb_env, ) agent = agent_cls.from_args( args, rwd_norm, env_mgr.action_space, builder ) self.agent = agent self.nb_step = args.nb_step self.env_mgr = env_mgr self.nb_env = args.nb_env self.network = net.to(device) self.optimizer = optim_fn(self.network.parameters()) self.device = device self.initial_step_count = initial_step_count self.log_id_dir = log_id_dir self.epoch_len = args.epoch_len self.summary_freq = args.summary_freq self.logger = logger self.summary_writer = SummaryWriter( os.path.join(log_id_dir, "rank{}".format(global_rank)) ) self.saver = SimpleModelSaver(log_id_dir) self.local_rank = local_rank self.global_rank = global_rank self.world_size = world_size self.updater = DistribUpdater( self.optimizer, self.network, args.grad_norm_clip, world_size, not args.no_divide, ) if args.load_network: self.network = self.load_network(self.network, args.load_network) logger.info("Reloaded network from {}".format(args.load_network)) if args.load_optim: self.optimizer = self.load_optim(self.optimizer, args.load_optim) logger.info("Reloaded optimizer from {}".format(args.load_optim)) self.network.train()
def __init__( self, args, logger, log_id_dir, initial_step_count, local_rank, global_rank, world_size ): seed = args.seed \ if global_rank == 0 \ else args.seed + args.nb_env * global_rank logger.info('Using {} for rank {} seed.'.format(seed, global_rank)) # ENV engine = REGISTRY.lookup_engine(args.env) env_cls = REGISTRY.lookup_env(args.env) env_mgr = SubProcEnvManager.from_args(args, engine, env_cls, seed=seed) # NETWORK torch.manual_seed(args.seed) device = torch.device("cuda:{}".format(local_rank)) output_space = REGISTRY.lookup_output_space( args.agent, env_mgr.action_space) if args.custom_network: net_cls = REGISTRY.lookup_network(args.custom_network) else: net_cls = ModularNetwork net = net_cls.from_args( args, env_mgr.observation_space, output_space, env_mgr.gpu_preprocessor, REGISTRY ) def optim_fn(x): return torch.optim.RMSprop(x, lr=args.lr, eps=1e-5, alpha=0.99) # AGENT rwd_norm = REGISTRY.lookup_reward_normalizer( args.rwd_norm).from_args(args) agent_cls = REGISTRY.lookup_agent(args.agent) builder = agent_cls.exp_spec_builder( env_mgr.observation_space, env_mgr.action_space, net.internal_space(), env_mgr.nb_env ) agent = agent_cls.from_args( args, rwd_norm, env_mgr.action_space, builder ) self.agent = agent self.nb_step = args.nb_step self.env_mgr = env_mgr self.nb_env = args.nb_env self.network = net.to(device) self.optimizer = optim_fn(self.network.parameters()) self.device = device self.initial_step_count = initial_step_count self.log_id_dir = log_id_dir self.epoch_len = args.epoch_len self.summary_freq = args.summary_freq self.logger = logger self.local_rank = local_rank self.global_rank = global_rank self.world_size = world_size if args.load_network: self.network = self.load_network(self.network, args.load_network) logger.info('Reloaded network from {}'.format(args.load_network)) if args.load_optim: self.optimizer = self.load_optim(self.optimizer, args.load_optim) logger.info('Reloaded optimizer from {}'.format(args.load_optim)) self.network.train()
def __init__(self, args, logger, log_id_dir, initial_step_count): # ENV engine = REGISTRY.lookup_engine(args.env) env_cls = REGISTRY.lookup_env(args.env) mgr_cls = REGISTRY.lookup_manager(args.manager) env_mgr = mgr_cls.from_args(args, engine, env_cls) # NETWORK torch.manual_seed(args.seed) if torch.cuda.is_available() and args.gpu_id >= 0: device = torch.device("cuda:{}".format(args.gpu_id)) torch.backends.cudnn.benchmark = True else: device = torch.device("cpu") output_space = REGISTRY.lookup_output_space(args.agent, env_mgr.action_space) if args.custom_network: net_cls = REGISTRY.lookup_network(args.custom_network) else: net_cls = ModularNetwork net = net_cls.from_args( args, env_mgr.gpu_preprocessor.observation_space, output_space, env_mgr.gpu_preprocessor, REGISTRY, ) logger.info("Network parameters: " + str(self.count_parameters(net))) def optim_fn(x): if args.optim == "RMSprop": return torch.optim.RMSprop(x, lr=args.lr, eps=1e-5, alpha=0.99) elif args.optim == "Adam": return torch.optim.Adam(x, lr=args.lr, eps=1e-5) def warmup_schedule(back_step): return back_step / args.warmup if back_step < args.warmup else 1.0 # AGENT rwd_norm = REGISTRY.lookup_reward_normalizer( args.rwd_norm).from_args(args) agent_cls = REGISTRY.lookup_agent(args.agent) builder = agent_cls.exp_spec_builder( env_mgr.observation_space, env_mgr.action_space, net.internal_space(), env_mgr.nb_env, ) agent = agent_cls.from_args(args, rwd_norm, env_mgr.action_space, builder) self.agent = agent.to(device) self.nb_step = args.nb_step self.env_mgr = env_mgr self.nb_env = args.nb_env self.network = net.to(device) self.optimizer = optim_fn(self.network.parameters()) self.scheduler = LambdaLR(self.optimizer, warmup_schedule) self.device = device self.initial_step_count = initial_step_count self.log_id_dir = log_id_dir self.epoch_len = args.epoch_len self.summary_freq = args.summary_freq self.logger = logger self.summary_writer = SummaryWriter(log_id_dir) self.saver = SimpleModelSaver(log_id_dir) self.updater = LocalUpdater(self.optimizer, self.network, args.grad_norm_clip) if args.load_network: self.network = self.load_network(self.network, args.load_network) logger.info("Reloaded network from {}".format(args.load_network)) if args.load_optim: self.optimizer = self.load_optim(self.optimizer, args.load_optim) logger.info("Reloaded optimizer from {}".format(args.load_optim)) self.network.train()