Rewrite state/reward normalizer

This commit is contained in:
Shangtong Zhang
2018-04-06 11:36:44 -06:00
parent 4c6481d5be
commit 8dfe6ff7c8
23 changed files with 125 additions and 926 deletions
+3 -4
View File
@@ -31,13 +31,12 @@ class A2CAgent(BaseAgent):
rollout = []
states = self.states
for i in range(config.rollout_length):
states = self.task.normalize_state(states)
states = config.state_normalizer(states)
prob, log_prob, value = self.network.predict(states)
actions = [self.policy.sample(p) for p in prob.data.cpu().numpy()]
actions = config.action_shift_fn(actions)
next_states, rewards, terminals, _ = self.task.step(actions)
self.episode_rewards += rewards
rewards = config.reward_shift_fn(rewards)
rewards = config.reward_normalizer(rewards)
for i, terminal in enumerate(terminals):
if terminals[i]:
self.last_episode_rewards[i] = self.episode_rewards[i]
@@ -47,7 +46,7 @@ class A2CAgent(BaseAgent):
states = next_states
self.states = states
_, _, pending_value = self.network.predict(self.task.normalize_state(states))
_, _, pending_value = self.network.predict(config.state_normalizer(states))
rollout.append([None, None, pending_value, None, None, None])
processed_rollout = [None] * (len(rollout) - 1)
+5 -1
View File
@@ -15,4 +15,8 @@ class BaseAgent:
self.task.close()
def save(self, filename):
pass
torch.save(self.network.state_dict(), filename)
def load(self, filename):
state_dict = torch.load(filename, map_location=lambda storage, loc: storage)
self.network.load_state_dict(state_dict)
+14 -14
View File
@@ -19,15 +19,15 @@ class CategoricalDQNAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.learning_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.optimizer = config.optimizer_fn(self.learning_network.parameters())
self.optimizer = config.optimizer_fn(self.network.parameters())
self.criterion = nn.MSELoss()
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
self.replay = config.replay_fn()
self.policy = config.policy_fn()
self.total_steps = 0
self.atoms = self.learning_network.tensor(
self.atoms = self.network.tensor(
np.linspace(config.categorical_v_min,
config.categorical_v_max,
config.categorical_n_atoms))
@@ -39,7 +39,7 @@ class CategoricalDQNAgent(BaseAgent):
total_reward = 0.0
steps = 0
while True:
value = self.learning_network.predict(np.stack([self.task.normalize_state(state)])).squeeze(0).data
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
# self.config.logger.histo_summary('prob', value, self.total_steps)
value = (value * self.atoms).sum(-1).cpu().numpy().flatten()
# self.config.logger.histo_summary('q', value, self.total_steps)
@@ -51,7 +51,7 @@ class CategoricalDQNAgent(BaseAgent):
action = self.policy.sample(value)
next_state, reward, done, _ = self.task.step(action)
total_reward += reward
reward = self.config.reward_shift_fn(reward)
reward = self.config.reward_normalizer(reward)
if not deterministic:
self.replay.feed([state, action, reward, next_state, int(done)])
self.total_steps += 1
@@ -62,8 +62,8 @@ class CategoricalDQNAgent(BaseAgent):
if not deterministic and self.total_steps > self.config.exploration_steps:
experiences = self.replay.sample()
states, actions, rewards, next_states, terminals = experiences
states = self.task.normalize_state(states)
next_states = self.task.normalize_state(next_states)
states = self.config.state_normalizer(states)
next_states = self.config.state_normalizer(next_states)
prob_next = self.target_network.predict(next_states).data
q_next = (prob_next * self.atoms).sum(-1)
# self.config.logger.histo_summary('q next', q_next.cpu().numpy(), self.total_steps)
@@ -72,8 +72,8 @@ class CategoricalDQNAgent(BaseAgent):
prob_next = prob_next.gather(1, a_next).squeeze(1)
# self.config.logger.histo_summary('prob next', prob_next.cpu().numpy(), self.total_steps)
rewards = self.learning_network.tensor(rewards)
terminals = self.learning_network.tensor(terminals)
rewards = self.network.tensor(rewards)
terminals = self.network.tensor(terminals)
atoms_next = rewards.view(-1, 1) + self.config.discount * (1 - terminals.view(-1, 1)) * self.atoms.view(1, -1)
# epsilon = 1e-5
atoms_next.clamp_(self.config.categorical_v_min, self.config.categorical_v_max)
@@ -82,13 +82,13 @@ class CategoricalDQNAgent(BaseAgent):
u = b.ceil()
d_m_l = (u + (l == u).float() - b) * prob_next
d_m_u = (b - l) * prob_next
target_prob = self.learning_network.tensor(np.zeros(prob_next.size()))
target_prob = self.network.tensor(np.zeros(prob_next.size()))
for i in range(target_prob.size(0)):
target_prob[i].index_add_(0, l[i].long(), d_m_l[i])
target_prob[i].index_add_(0, u[i].long(), d_m_u[i])
prob = self.learning_network.predict(states)
actions = self.learning_network.tensor(actions, torch.LongTensor)
prob = self.network.predict(states)
actions = self.network.tensor(actions, torch.LongTensor)
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()
@@ -97,7 +97,7 @@ class CategoricalDQNAgent(BaseAgent):
loss.backward()
self.optimizer.step()
if not deterministic and self.total_steps % self.config.target_network_update_freq == 0:
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
if not deterministic and self.total_steps > self.config.exploration_steps:
self.policy.update_epsilon()
episode_time = time.time() - episode_start_time
+15 -17
View File
@@ -19,19 +19,20 @@ class DDPGAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.worker_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network.load_state_dict(self.worker_network.state_dict())
self.actor_opt = config.actor_optimizer_fn(self.worker_network.actor.parameters())
self.critic_opt = config.critic_optimizer_fn(self.worker_network.critic.parameters())
self.network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim,
config.actor_network_fn, config.critic_network_fn)
self.actor = self.network.actor
self.critic = self.network.critic
self.target_network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim,
config.actor_network_fn, config.critic_network_fn)
self.target_network.load_state_dict(self.network.state_dict())
self.actor_opt = config.actor_optimizer_fn(self.actor.parameters())
self.critic_opt = config.critic_optimizer_fn(self.critic.parameters())
self.replay = config.replay_fn()
self.random_process = config.random_process_fn(self.task.action_dim)
self.criterion = nn.MSELoss()
self.total_steps = 0
# self.state_normalizer = Normalizer(self.task.state_dim)
# self.reward_normalizer = Normalizer(1)
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) +
@@ -40,28 +41,26 @@ class DDPGAgent(BaseAgent):
def episode(self, deterministic=False, video_recorder=None):
self.random_process.reset_states()
state = self.task.reset()
# state = self.state_normalizer(state)
state = self.config.state_normalizer(state)
config = self.config
actor = self.worker_network.actor
critic = self.worker_network.critic
actor = self.network.actor
critic = self.network.critic
target_actor = self.target_network.actor
target_critic = self.target_network.critic
steps = 0
total_reward = 0.0
while True:
actor.eval()
action = actor.predict(np.stack([state]), True).flatten()
if not deterministic:
# action += config.gaussian_noise_scale * np.random.randn(*action.shape)
action += self.random_process.sample()
next_state, reward, done, info = self.task.step(action)
if video_recorder is not None:
video_recorder.capture_frame()
# next_state = self.state_normalizer(next_state)
next_state = self.config.state_normalizer(next_state)
total_reward += reward
# reward = self.reward_normalizer(reward)
reward = self.config.reward_normalizer(reward)
if not deterministic:
self.replay.feed([state, action, reward, next_state, int(done)])
@@ -74,7 +73,6 @@ class DDPGAgent(BaseAgent):
break
if not deterministic and self.replay.size() >= config.min_memory_size:
self.worker_network.train()
experiences = self.replay.sample()
states, actions, rewards, next_states, terminals = experiences
q_next = target_critic.predict(next_states, target_actor.predict(next_states))
@@ -103,6 +101,6 @@ class DDPGAgent(BaseAgent):
param.grad.data.clamp(-config.gradient_clip, config.gradient_clip)
self.actor_opt.step()
self.soft_update(self.target_network, self.worker_network)
self.soft_update(self.target_network, self.network)
return total_reward, steps
+13 -13
View File
@@ -19,11 +19,11 @@ class DQNAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.learning_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.optimizer = config.optimizer_fn(self.learning_network.parameters())
self.optimizer = config.optimizer_fn(self.network.parameters())
self.criterion = nn.MSELoss()
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
self.replay = config.replay_fn()
self.policy = config.policy_fn()
self.total_steps = 0
@@ -34,7 +34,7 @@ class DQNAgent(BaseAgent):
total_reward = 0.0
steps = 0
while True:
value = self.learning_network.predict(np.stack([self.task.normalize_state(state)]), True).flatten()
value = self.network.predict(np.stack([self.config.state_normalizer(state)]), True).flatten()
if deterministic:
action = np.argmax(value)
elif self.total_steps < self.config.exploration_steps:
@@ -43,7 +43,7 @@ class DQNAgent(BaseAgent):
action = self.policy.sample(value)
next_state, reward, done, _ = self.task.step(action)
total_reward += reward
reward = self.config.reward_shift_fn(reward)
reward = self.config.reward_normalizer(reward)
if not deterministic:
self.replay.feed([state, action, reward, next_state, int(done)])
self.total_steps += 1
@@ -54,27 +54,27 @@ class DQNAgent(BaseAgent):
if not deterministic and self.total_steps > self.config.exploration_steps:
experiences = self.replay.sample()
states, actions, rewards, next_states, terminals = experiences
states = self.task.normalize_state(states)
next_states = self.task.normalize_state(next_states)
states = self.config.state_normalizer(states)
next_states = self.config.state_normalizer(next_states)
q_next = self.target_network.predict(next_states, False).detach()
if self.config.double_q:
_, best_actions = self.learning_network.predict(next_states).detach().max(1)
_, best_actions = self.network.predict(next_states).detach().max(1)
q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
else:
q_next, _ = q_next.max(1)
terminals = self.learning_network.variable(terminals)
rewards = self.learning_network.variable(rewards)
terminals = self.network.variable(terminals)
rewards = self.network.variable(rewards)
q_next = self.config.discount * q_next * (1 - terminals)
q_next.add_(rewards)
actions = self.learning_network.variable(actions, torch.LongTensor).unsqueeze(1)
q = self.learning_network.predict(states, False)
actions = self.network.variable(actions, torch.LongTensor).unsqueeze(1)
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()
self.optimizer.step()
if not deterministic and self.total_steps % self.config.target_network_update_freq == 0:
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
if not deterministic and self.total_steps > self.config.exploration_steps:
self.policy.update_epsilon()
episode_time = time.time() - episode_start_time
+10 -11
View File
@@ -19,10 +19,10 @@ class NStepDQNAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.learning_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.optimizer = config.optimizer_fn(self.learning_network.parameters())
self.target_network.load_state_dict(self.learning_network.state_dict())
self.optimizer = config.optimizer_fn(self.network.parameters())
self.target_network.load_state_dict(self.network.state_dict())
self.policy = config.policy_fn()
self.total_steps = 0
@@ -35,12 +35,11 @@ class NStepDQNAgent(BaseAgent):
rollout = []
states = self.states
for i in range(config.rollout_length):
q = self.learning_network.predict(self.task.normalize_state(states))
q = self.network.predict(self.config.state_normalizer(states))
actions = [self.policy.sample(v) for v in q.data.cpu().numpy()]
actions = config.action_shift_fn(actions)
next_states, rewards, terminals, _ = self.task.step(actions)
self.episode_rewards += rewards
rewards = config.reward_shift_fn(rewards)
rewards = config.reward_normalizer(rewards)
for i, terminal in enumerate(terminals):
if terminals[i]:
self.last_episode_rewards[i] = self.episode_rewards[i]
@@ -52,19 +51,19 @@ class NStepDQNAgent(BaseAgent):
self.policy.update_epsilon()
self.total_steps += config.num_workers
if self.total_steps / config.num_workers % config.target_network_update_freq == 0:
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
self.states = states
processed_rollout = [None] * (len(rollout))
returns = self.target_network.predict(self.task.normalize_state(states)).data
returns = self.target_network.predict(config.state_normalizer(states)).data
returns, _ = torch.max(returns, dim=1, keepdim=True)
for i in reversed(range(len(rollout))):
q, actions, rewards, terminals = rollout[i]
actions = self.learning_network.tensor(actions, torch.LongTensor).unsqueeze(1)
actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1)
q = q.gather(1, Variable(actions))
terminals = self.learning_network.tensor(terminals).unsqueeze(1)
rewards = self.learning_network.tensor(rewards).unsqueeze(1)
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]
+7 -6
View File
@@ -19,16 +19,17 @@ class PPOAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.actor = config.actor_network_fn(self.task.state_dim, self.task.action_dim)
self.critic = config.critic_network_fn(self.task.state_dim, self.task.action_dim)
self.network = DisjointActorCriticNet(self.task.state_dim, self.task.action_dim,
config.actor_network_fn, config.critic_network_fn)
self.actor = self.network.actor
self.critic = self.network.critic
self.actor_opt = config.actor_optimizer_fn(self.actor.parameters())
self.critic_opt = config.critic_optimizer_fn(self.critic.parameters())
self.total_steps = 0
self.episode_rewards = np.zeros(config.num_workers)
self.last_episode_rewards = np.zeros(config.num_workers)
self.state_normalizer = Normalizer(self.task.state_dim)
self.states = self.task.reset()
self.states = self.state_normalizer(self.states)
self.states = config.state_normalizer(self.states)
def iteration(self):
config = self.config
@@ -43,12 +44,12 @@ class PPOAgent(BaseAgent):
log_probs = torch.sum(log_probs, dim=1, keepdim=True)
next_states, rewards, terminals, _ = self.task.step(actions.data.cpu().numpy())
self.episode_rewards += rewards
rewards = config.reward_shift_fn(rewards)
rewards = config.reward_normalizer(rewards)
for i, terminal in enumerate(terminals):
if terminals[i]:
self.last_episode_rewards[i] = self.episode_rewards[i]
self.episode_rewards[i] = 0
next_states = self.state_normalizer(next_states)
next_states = config.state_normalizer(next_states)
rollout.append([states, values, actions, log_probs, rewards, 1 - terminals])
states = next_states
+13 -13
View File
@@ -19,16 +19,16 @@ class QuantileRegressionDQNAgent(BaseAgent):
BaseAgent.__init__(self)
self.config = config
self.task = config.task_fn()
self.learning_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.target_network = config.network_fn(self.task.state_dim, self.task.action_dim)
self.optimizer = config.optimizer_fn(self.learning_network.parameters())
self.optimizer = config.optimizer_fn(self.network.parameters())
self.criterion = nn.MSELoss()
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
self.replay = config.replay_fn()
self.policy = config.policy_fn()
self.total_steps = 0
self.quantile_weight = 1.0 / self.config.num_quantiles
self.cumulative_density = self.learning_network.tensor(
self.cumulative_density = self.network.tensor(
(2 * np.arange(self.config.num_quantiles) + 1) / (2.0 * self.config.num_quantiles))
def huber(self, x):
@@ -41,7 +41,7 @@ class QuantileRegressionDQNAgent(BaseAgent):
total_reward = 0.0
steps = 0
while True:
value = self.learning_network.predict(np.stack([self.task.normalize_state(state)])).squeeze(0).data
value = self.network.predict(np.stack([self.config.state_normalizer(state)])).squeeze(0).data
value = (value * self.quantile_weight).sum(-1).cpu().numpy().flatten()
if deterministic:
action = np.argmax(value)
@@ -51,7 +51,7 @@ class QuantileRegressionDQNAgent(BaseAgent):
action = self.policy.sample(value)
next_state, reward, done, _ = self.task.step(action)
total_reward += reward
reward = self.config.reward_shift_fn(reward)
reward = self.config.reward_normalizer(reward)
if not deterministic:
self.replay.feed([state, action, reward, next_state, int(done)])
self.total_steps += 1
@@ -62,8 +62,8 @@ class QuantileRegressionDQNAgent(BaseAgent):
if not deterministic and self.total_steps > self.config.exploration_steps:
experiences = self.replay.sample()
states, actions, rewards, next_states, terminals = experiences
states = self.task.normalize_state(states)
next_states = self.task.normalize_state(next_states)
states = self.config.state_normalizer(states)
next_states = self.config.state_normalizer(next_states)
quantiles_next = self.target_network.predict(next_states).data
q_next = (quantiles_next * self.quantile_weight).sum(-1)
@@ -71,12 +71,12 @@ class QuantileRegressionDQNAgent(BaseAgent):
a_next = a_next.view(-1, 1, 1).expand(-1, -1, quantiles_next.size(2))
quantiles_next = quantiles_next.gather(1, a_next).squeeze(1)
rewards = self.learning_network.tensor(rewards)
terminals = self.learning_network.tensor(terminals)
rewards = self.network.tensor(rewards)
terminals = self.network.tensor(terminals)
quantiles_next = rewards.view(-1, 1) + self.config.discount * (1 - terminals.view(-1, 1)) * quantiles_next
quantiles = self.learning_network.predict(states)
actions = self.learning_network.tensor(actions, torch.LongTensor)
quantiles = self.network.predict(states)
actions = self.network.tensor(actions, torch.LongTensor)
actions = actions.view(-1, 1, 1).expand(-1, -1, quantiles.size(2))
quantiles = quantiles.gather(1, Variable(actions)).squeeze(1)
@@ -88,7 +88,7 @@ class QuantileRegressionDQNAgent(BaseAgent):
loss.mean(1).sum().backward()
self.optimizer.step()
if not deterministic and self.total_steps % self.config.target_network_update_freq == 0:
self.target_network.load_state_dict(self.learning_network.state_dict())
self.target_network.load_state_dict(self.network.state_dict())
if not deterministic and self.total_steps > self.config.exploration_steps:
self.policy.update_epsilon()
episode_time = time.time() - episode_start_time