# Load model
try:
policy_net = utils.load_model(model_dir)
target_net = DQNModel(env.action_space, env=args.env)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()
print("Model successfully loaded\n")
except OSError:
policy_net = DQNModel(env.action_space, env=args.env)
target_net = DQNModel(env.action_space, env=args.env)
target_net.load_state_dict(policy_net.state_dict())
print("Model successfully created\n")
if torch.cuda.is_available():
policy_net.cuda()
target_net.cuda()
target_net.eval()
print("CUDA available: {}\n".format(torch.cuda.is_available()))
# Init Algorithm
algo = torch_rl.DQNAlgo_new(env,
policy_net,
target_net,
args.frames,
args.discount,
args.lr,
args.optim_eps,
args.batch_size,
preprocess_obss,
record_qvals=args.debug)
try:
base_model = utils.load_model(model_dir)
logger.info("Model successfully loaded\n")
except OSError:
if args.algo == "dqn":
base_model = DQNModel(obs_space, envs[0].action_space, args.mem,
args.text)
else:
base_model = ACModel(obs_space, envs[0].action_space, args.mem,
args.text)
logger.info("Model successfully created\n")
logger.info("{}\n".format(base_model))
if torch.cuda.is_available():
base_model.cuda()
logger.info("CUDA available: {}\n".format(torch.cuda.is_available()))
# Train model
num_frames = status["num_frames"]
total_start_time = time.time()
update = status["update"]
best_val = 0
if args.algo == "a2c":
algo = torch_rl.A2CAlgo(envs, base_model, args.frames_per_proc,
args.discount, args.lr, args.gae_lambda,
args.entropy_coef, args.value_loss_coef,
args.max_grad_norm, args.recurrence,
args.optim_alpha, args.optim_eps, preprocess_obss)