From 8dfe6ff7c8d61a6d3392c7efade679f6580f661f Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Fri, 6 Apr 2018 11:36:44 -0600 Subject: [PATCH] Rewrite state/reward normalizer --- agent/A2C_agent.py | 7 +- agent/BaseAgent.py | 6 +- agent/CategoricalDQN_agent.py | 28 ++--- agent/DDPG_agent.py | 32 +++-- agent/DQN_agent.py | 26 ++-- agent/NStepDQN_agent.py | 21 ++-- agent/PPO_agent.py | 13 +- agent/QuantileRegressionDQN_agent.py | 26 ++-- async_worker/__init__.py | 7 -- async_worker/actor_critic.py | 81 ------------ async_worker/continuous_actor_critic.py | 109 ---------------- async_worker/dpg.py | 118 ------------------ async_worker/n_step_q.py | 79 ------------ async_worker/one_step_q.py | 76 ----------- async_worker/one_step_sarsa.py | 83 ------------- async_worker/ppo.py | 159 ------------------------ component/atari_wrapper.py | 12 +- component/task.py | 3 - main.py | 28 +++-- network/continuous_action_network.py | 5 - utils/config.py | 11 +- utils/misc.py | 25 +--- utils/normalizer.py | 96 +++----------- 23 files changed, 125 insertions(+), 926 deletions(-) delete mode 100644 async_worker/__init__.py delete mode 100644 async_worker/actor_critic.py delete mode 100644 async_worker/continuous_actor_critic.py delete mode 100644 async_worker/dpg.py delete mode 100644 async_worker/n_step_q.py delete mode 100644 async_worker/one_step_q.py delete mode 100644 async_worker/one_step_sarsa.py delete mode 100644 async_worker/ppo.py diff --git a/agent/A2C_agent.py b/agent/A2C_agent.py index ead73dc..45a0652 100644 --- a/agent/A2C_agent.py +++ b/agent/A2C_agent.py @@ -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) diff --git a/agent/BaseAgent.py b/agent/BaseAgent.py index 6c8feae..e66899a 100644 --- a/agent/BaseAgent.py +++ b/agent/BaseAgent.py @@ -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) \ No newline at end of file diff --git a/agent/CategoricalDQN_agent.py b/agent/CategoricalDQN_agent.py index af8a081..34e9922 100644 --- a/agent/CategoricalDQN_agent.py +++ b/agent/CategoricalDQN_agent.py @@ -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 diff --git a/agent/DDPG_agent.py b/agent/DDPG_agent.py index 4aa01ee..26ac3ff 100644 --- a/agent/DDPG_agent.py +++ b/agent/DDPG_agent.py @@ -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 diff --git a/agent/DQN_agent.py b/agent/DQN_agent.py index 2404ee4..e51f734 100644 --- a/agent/DQN_agent.py +++ b/agent/DQN_agent.py @@ -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 diff --git a/agent/NStepDQN_agent.py b/agent/NStepDQN_agent.py index 599f676..15a4776 100644 --- a/agent/NStepDQN_agent.py +++ b/agent/NStepDQN_agent.py @@ -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] diff --git a/agent/PPO_agent.py b/agent/PPO_agent.py index f82e4da..28b5684 100644 --- a/agent/PPO_agent.py +++ b/agent/PPO_agent.py @@ -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 diff --git a/agent/QuantileRegressionDQN_agent.py b/agent/QuantileRegressionDQN_agent.py index 1d13705..7813667 100644 --- a/agent/QuantileRegressionDQN_agent.py +++ b/agent/QuantileRegressionDQN_agent.py @@ -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 diff --git a/async_worker/__init__.py b/async_worker/__init__.py deleted file mode 100644 index 9a2160d..0000000 --- a/async_worker/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -from .actor_critic import * -from .continuous_actor_critic import * -from .n_step_q import * -from .one_step_sarsa import * -from .one_step_q import * -from .ppo import * -from .dpg import * \ No newline at end of file diff --git a/async_worker/actor_critic.py b/async_worker/actor_critic.py deleted file mode 100644 index 07f1cb1..0000000 --- a/async_worker/actor_critic.py +++ /dev/null @@ -1,81 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable -import torch.nn as nn -from utils import * - -class AdvantageActorCritic: - def __init__(self, config, learning_network, target_network): - self.config = config - self.optimizer = config.optimizer_fn(learning_network.parameters()) - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(learning_network.state_dict()) - self.task = config.task_fn() - self.policy = config.policy_fn() - self.learning_network = learning_network - - def episode(self, deterministic=False): - config = self.config - state = self.task.reset() - steps = 0 - total_reward = 0 - pending = [] - while not config.stop_signal.value: - prob, log_prob, value = self.worker_network.predict(np.stack([state])) - action = self.policy.sample(prob.data.numpy().flatten(), deterministic) - next_state, reward, terminal, _ = self.task.step(action) - - steps += 1 - total_reward += reward - reward = config.reward_shift_fn(reward) - - if deterministic: - if terminal: - break - state = next_state - continue - - pending.append([prob, log_prob, value, action, reward]) - with config.steps_lock: - config.total_steps.value += 1 - - if terminal or len(pending) >= config.update_interval: - loss = 0 - if terminal: - R = torch.FloatTensor([[0]]) - else: - R = self.worker_network.critic(np.stack([next_state])).data - GAE = torch.FloatTensor([[0]]) - for i in reversed(range(len(pending))): - prob, log_prob, value, action, reward = pending[i] - if i == len(pending) - 1: - delta = reward + config.discount * R - value.data - else: - delta = reward + config.discount * pending[i + 1][2].data - value.data - GAE = config.discount * config.gae_tau * GAE + delta - loss += -log_prob.gather(1, Variable(torch.LongTensor([[action]]))) * Variable(GAE) - loss += config.entropy_weight * torch.sum(torch.mul(prob, log_prob)) - - R = reward + config.discount * R - loss += 0.5 * (Variable(R) - value).pow(2) - - pending = [] - self.worker_network.zero_grad() - self.optimizer.zero_grad() - loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - sync_grad(self.learning_network, self.worker_network) - self.optimizer.step() - self.worker_network.load_state_dict(self.learning_network.state_dict()) - self.worker_network.reset(terminal) - - if terminal: - break - state = next_state - - return steps, total_reward \ No newline at end of file diff --git a/async_worker/continuous_actor_critic.py b/async_worker/continuous_actor_critic.py deleted file mode 100644 index b62bc5e..0000000 --- a/async_worker/continuous_actor_critic.py +++ /dev/null @@ -1,109 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable -import torch.nn as nn -from utils import * - -class ContinuousAdvantageActorCritic: - def __init__(self, config, learning_network, extra): - self.config = config - self.actor_opt = config.actor_optimizer_fn(learning_network.actor.parameters()) - self.critic_opt = config.critic_optimizer_fn(learning_network.critic.parameters()) - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(learning_network.state_dict()) - self.task = config.task_fn() - self.policy = config.policy_fn() - self.learning_network = learning_network - self.counter = 0 - - self.shared_state_normalizer = extra[0] - self.state_normalizer = StaticNormalizer(self.task.state_dim) - self.shared_reward_normalizer = extra[1] - self.reward_normalizer = StaticNormalizer(1) - - def episode(self, deterministic=False): - config = self.config - self.state_normalizer.offline_stats.load(self.shared_state_normalizer.offline_stats) - self.reward_normalizer.offline_stats.load(self.shared_reward_normalizer.offline_stats) - state = self.task.reset() - state = self.state_normalizer(state) - steps = 0 - total_reward = 0 - pending = [] - while not config.stop_signal.value: - mean, std, log_std = self.worker_network.actor.predict(np.stack([state])) - value = self.worker_network.critic.predict(np.stack([state])) - action = self.policy.sample(mean.data.numpy().flatten(), - std.data.numpy().flatten(), - False) - action = self.config.action_shift_fn(action) - next_state, reward, terminal, _ = self.task.step(action) - next_state = self.state_normalizer(next_state) - - steps += 1 - total_reward += reward - reward = self.reward_normalizer(reward) - - if deterministic: - if terminal: - break - state = next_state - continue - - pending.append([mean, std, log_std, value, action, reward]) - with config.steps_lock: - config.total_steps.value += 1 - - if terminal or len(pending) >= config.update_interval: - critic_loss = 0 - actor_loss = 0 - if terminal: - R = torch.FloatTensor([[0]]) - else: - R = self.worker_network.critic(np.stack([next_state])).data - GAE = torch.FloatTensor([[0]]) - for i in reversed(range(len(pending))): - mean, std, log_std, value, action, reward = pending[i] - if i == len(pending) - 1: - delta = reward + config.discount * R - value.data - else: - delta = reward + pending[i + 1][3].data - value.data - GAE = config.discount * config.gae_tau * GAE + delta - - action = Variable(torch.FloatTensor([action])) - log_density = self.worker_network.actor.log_density(action, mean, log_std, std) - actor_loss += -torch.sum(log_density) * Variable(GAE) - if config.entropy_weight: - actor_loss += -config.entropy_weight * self.worker_network.actor.entropy(std) - - R = reward + config.discount * R - critic_loss += 0.5 * (Variable(R) - value).pow(2) - - pending = [] - self.worker_network.zero_grad() - self.actor_opt.zero_grad() - self.critic_opt.zero_grad() - actor_loss.backward() - critic_loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - sync_grad(self.learning_network, self.worker_network) - self.actor_opt.step() - self.critic_opt.step() - self.worker_network.load_state_dict(self.learning_network.state_dict()) - - if terminal: - break - state = next_state - - self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats) - self.state_normalizer.online_stats.zero() - - self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats) - self.reward_normalizer.online_stats.zero() - - return steps, total_reward \ No newline at end of file diff --git a/async_worker/dpg.py b/async_worker/dpg.py deleted file mode 100644 index 6187904..0000000 --- a/async_worker/dpg.py +++ /dev/null @@ -1,118 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### - -import numpy as np -import torch.multiprocessing as mp -from network import * -from utils import * -from component import * -from async_worker import * -import pickle -import os -import time - -class DeterministicPolicyGradient: - def __init__(self, config, shared_network, extra): - self.config = config - self.task = config.task_fn() - - self.shared_network = shared_network - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(self.shared_network.state_dict()) - self.target_network = config.network_fn() - self.target_network.load_state_dict(self.worker_network.state_dict()) - self.actor_opt = config.actor_optimizer_fn(self.shared_network.actor.parameters()) - self.critic_opt = config.critic_optimizer_fn(self.shared_network.critic.parameters()) - - self.random_process = config.random_process_fn() - self.criterion = nn.MSELoss() - - self.shared_state_normalizer, self.shared_reward_normalizer, self.replay = extra - self.state_normalizer = StaticNormalizer(self.task.state_dim) - self.reward_normalizer = StaticNormalizer(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) + - param.data * self.config.target_network_mix) - - def episode(self, deterministic=False): - self.random_process.reset_states() - state = self.task.reset() - state = self.state_normalizer(state) - - config = self.config - actor = self.worker_network.actor - critic = self.worker_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])).flatten() - if not deterministic: - action += self.random_process.sample() - next_state, reward, done, info = self.task.step(action) - next_state = self.state_normalizer(next_state) - total_reward += reward - reward = self.reward_normalizer(reward) - - if not deterministic: - self.replay.feed([state, action, reward, next_state, int(done)]) - with config.steps_lock: - config.total_steps.value += 1 - - steps += 1 - state = next_state - - if done: - 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)) - 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() - q = critic.predict(states, actions) - critic_loss = self.criterion(q, q_next) - - critic.zero_grad() - self.critic_opt.zero_grad() - critic_loss.backward() - with config.network_lock: - sync_grad(self.shared_network.critic, critic) - self.critic_opt.step() - - actions = actor.predict(states, False) - var_actions = Variable(actions.data, requires_grad=True) - q = critic.predict(states, var_actions) - q.backward(torch.ones(q.size())) - - actor.zero_grad() - self.actor_opt.zero_grad() - actions.backward(-var_actions.grad.data) - with config.network_lock: - sync_grad(self.shared_network.actor, actor) - self.actor_opt.step() - - self.worker_network.load_state_dict(self.shared_network.state_dict()) - - self.soft_update(self.target_network, self.worker_network) - - self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats) - self.state_normalizer.online_stats.zero() - - self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats) - self.reward_normalizer.online_stats.zero() - - return steps, total_reward diff --git a/async_worker/n_step_q.py b/async_worker/n_step_q.py deleted file mode 100644 index a13def5..0000000 --- a/async_worker/n_step_q.py +++ /dev/null @@ -1,79 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable -import torch.nn as nn -from utils import * - -class NStepQLearning: - def __init__(self, config, learning_network, target_network): - self.config = config - self.optimizer = config.optimizer_fn(learning_network.parameters()) - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(learning_network.state_dict()) - self.task = config.task_fn() - self.policy = config.policy_fn() - self.learning_network = learning_network - self.target_network = target_network - - def episode(self, deterministic=False): - config = self.config - state = self.task.reset() - steps = 0 - total_reward = 0 - pending = [] - while not config.stop_signal.value: - q = self.worker_network.predict(np.stack([state])) - action = self.policy.sample(q.data.numpy().flatten(), deterministic) - next_state, reward, terminal, _ = self.task.step(action) - - steps += 1 - total_reward += reward - reward = config.reward_shift_fn(reward) - - if deterministic: - if terminal: - break - state = next_state - continue - - with config.steps_lock: - config.total_steps.value += 1 - pending.append([q, action, reward]) - - if terminal or len(pending) >= config.update_interval: - loss = 0 - if terminal: - R = torch.FloatTensor([0]) - else: - R, _ = self.target_network.predict( - np.stack([next_state])).data.max(1) - - for i in reversed(range(len(pending))): - q, action, reward = pending[i] - R = reward + config.discount * R - q = q.gather(1, Variable(torch.LongTensor([[action]]))).unsqueeze(1) - loss += 0.5 * (Variable(R) - q).pow(2) - - pending = [] - self.worker_network.zero_grad() - self.optimizer.zero_grad() - loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - sync_grad(self.learning_network, self.worker_network) - self.optimizer.step() - self.worker_network.load_state_dict(self.learning_network.state_dict()) - self.worker_network.reset(terminal) - - if terminal: - break - state = next_state - - if config.total_steps.value % config.target_network_update_freq == 0: - self.target_network.load_state_dict(self.learning_network.state_dict()) - - return steps, total_reward \ No newline at end of file diff --git a/async_worker/one_step_q.py b/async_worker/one_step_q.py deleted file mode 100644 index cbf9e14..0000000 --- a/async_worker/one_step_q.py +++ /dev/null @@ -1,76 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable -import torch.nn as nn -from utils import * - -class OneStepQLearning: - def __init__(self, config, learning_network, target_network): - self.config = config - self.optimizer = config.optimizer_fn(learning_network.parameters()) - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(learning_network.state_dict()) - self.task = config.task_fn() - self.policy = config.policy_fn() - self.learning_network = learning_network - self.target_network = target_network - - def episode(self, deterministic=False): - config = self.config - state = self.task.reset() - steps = 0 - total_reward = 0 - pending = [] - while not config.stop_signal.value: - q = self.worker_network.predict(np.stack([state])) - action = self.policy.sample(q.data.numpy().flatten(), deterministic) - next_state, reward, terminal, _ = self.task.step(action) - - steps += 1 - total_reward += reward - reward = config.reward_shift_fn(reward) - - if deterministic: - if terminal: - break - state = next_state - continue - - with config.steps_lock: - config.total_steps.value += 1 - pending.append([q, action, reward, next_state]) - - if terminal or len(pending) >= config.update_interval: - loss = 0 - for i in range(len(pending)): - q, action, reward, next_state = pending[i] - q_next, _ = self.target_network.predict(np.stack([next_state])).data.max(1) - if terminal and i == len(pending) - 1: - q_next = torch.FloatTensor([[0]]) - q_next = config.discount * q_next + reward - q = q.gather(1, Variable(torch.LongTensor([[action]]))).unsqueeze(1) - loss += 0.5 * (q - Variable(q_next)).pow(2) - - pending = [] - self.worker_network.zero_grad() - self.optimizer.zero_grad() - loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - sync_grad(self.learning_network, self.worker_network) - self.optimizer.step() - self.worker_network.load_state_dict(self.learning_network.state_dict()) - self.worker_network.reset(terminal) - - if terminal: - break - state = next_state - - if config.total_steps.value % config.target_network_update_freq == 0: - self.target_network.load_state_dict(self.learning_network.state_dict()) - - return steps, total_reward \ No newline at end of file diff --git a/async_worker/one_step_sarsa.py b/async_worker/one_step_sarsa.py deleted file mode 100644 index 7d2a5e8..0000000 --- a/async_worker/one_step_sarsa.py +++ /dev/null @@ -1,83 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### -import numpy as np -import torch -from torch.autograd import Variable -import torch.nn as nn -from utils import * - -class OneStepSarsa: - def __init__(self, config, learning_network, target_network): - self.config = config - self.optimizer = config.optimizer_fn(learning_network.parameters()) - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(learning_network.state_dict()) - self.task = config.task_fn() - self.policy = config.policy_fn() - self.learning_network = learning_network - self.target_network = target_network - - def episode(self, deterministic=False): - config = self.config - state = self.task.reset() - q = self.worker_network.predict(np.stack([state])) - action = self.policy.sample(q.data.numpy().flatten(), deterministic) - steps = 0 - total_reward = 0 - pending = [] - while not config.stop_signal.value: - next_state, reward, terminal, _ = self.task.step(action) - next_q = self.worker_network.predict(np.stack([next_state])) - next_action = self.policy.sample(next_q.data.numpy().flatten(), deterministic) - pending.append([q, action, reward, next_state, next_action]) - - steps += 1 - total_reward += reward - reward = config.reward_shift_fn(reward) - - if deterministic: - if terminal: - break - state = next_state - action = next_action - continue - - with config.steps_lock: - config.total_steps.value += 1 - - if terminal or len(pending) >= config.update_interval: - loss = 0 - for i in range(len(pending)): - q, action, reward, next_state, next_action = pending[i] - q_next = self.target_network.predict(np.stack([next_state])).data - if terminal and i == len(pending) - 1: - q_next = torch.FloatTensor([[0]]) - else: - q_next = q_next.gather(1, torch.LongTensor([[next_action]])) - q_next = config.discount * q_next + reward - q = q.gather(1, Variable(torch.LongTensor([[action]]))) - loss += 0.5 * (q - Variable(q_next)).pow(2) - - pending = [] - self.worker_network.zero_grad() - self.optimizer.zero_grad() - loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - sync_grad(self.learning_network, self.worker_network) - self.optimizer.step() - self.worker_network.load_state_dict(self.learning_network.state_dict()) - self.worker_network.reset(terminal) - - if terminal: - break - else: - q = next_q - action = next_action - - if config.total_steps.value % config.target_network_update_freq == 0: - self.target_network.load_state_dict(self.learning_network.state_dict()) - - return steps, total_reward diff --git a/async_worker/ppo.py b/async_worker/ppo.py deleted file mode 100644 index 5edd8a0..0000000 --- a/async_worker/ppo.py +++ /dev/null @@ -1,159 +0,0 @@ -####################################################################### -# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) # -# Permission given to modify the code as long as you keep this # -# declaration at the top # -####################################################################### - -import numpy as np -import torch.multiprocessing as mp -from network import * -from utils import * -from component import * -from async_worker import * -import pickle -import os -import time - -class ProximalPolicyOptimization: - def __init__(self, config, shared_network, extra): - self.config = config - self.task = config.task_fn() - self.policy = config.policy_fn() - - self.shared_network = shared_network - self.actor_opt = config.actor_optimizer_fn(shared_network.actor.parameters()) - self.critic_opt = config.critic_optimizer_fn(shared_network.critic.parameters()) - - self.worker_network = config.network_fn() - self.worker_network.load_state_dict(shared_network.state_dict()) - - self.shared_state_normalizer = extra[0] - self.state_normalizer = StaticNormalizer(self.task.state_dim) - self.shared_reward_normalizer = extra[1] - self.reward_normalizer = StaticNormalizer(1) - - def episode(self, deterministic=False): - config = self.config - self.state_normalizer.offline_stats.load(self.shared_state_normalizer.offline_stats) - self.reward_normalizer.offline_stats.load(self.shared_reward_normalizer.offline_stats) - - replay = config.replay_fn() - state = self.task.reset() - state = self.state_normalizer(state) - - episode_length = 0 - batched_rewards = 0 - batched_steps = 0 - batched_episode = 0 - - actor_net = self.worker_network.actor - critic_net = self.worker_network.critic - - actor_net_old = config.actor_network_fn() - actor_net_old.load_state_dict(actor_net.state_dict()) - - self.worker_network.load_state_dict(self.shared_network.state_dict()) - - while not replay.full(): - states = [] - actions = [] - rewards = [] - values = [] - returns = [] - advantages = [] - - for i in range(config.rollout_length): - mean, std, log_std = actor_net.predict(np.stack([state])) - value = critic_net.predict(np.stack([state])) - action = self.policy.sample(mean.data.cpu().numpy().flatten(), std.data.cpu().numpy().flatten(), deterministic) - action = self.config.action_shift_fn(action) - states.append(state) - actions.append(action) - values.append(value) - state, reward, done, _ = self.task.step(action) - state = self.state_normalizer(state) - - batched_rewards += reward - batched_steps += 1 - episode_length += 1 - - reward = self.reward_normalizer(reward) - rewards.append(reward) - - if done: - episode_length = 0 - batched_episode += 1 - state = self.task.reset() - state = self.state_normalizer(state) - break - - R = torch.zeros((1, 1)) - if not done: - R = critic_net.predict(np.stack([state])).data - - - values.append(actor_net.variable(R)) - A = actor_net.variable(torch.zeros((1, 1))) - for i in reversed(range(len(rewards))): - R = actor_net.variable([[rewards[i]]]) - ret = R + self.config.discount * values[i + 1] - A = ret - values[i] + self.config.discount * self.config.gae_tau * A - advantages.append(A.detach()) - returns.append(ret.detach()) - advantages = list(reversed(advantages)) - returns = list(reversed(returns)) - replay.feed([states, actions, returns, advantages]) - - batched_rewards /= batched_episode - batched_steps /= batched_episode - - if deterministic: - return batched_steps, batched_rewards - - with config.steps_lock: - config.total_steps.value += replay.memory_size - - self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats) - self.state_normalizer.online_stats.zero() - - self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats) - self.reward_normalizer.online_stats.zero() - - for _ in np.arange(self.config.optimize_epochs): - self.worker_network.load_state_dict(self.shared_network.state_dict()) - - states, actions, returns, advantages = replay.sample() - states = actor_net.variable(np.stack(states)) - actions = actor_net.variable(np.stack(actions)) - returns = torch.cat(returns, 0) - advantages = torch.cat(advantages, 0).squeeze(1) - advantages = (advantages - advantages.mean()) / advantages.std() - - mean_old, std_old, log_std_old = actor_net_old.predict(states) - probs_old = actor_net.log_density(actions, mean_old, log_std_old, std_old) - mean, std, log_std = actor_net.predict(states) - probs = actor_net.log_density(actions, mean, log_std, std) - ratio = (probs - probs_old).exp() - obj = ratio * advantages - obj_clipped = ratio.clamp(1.0 - self.config.ppo_ratio_clip, 1.0 + self.config.ppo_ratio_clip) * advantages - policy_loss = -torch.min(obj, obj_clipped).mean(0) - if config.entropy_weight: - policy_loss += -config.entropy_weight * actor_net.entropy(std) - - v = critic_net.predict(states) - value_loss = 0.5 * (returns - v).pow(2).mean() - actor_net_old.load_state_dict(actor_net.state_dict()) - - self.worker_network.zero_grad() - policy_loss.backward() - value_loss.backward() - nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip) - with config.network_lock: - self.shared_network.zero_grad() - self.actor_opt.zero_grad() - self.critic_opt.zero_grad() - sync_grad(self.shared_network, self.worker_network) - self.actor_opt.step() - self.critic_opt.step() - - return batched_steps, batched_rewards diff --git a/component/atari_wrapper.py b/component/atari_wrapper.py index f626b88..2b575fc 100644 --- a/component/atari_wrapper.py +++ b/component/atari_wrapper.py @@ -142,14 +142,6 @@ class SkipEnv(gym.Wrapper): obs = self.env.reset() return obs -class ClipRewardEnv(gym.RewardWrapper): - def __init__(self, env): - gym.RewardWrapper.__init__(self, env) - - def reward(self, reward): - """Bin reward to {+1, 0, -1} by its sign.""" - return np.sign(reward) - class WarpFrame(gym.ObservationWrapper): def __init__(self, env): """Warp frames to 84x84 as done in the Nature paper and later work.""" @@ -259,7 +251,7 @@ def make_atari(env_id, frame_skip=4): env = MaxAndSkipEnv(env, skip=4) return env -def wrap_deepmind(env, episode_life=True, clip_rewards=True, history_length=0): +def wrap_deepmind(env, episode_life=True, history_length=0): """Configure environment for DeepMind-style Atari. """ if episode_life: @@ -267,8 +259,6 @@ def wrap_deepmind(env, episode_life=True, clip_rewards=True, history_length=0): if 'FIRE' in env.unwrapped.get_action_meanings(): env = FireResetEnv(env) env = WarpFrame(env) - if clip_rewards: - env = ClipRewardEnv(env) env = WrapPyTorch(env) if history_length: env = StackFrame(env, history_length) diff --git a/component/task.py b/component/task.py index 8a878b0..da52aad 100644 --- a/component/task.py +++ b/component/task.py @@ -24,9 +24,6 @@ class BasicTask: state = self.env.reset() return state - def normalize_state(self, state): - return state - def step(self, action): next_state, reward, done, info = self.env.step(action) self.steps += 1 diff --git a/main.py b/main.py index 9c35374..87508ea 100644 --- a/main.py +++ b/main.py @@ -106,6 +106,8 @@ def dqn_pixel_atari(name): # config.network_fn = lambda state_dim, action_dim: DuelingConvNet(config.history_length, action_dim) config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.state_normalizer = ImageNormalizer() + config.reward_normalizer = SignNormalizer() config.discount = 0.99 config.target_network_update_freq = 10000 config.exploration_steps= 50000 @@ -124,6 +126,8 @@ def a2c_pixel_atari(name): config.network_fn = lambda state_dim, action_dim: ActorCriticConvNet( config.history_length, action_dim, gpu=3) config.policy_fn = SamplePolicy + config.state_normalizer = ImageNormalizer() + config.reward_normalizer = SignNormalizer() config.discount = 0.99 config.use_gae = False config.gae_tau = 0.97 @@ -143,6 +147,8 @@ def categorical_dqn_pixel_atari(name): config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) config.discount = 0.99 + config.state_normalizer = ImageNormalizer() + config.reward_normalizer = SignNormalizer() config.target_network_update_freq = 10000 config.exploration_steps= 50000 config.logger = Logger('./log', logger) @@ -161,6 +167,8 @@ def quantile_regression_dqn_pixel_atari(name): config.network_fn = lambda state_dim, action_dim: QuantileConvNet(config.history_length, action_dim, config.num_quantiles, gpu=0) config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.01) config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.state_normalizer = ImageNormalizer() + config.reward_normalizer = SignNormalizer() config.discount = 0.99 config.target_network_update_freq = 10000 config.exploration_steps= 50000 @@ -179,6 +187,8 @@ def n_step_dqn_pixel_atari(name): config.optimizer_fn = lambda params: torch.optim.RMSprop(params, lr=0.00025, alpha=0.95, eps=0.01) config.network_fn = lambda state_dim, action_dim: ConvNet(config.history_length, action_dim, gpu=0) config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) + config.state_normalizer = ImageNormalizer() + config.reward_normalizer = SignNormalizer() config.discount = 0.99 config.target_network_update_freq = 10000 config.rollout_length = 5 @@ -193,7 +203,8 @@ def dqn_ram_atari(name): config.network_fn = lambda state_dim, action_dim: FCNet(state_dim, 64, action_dim, gpu=2) config.policy_fn = lambda: GreedyPolicy(epsilon=0.1, final_step=1000000, min_epsilon=0.1) config.replay_fn = lambda: Replay(memory_size=100000, batch_size=32, dtype=np.uint8) - config.reward_shift_fn = lambda r: np.sign(r) + config.state_normalizer = RescaleNormalizer(1.0 / 128) + config.reward_normalizer = SignNormalizer() config.discount = 0.99 config.target_network_update_freq = 10000 config.max_episode_length = 0 @@ -216,6 +227,7 @@ def ppo_continuous(): config.critic_network_fn = lambda state_dim, action_dim: GaussianCriticNet(state_dim) config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001) config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001) + config.state_normalizer = RunningStatsNormalizer() config.discount = 0.99 config.use_gae = True config.gae_tau = 0.97 @@ -228,19 +240,17 @@ def ppo_continuous(): def ddpg_continuous(): config = Config() - config.task_fn = lambda: Pendulum() + log_dir = get_default_log_dir(ddpg_continuous.__name__) + config.task_fn = lambda: Pendulum(log_dir=log_dir) # config.task_fn = lambda: Roboschool('RoboschoolInvertedPendulum-v1') # config.task_fn = lambda: Roboschool('RoboschoolReacher-v1') # config.task_fn = lambda: Roboschool('RoboschoolHopper-v1') # config.task_fn = lambda: Roboschool('RoboschoolAnt-v1') # config.task_fn = lambda: Roboschool('RoboschoolWalker2d-v1') - actor_network_fn = lambda state_dim, action_dim: DeterministicActorNet(state_dim, action_dim) - critic_network_fn = lambda state_dim, action_dim: DeterministicCriticNet(state_dim, action_dim) - config.network_fn = lambda state_dim, action_dim: \ - DisjointActorCriticNet(state_dim, action_dim, actor_network_fn, critic_network_fn) + config.actor_network_fn = lambda state_dim, action_dim: DeterministicActorNet(state_dim, action_dim) + config.critic_network_fn = lambda state_dim, action_dim: DeterministicCriticNet(state_dim, action_dim) config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, lr=1e-4) - config.critic_optimizer_fn =\ - lambda params: torch.optim.Adam(params, lr=1e-4) + config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, lr=1e-4) config.replay_fn = lambda: HighDimActionReplay(memory_size=1000000, batch_size=64) config.discount = 0.99 config.random_process_fn = \ @@ -273,7 +283,7 @@ if __name__ == '__main__': # n_step_dqn_pixel_atari('BreakoutNoFrameskip-v4') # dqn_ram_atari('Breakout-ramNoFrameskip-v4') - ddpg_continuous() + # ddpg_continuous() # ppo_continuous() # acvp.train('PongNoFrameskip-v4') diff --git a/network/continuous_action_network.py b/network/continuous_action_network.py index cfe9c11..c1f6277 100644 --- a/network/continuous_action_network.py +++ b/network/continuous_action_network.py @@ -90,7 +90,6 @@ class GaussianActorNet(nn.Module, BasicNet): action_scale=1, action_gate=F.tanh, gpu=-1, - # unit_std=True, hidden_size=64, non_linear=F.tanh): super(GaussianActorNet, self).__init__() @@ -161,7 +160,3 @@ class DisjointActorCriticNet: def zero_grad(self): self.actor.zero_grad() self.critic.zero_grad() - - def train(self): - self.actor.train() - self.critic.train() diff --git a/utils/config.py b/utils/config.py index d3b0038..5c23b5c 100644 --- a/utils/config.py +++ b/utils/config.py @@ -3,6 +3,7 @@ # Permission given to modify the code as long as you keep this # # declaration at the top # ####################################################################### +from .normalizer import * class Config: def __init__(self): @@ -22,22 +23,18 @@ class Config: self.exploration_steps = 0 self.logger = None self.history_length = 1 - self.test_interval = 0 - self.test_repetitions = 50 self.double_q = False self.tag = 'vanilla' self.num_workers = 1 - self.worker = None self.update_interval = 1 - self.gradient_clip = 40 + self.gradient_clip = 0.5 self.entropy_weight = 0.01 self.use_gae = False self.gae_tau = 1.0 self.noise_decay_interval = 0 self.target_network_mix = 0.001 - self.action_shift_fn = lambda a: a - self.reward_shift_fn = lambda r: r - self.reward_weight = 1 + self.state_normalizer = RescaleNormalizer() + self.reward_normalizer = RescaleNormalizer() self.hybrid_reward = False self.episode_limit = 0 self.min_memory_size = 200 diff --git a/utils/misc.py b/utils/misc.py index 55ceaa6..883dedb 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -16,7 +16,6 @@ def run_episodes(agent): ep = 0 rewards = [] steps = [] - avg_test_rewards = [] agent_type = agent.__class__.__name__ while True: ep += 1 @@ -38,25 +37,8 @@ def run_episodes(agent): if config.max_steps and agent.total_steps > config.max_steps: break - if config.test_interval and ep % config.test_interval == 0: - config.logger.info('Testing...') - agent.save('data/%s-%s-model-%s.bin' % (agent_type, config.tag, agent.task.name)) - test_rewards = [] - for _ in range(config.test_repetitions): - test_rewards.append(agent.episode(True)[0]) - avg_reward = np.mean(test_rewards) - avg_test_rewards.append(avg_reward) - config.logger.info('Avg reward %f(%f)' % ( - avg_reward, np.std(test_rewards) / np.sqrt(config.test_repetitions))) - with open('data/%s-%s-all-stats-%s.bin' % (agent_type, config.tag, agent.task.name), 'wb') as f: - pickle.dump({'rewards': rewards, - 'steps': steps, - 'test_rewards': avg_test_rewards}, f) - if avg_reward > config.success_threshold: - break - agent.close() - return steps, rewards, avg_test_rewards + return steps, rewards def run_iterations(agent): config = agent.config @@ -79,11 +61,6 @@ def run_iterations(agent): pickle.dump({'rewards': rewards, 'steps': steps}, f) agent.save('data/%s-%s-model-%s.bin' % (agent_name, config.tag, agent.task.name)) - if config.test_interval and iteration % config.test_interval == 0: - test_rewards, test_steps = agent.evaluate() - config.logger.info('total steps %d, test reward %f, test steps %d' % ( - agent.total_steps, test_rewards, test_steps - )) iteration += 1 if config.max_steps and agent.total_steps >= config.max_steps: agent.close() diff --git a/utils/normalizer.py b/utils/normalizer.py index 99cb98f..f28d68f 100644 --- a/utils/normalizer.py +++ b/utils/normalizer.py @@ -6,16 +6,22 @@ import torch import numpy as np -class Normalizer: - def __init__(self, x_size): +class RunningStatsNormalizer: + def __init__(self): + self.needs_reset = True + + def reset(self, x_size): self.m = np.zeros(x_size) self.v = np.zeros(x_size) - self.n = 1.0 + self.n = 0.0 + self.needs_reset = False def __call__(self, x): if np.isscalar(x) or len(x.shape) == 1: + if self.needs_reset: self.reset(1) return self.nomalize_single(x) elif len(x.shape) == 2: + if self.needs_reset: self.reset(x.shape[1]) new_x = np.zeros(x.shape) for i in range(x.shape[0]): new_x[i] = self.nomalize_single(x[i]) @@ -38,79 +44,17 @@ class Normalizer: x = np.asscalar(x) return x -class StaticNormalizer: - def __init__(self, o_size): - self.offline_stats = SharedStats(o_size) - self.online_stats = SharedStats(o_size) +class RescaleNormalizer: + def __init__(self, coef=1.0): + self.coef = coef - def __call__(self, o_): - if np.isscalar(o_): - o = torch.FloatTensor([o_]) - else: - o = torch.FloatTensor(o_) - self.online_stats.feed(o) - if self.offline_stats.n[0] == 0: - return o_ - std = (self.offline_stats.v + 1e-6) ** .5 - o = (o - self.offline_stats.m) / std - o = o.numpy() - if np.isscalar(o_): - o = np.asscalar(o) - else: - o = o.reshape(o_.shape) - return o - - def state_dict(self): - return self.offline_stats.state_dict() + def __call__(self, x): + return self.coef * x - def load_state_dict(self, saved): - self.offline_stats.load_state_dict(saved) +class ImageNormalizer(RescaleNormalizer): + def __init__(self): + RescaleNormalizer.__init__(self, 1.0 / 255) -class SharedStats: - def __init__(self, o_size): - self.m = torch.zeros(o_size) - self.v = torch.zeros(o_size) - self.n = torch.zeros(1) - self.m.share_memory_() - self.v.share_memory_() - self.n.share_memory_() - - def feed(self, o): - n = self.n[0] - new_m = self.m * (n / (n + 1)) + o / (n + 1) - self.v.copy_(self.v * (n / (n + 1)) + (o - self.m) * (o - new_m) / (n + 1)) - self.m.copy_(new_m) - self.n.add_(1) - - def zero(self): - self.m.zero_() - self.v.zero_() - self.n.zero_() - - def load(self, stats): - self.m.copy_(stats.m) - self.v.copy_(stats.v) - self.n.copy_(stats.n) - - def merge(self, B): - A = self - n_A = self.n[0] - n_B = B.n[0] - n = n_A + n_B - delta = B.m - A.m - m = A.m + delta * n_B / n - v = A.v * n_A + B.v * n_B + delta * delta * n_A * n_B / n - v /= n - self.m.copy_(m) - self.v.copy_(v) - self.n.add_(B.n) - - def state_dict(self): - return {'m': self.m.numpy(), - 'v': self.v.numpy(), - 'n': self.n.numpy()} - - def load_state_dict(self, saved): - self.m = torch.FloatTensor(saved['m']) - self.v = torch.FloatTensor(saved['v']) - self.n = torch.FloatTensor(saved['n']) +class SignNormalizer: + def __call__(self, x): + return np.sign(x) \ No newline at end of file