Specify a gpu for a network

This commit is contained in:
Shangtong Zhang
2018-02-03 10:53:14 -07:00
parent fef04b6bf6
commit e91ec3be45
12 changed files with 93 additions and 103 deletions
+5 -5
View File
@@ -94,15 +94,15 @@ class A2CAgent:
rollout.append([None, None, pending_value, None, None, None])
processed_rollout = [None] * (len(rollout) - 1)
advantages = self.network.FloatTensor(np.zeros((config.num_workers, 1)))
advantages = self.network.tensor(np.zeros((config.num_workers, 1)))
returns = pending_value.data
for i in reversed(range(len(rollout) - 1)):
prob, log_prob, value, actions, rewards, terminals = rollout[i]
terminals = self.network.FloatTensor(terminals).unsqueeze(1)
rewards = self.network.FloatTensor(rewards).unsqueeze(1)
actions = self.network.LongTensor(actions).unsqueeze(1)
terminals = self.network.tensor(terminals).unsqueeze(1)
rewards = self.network.tensor(rewards).unsqueeze(1)
actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1)
next_value = rollout[i + 1][2]
returns = rewards + terminals * config.discount * returns
returns = rewards + config.discount * terminals * returns
td_error = rewards + config.discount * terminals * next_value.data - value.data
advantages = advantages * config.gae_tau * config.discount * terminals + td_error
processed_rollout[i] = [prob, log_prob, value, actions, returns, advantages]
+3 -3
View File
@@ -83,8 +83,8 @@ class DDPGAgent:
experiences = self.replay.sample()
states, actions, rewards, next_states, terminals = experiences
q_next = target_critic.predict(next_states, target_actor.predict(next_states))
terminals = critic.to_torch_variable(terminals).unsqueeze(1)
rewards = critic.to_torch_variable(rewards).unsqueeze(1)
terminals = critic.variable(terminals).unsqueeze(1)
rewards = critic.variable(rewards).unsqueeze(1)
q_next = config.discount * q_next * (1 - terminals)
q_next.add_(rewards)
q_next = q_next.detach()
@@ -99,7 +99,7 @@ class DDPGAgent:
actions = actor.predict(states, False)
var_actions = Variable(actions.data, requires_grad=True)
q = critic.predict(states, var_actions)
q.backward(critic.FloatTensor(np.ones(q.size())))
q.backward(critic.tensor(np.ones(q.size())))
actor.zero_grad()
self.actor_opt.zero_grad()
+3 -3
View File
@@ -60,11 +60,11 @@ class DQNAgent:
q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
else:
q_next, _ = q_next.max(1)
terminals = self.learning_network.to_torch_variable(terminals)
rewards = self.learning_network.to_torch_variable(rewards)
terminals = self.learning_network.variable(terminals)
rewards = self.learning_network.variable(rewards)
q_next = self.config.discount * q_next * (1 - terminals)
q_next.add_(rewards)
actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
actions = self.learning_network.variable(actions, torch.LongTensor).unsqueeze(1)
q = self.learning_network.predict(states, False)
q = q.gather(1, actions).squeeze(1)
loss = self.criterion(q, q_next)