Upgrade to PyTorch v0.4

This commit is contained in:
Shangtong Zhang
2018-04-24 23:10:09 -06:00
parent fc7cfdf095
commit c7ceb876dd
13 changed files with 138 additions and 165 deletions
+11 -11
View File
@@ -32,7 +32,7 @@ class A2CAgent(BaseAgent):
states = self.states
for _ in range(config.rollout_length):
prob, log_prob, value = self.network.predict(config.state_normalizer(states))
actions = [self.policy.sample(p) for p in prob.data.cpu().numpy()]
actions = [self.policy.sample(p) for p in prob.cpu().detach().numpy()]
next_states, rewards, terminals, _ = self.task.step(actions)
self.episode_rewards += rewards
rewards = config.reward_normalizer(rewards)
@@ -50,34 +50,34 @@ class A2CAgent(BaseAgent):
processed_rollout = [None] * (len(rollout) - 1)
advantages = self.network.tensor(np.zeros((config.num_workers, 1)))
returns = pending_value.data
returns = pending_value.detach()
for i in reversed(range(len(rollout) - 1)):
prob, log_prob, value, actions, rewards, terminals = rollout[i]
terminals = self.network.tensor(terminals).unsqueeze(1)
rewards = self.network.tensor(rewards).unsqueeze(1)
actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1)
actions = self.network.tensor(actions).unsqueeze(1).long()
next_value = rollout[i + 1][2]
returns = rewards + config.discount * terminals * returns
if not config.use_gae:
advantages = returns - value.data
advantages = returns - value.detach()
else:
td_error = rewards + config.discount * terminals * next_value.data - value.data
td_error = rewards + config.discount * terminals * next_value.detach() - value.detach()
advantages = advantages * config.gae_tau * config.discount * terminals + td_error
processed_rollout[i] = [prob, log_prob, value, actions, returns, advantages]
prob, log_prob, value, actions, returns, advantages = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout))
policy_loss = -log_prob.gather(1, Variable(actions)) * Variable(advantages)
policy_loss = -log_prob.gather(1, actions) * advantages
entropy_loss = torch.sum(prob * log_prob, dim=1, keepdim=True)
value_loss = 0.5 * (Variable(returns) - value).pow(2)
value_loss = 0.5 * (returns - value).pow(2)
self.policy_loss = np.mean(policy_loss.data.cpu().numpy())
self.entropy_loss = np.mean(entropy_loss.data.cpu().numpy())
self.value_loss = np.mean(value_loss.data.cpu().numpy())
self.policy_loss = np.mean(policy_loss.cpu().detach().numpy())
self.entropy_loss = np.mean(entropy_loss.cpu().detach().numpy())
self.value_loss = np.mean(value_loss.cpu().detach().numpy())
self.optimizer.zero_grad()
(policy_loss + config.entropy_weight * entropy_loss +
config.value_loss_weight * value_loss).mean().backward()
nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip)
nn.utils.clip_grad_norm_(self.network.parameters(), config.gradient_clip)
self.optimizer.step()
self.evaluate(config.rollout_length)
+11 -12
View File
@@ -34,8 +34,8 @@ class CategoricalDQNAgent(BaseAgent):
self.delta_atom = (config.categorical_v_max - config.categorical_v_min) / float(config.categorical_n_atoms - 1)
def evaluation_action(self, state):
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
value = (value * self.atoms).sum(-1).cpu().numpy().flatten()
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).detach()
value = (value * self.atoms).sum(-1).cpu().detach().numpy().flatten()
return np.argmax(value)
def episode(self, deterministic=False):
@@ -44,9 +44,9 @@ class CategoricalDQNAgent(BaseAgent):
total_reward = 0.0
steps = 0
while True:
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).detach()
# self.config.logger.histo_summary('prob', value, self.total_steps)
value = (value * self.atoms).sum(-1).cpu().numpy().flatten()
value = (value * self.atoms).sum(-1).cpu().detach().numpy().flatten()
# self.config.logger.histo_summary('q', value, self.total_steps)
if deterministic:
action = np.argmax(value)
@@ -68,13 +68,13 @@ class CategoricalDQNAgent(BaseAgent):
states, actions, rewards, next_states, terminals = experiences
states = self.config.state_normalizer(states)
next_states = self.config.state_normalizer(next_states)
prob_next = self.target_network.predict(next_states).data
prob_next = self.target_network.predict(next_states).detach()
q_next = (prob_next * self.atoms).sum(-1)
# self.config.logger.histo_summary('q next', q_next.cpu().numpy(), self.total_steps)
# self.config.logger.histo_summary('q next', q_next.cpu().detach().numpy(), self.total_steps)
_, a_next = torch.max(q_next, dim=1)
a_next = a_next.view(-1, 1, 1).expand(-1, -1, prob_next.size(2))
prob_next = prob_next.gather(1, a_next).squeeze(1)
# self.config.logger.histo_summary('prob next', prob_next.cpu().numpy(), self.total_steps)
# self.config.logger.histo_summary('prob next', prob_next.cpu().detach().numpy(), self.total_steps)
rewards = self.network.tensor(rewards)
terminals = self.network.tensor(terminals)
@@ -92,14 +92,13 @@ class CategoricalDQNAgent(BaseAgent):
target_prob[i].index_add_(0, u[i].long(), d_m_u[i])
prob = self.network.predict(states)
actions = self.network.tensor(actions, torch.LongTensor)
actions = self.network.tensor(actions).long()
actions = actions.view(-1, 1, 1).expand(-1, -1, prob.size(2))
prob = prob.gather(1, Variable(actions)).squeeze(1)
loss = -(Variable(target_prob) * prob.log()).sum(-1).mean()
# self.config.logger.scalar_summary('loss', loss.data.cpu().numpy().flatten(), self.total_steps)
prob = prob.gather(1, actions).squeeze(1)
loss = -(target_prob * prob.log()).sum(-1).mean()
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm(self.network.parameters(), self.config.gradient_clip)
nn.utils.clip_grad_norm_(self.network.parameters(), self.config.gradient_clip)
self.optimizer.step()
self.evaluate()
+8 -8
View File
@@ -35,8 +35,9 @@ class DDPGAgent(BaseAgent):
def soft_update(self, target, src):
for target_param, param in zip(target.parameters(), src.parameters()):
target_param.data.copy_(target_param.data * (1.0 - self.config.target_network_mix) +
param.data * self.config.target_network_mix)
target_param.detach_()
target_param.copy_(target_param * (1.0 - self.config.target_network_mix) +
param * self.config.target_network_mix)
def evaluation_action(self, state):
self.config.state_normalizer.set_read_only()
@@ -82,8 +83,8 @@ class DDPGAgent(BaseAgent):
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.variable(terminals).unsqueeze(1)
rewards = critic.variable(rewards).unsqueeze(1)
terminals = critic.tensor(terminals).unsqueeze(1)
rewards = critic.tensor(rewards).unsqueeze(1)
q_next = config.discount * q_next * (1 - terminals)
q_next.add_(rewards)
q_next = q_next.detach()
@@ -96,15 +97,14 @@ class DDPGAgent(BaseAgent):
self.critic_opt.step()
actions = actor.predict(states, False)
var_actions = Variable(actions.data, requires_grad=True)
var_actions = actions.detach().requires_grad_()
q = critic.predict(states, var_actions)
q.backward(critic.tensor(np.ones(q.size())))
actor.zero_grad()
self.actor_opt.zero_grad()
actions.backward(-var_actions.grad.data)
for param in actor.parameters():
param.grad.data.clamp(-config.gradient_clip, config.gradient_clip)
actions.backward(-var_actions.grad)
torch.nn.utils.clip_grad_value_(actor.parameters(), config.gradient_clip)
self.actor_opt.step()
self.soft_update(self.target_network, self.network)
+4 -4
View File
@@ -61,17 +61,17 @@ class DQNAgent(BaseAgent):
q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
else:
q_next, _ = q_next.max(1)
terminals = self.network.variable(terminals)
rewards = self.network.variable(rewards)
terminals = self.network.tensor(terminals)
rewards = self.network.tensor(rewards)
q_next = self.config.discount * q_next * (1 - terminals)
q_next.add_(rewards)
actions = self.network.variable(actions, torch.LongTensor).unsqueeze(1)
actions = self.network.tensor(actions).unsqueeze(1).long()
q = self.network.predict(states, False)
q = q.gather(1, actions).squeeze(1)
loss = self.criterion(q, q_next)
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm(self.network.parameters(), self.config.gradient_clip)
nn.utils.clip_grad_norm_(self.network.parameters(), self.config.gradient_clip)
self.optimizer.step()
self.evaluate()
+6 -6
View File
@@ -36,7 +36,7 @@ class NStepDQNAgent(BaseAgent):
states = self.states
for _ in range(config.rollout_length):
q = self.network.predict(self.config.state_normalizer(states))
actions = [self.policy.sample(v) for v in q.data.cpu().numpy()]
actions = [self.policy.sample(v) for v in q.cpu().detach().numpy()]
next_states, rewards, terminals, _ = self.task.step(actions)
self.episode_rewards += rewards
rewards = config.reward_normalizer(rewards)
@@ -56,22 +56,22 @@ class NStepDQNAgent(BaseAgent):
self.states = states
processed_rollout = [None] * (len(rollout))
returns = self.target_network.predict(config.state_normalizer(states)).data
returns = self.target_network.predict(config.state_normalizer(states)).detach()
returns, _ = torch.max(returns, dim=1, keepdim=True)
for i in reversed(range(len(rollout))):
q, actions, rewards, terminals = rollout[i]
actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1)
q = q.gather(1, Variable(actions))
actions = self.network.tensor(actions).unsqueeze(1).long()
q = q.gather(1, actions)
terminals = self.network.tensor(terminals).unsqueeze(1)
rewards = self.network.tensor(rewards).unsqueeze(1)
returns = rewards + config.discount * terminals * returns
processed_rollout[i] = [q, returns]
q, returns= map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout))
loss = 0.5 * (q - Variable(returns)).pow(2).mean()
loss = 0.5 * (q - returns).pow(2).mean()
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip)
nn.utils.clip_grad_norm_(self.network.parameters(), config.gradient_clip)
self.optimizer.step()
self.evaluate(config.rollout_length)
+8 -10
View File
@@ -32,7 +32,7 @@ class PPOAgent(BaseAgent):
states = self.states
for _ in range(config.rollout_length):
actions, log_probs, _, values = self.network.predict(states)
next_states, rewards, terminals, _ = self.task.step(actions.data.cpu().numpy())
next_states, rewards, terminals, _ = self.task.step(actions.cpu().detach().numpy())
self.episode_rewards += rewards
rewards = config.reward_normalizer(rewards)
for i, terminal in enumerate(terminals):
@@ -49,33 +49,31 @@ class PPOAgent(BaseAgent):
processed_rollout = [None] * (len(rollout) - 1)
advantages = self.network.tensor(np.zeros((config.num_workers, 1)))
returns = pending_value.data
returns = pending_value.detach()
for i in reversed(range(len(rollout) - 1)):
states, value, actions, log_probs, rewards, terminals = rollout[i]
terminals = self.network.tensor(terminals).unsqueeze(1)
rewards = self.network.tensor(rewards).unsqueeze(1)
actions = self.network.variable(actions)
states = self.network.variable(states)
actions = self.network.tensor(actions)
states = self.network.tensor(states)
next_value = rollout[i + 1][1]
returns = rewards + config.discount * terminals * returns
if not config.use_gae:
advantages = returns - value.data
advantages = returns - value.detach()
else:
td_error = rewards + config.discount * terminals * next_value.data - value.data
td_error = rewards + config.discount * terminals * next_value.detach() - value.detach()
advantages = advantages * config.gae_tau * config.discount * terminals + td_error
processed_rollout[i] = [states, actions, log_probs, returns, advantages]
states, actions, log_probs_old, returns, advantages = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout))
advantages = (advantages - advantages.mean()) / advantages.std()
advantages = Variable(advantages)
returns = Variable(returns)
batcher = Batcher(states.size(0) // config.num_mini_batches, [np.arange(states.size(0))])
for _ in range(config.optimization_epochs):
batcher.shuffle()
while not batcher.end():
batch_indices = batcher.next_batch()[0]
batch_indices = self.network.variable(batch_indices, torch.LongTensor)
batch_indices = self.network.tensor(batch_indices).long()
sampled_states = states[batch_indices]
sampled_actions = actions[batch_indices]
sampled_log_probs_old = log_probs_old[batch_indices]
@@ -93,7 +91,7 @@ class PPOAgent(BaseAgent):
self.network.zero_grad()
(policy_loss + value_loss).backward()
nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip)
nn.utils.clip_grad_norm_(self.network.parameters(), config.gradient_clip)
self.network.step()
steps = config.rollout_length * config.num_workers
+9 -9
View File
@@ -36,8 +36,8 @@ class QuantileRegressionDQNAgent(BaseAgent):
return 0.5 * x.pow(2) * cond + (x.abs() - 0.5) * (1 - cond)
def evaluation_action(self, state):
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
value = (value * self.quantile_weight).sum(-1).cpu().numpy().flatten()
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).detach()
value = (value * self.quantile_weight).sum(-1).cpu().detach().numpy().flatten()
return np.argmax(value)
def episode(self, deterministic=False):
@@ -46,8 +46,8 @@ class QuantileRegressionDQNAgent(BaseAgent):
total_reward = 0.0
steps = 0
while True:
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
value = (value * self.quantile_weight).sum(-1).cpu().numpy().flatten()
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).detach()
value = (value * self.quantile_weight).sum(-1).cpu().detach().numpy().flatten()
if deterministic:
action = np.argmax(value)
elif self.total_steps < self.config.exploration_steps:
@@ -69,7 +69,7 @@ class QuantileRegressionDQNAgent(BaseAgent):
states = self.config.state_normalizer(states)
next_states = self.config.state_normalizer(next_states)
quantiles_next = self.target_network.predict(next_states).data
quantiles_next = self.target_network.predict(next_states).detach()
q_next = (quantiles_next * self.quantile_weight).sum(-1)
_, a_next = torch.max(q_next, dim=1)
a_next = a_next.view(-1, 1, 1).expand(-1, -1, quantiles_next.size(2))
@@ -80,13 +80,13 @@ class QuantileRegressionDQNAgent(BaseAgent):
quantiles_next = rewards.view(-1, 1) + self.config.discount * (1 - terminals.view(-1, 1)) * quantiles_next
quantiles = self.network.predict(states)
actions = self.network.tensor(actions, torch.LongTensor)
actions = self.network.tensor(actions).long()
actions = actions.view(-1, 1, 1).expand(-1, -1, quantiles.size(2))
quantiles = quantiles.gather(1, Variable(actions)).squeeze(1)
quantiles = quantiles.gather(1, actions).squeeze(1)
quantiles_next = quantiles_next.t().unsqueeze(-1)
diff = Variable(quantiles_next) - quantiles
loss = self.huber(diff) * Variable(self.cumulative_density.view(1, -1) - (diff.data < 0).float()).abs()
diff = quantiles_next - quantiles
loss = self.huber(diff) * (self.cumulative_density.view(1, -1) - (diff.detach() < 0).float()).abs()
self.optimizer.zero_grad()
loss.mean(1).sum().backward()