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 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 __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)