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, log_id_dir, initial_step_count, rank=0, ): # ARGS TO STATE VARS self._args = args self.nb_learners = args.nb_learners self.nb_workers = args.nb_workers self.rank = rank self.nb_step = args.nb_step self.nb_env = args.nb_env self.initial_step_count = initial_step_count self.epoch_len = args.epoch_len self.summary_freq = args.summary_freq self.nb_learn_batch = args.nb_learn_batch self.rollout_queue_size = args.rollout_queue_size # can be none if rank != 0 self.log_id_dir = log_id_dir # load saved registry classes REGISTRY.load_extern_classes(log_id_dir) # ENV (temporary) env_cls = REGISTRY.lookup_env(args.env) env = env_cls.from_args(args, 0) env_action_space, env_observation_space, env_gpu_preprocessor = \ env.action_space, env.observation_space, env.gpu_preprocessor env.close() # NETWORK torch.manual_seed(args.seed) device = torch.device("cuda") # ray handles gpus torch.backends.cudnn.benchmark = True output_space = REGISTRY.lookup_output_space( args.actor_worker, env_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_observation_space, output_space, env_gpu_preprocessor, REGISTRY ) self.network = net.to(device) # TODO: this is a hack, remove once queuer puts rollouts on the correct device self.network.device = device self.device = device self.network.train() # OPTIMIZER def optim_fn(x): return torch.optim.RMSprop(x, lr=args.lr, eps=1e-5, alpha=0.99) if args.nb_learners > 1: self.optimizer = NCCLOptimizer(optim_fn, self.network, self.nb_learners) else: self.optimizer = optim_fn(self.network.parameters()) # LEARNER / EXP rwd_norm = REGISTRY.lookup_reward_normalizer( args.rwd_norm).from_args(args) actor_cls = REGISTRY.lookup_actor(args.actor_host) builder = actor_cls.exp_spec_builder( env.observation_space, env.action_space, net.internal_space(), args.nb_env * args.nb_learn_batch ) w_builder = REGISTRY.lookup_actor(args.actor_worker).exp_spec_builder( env.observation_space, env.action_space, net.internal_space(), args.nb_env ) actor = actor_cls.from_args(args, env.action_space) learner = REGISTRY.lookup_learner(args.learner).from_args(args, rwd_norm) exp_cls = REGISTRY.lookup_exp(args.exp).from_args(args, builder) self.actor = actor self.learner = learner self.exp = exp_cls.from_args(args, builder).to(device) # Rank 0 setup, load network/optimizer and create SummaryWriter/Saver if rank == 0: if args.load_network: self.network = self.load_network(self.network, args.load_network) print('Reloaded network from {}'.format(args.load_network)) if args.load_optim: self.optimizer = self.load_optim(self.optimizer, args.load_optim) print('Reloaded optimizer from {}'.format(args.load_optim)) print('Network parameters: ' + str(self.count_parameters(net))) self.summary_writer = SummaryWriter(log_id_dir) self.saver = SimpleModelSaver(log_id_dir)
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()