def __init__(self, state_size, use_network="", device="cuda",lr=5e-4):
super(Agent, self).__init__()
self.device = device
if use_network == "":
net = Policy(state_size[0]).to(self.device)
print("Build a new network!")
else:
try:
net = torch.load("./Networks/" + use_network).to(self.device)
net.classifier = nn.Sequential(*list(net.classifier.children())[:-1]) # delete the softmax layer
print("Loaded %s." % use_network)
except:
net = Policy(state_size[0]).to(self.device)
print("No such network named %s. Rebuild a new network!" % use_network)
self.policy = net
# self.policy = net.eval() # avoid dropout
self.optimizer = optim.Adam(self.policy.parameters(),lr=lr)