From 6de80627020f64217ac3f6cbd25aa684c4ec903b Mon Sep 17 00:00:00 2001 From: Shangtong Zhang Date: Mon, 12 Feb 2018 14:27:40 -0700 Subject: [PATCH] Code cleanup --- agent/A2C_agent.py | 20 +++----------------- component/task.py | 28 ++++++---------------------- main.py | 6 +++--- network/base_network.py | 8 ++++++++ network/conv_network.py | 8 -------- utils/config.py | 2 +- utils/misc.py | 14 ++++++++++++++ 7 files changed, 35 insertions(+), 51 deletions(-) diff --git a/agent/A2C_agent.py b/agent/A2C_agent.py index 3de9a8b..77506cd 100644 --- a/agent/A2C_agent.py +++ b/agent/A2C_agent.py @@ -23,7 +23,6 @@ class A2CAgent: self.policy = config.policy_fn() self.total_steps = 0 self.states = self.task.reset() - self.episode_rewards = np.zeros(config.num_workers) self.last_episode_rewards = np.zeros(config.num_workers) @@ -48,26 +47,13 @@ class A2CAgent: break return total_rewards, steps - def episode(self, deterministic=False): - config = self.config - for _ in range(config.iteration_log_interval): - self.iteration(deterministic) - config.logger.info('max/min reward %f/%f, policy loss %f, entropy loss %f, value loss %f' % - (np.max(self.last_episode_rewards), np.min(self.last_episode_rewards), - self.policy_loss, self.entropy_loss, self.value_loss)) - return self.last_episode_rewards.mean(), config.rollout_length * config.num_workers * \ - config.iteration_log_interval - - def iteration(self, deterministic=False): - if deterministic: - return self.evaluate() - + def iteration(self): config = self.config rollout = [] states = self.states for i in range(config.rollout_length): prob, log_prob, value = self.network.predict(states) - actions = [self.policy.sample(p, deterministic) for p in prob.data.cpu().numpy()] + 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 @@ -95,7 +81,7 @@ class A2CAgent: actions = self.network.tensor(actions, torch.LongTensor).unsqueeze(1) next_value = rollout[i + 1][2] returns = rewards + config.discount * terminals * returns - if config.no_gae: + if config.use_gae: advantages = returns - value.data else: td_error = rewards + config.discount * terminals * next_value.data - value.data diff --git a/component/task.py b/component/task.py index 1544c34..c292441 100644 --- a/component/task.py +++ b/component/task.py @@ -12,26 +12,21 @@ import sys class BasicTask: def __init__(self, max_steps=sys.maxsize): - self.normalized_state = True self.steps = 0 self.max_steps = max_steps - def normalize_state(self, state): - return state - def reset(self): self.steps = 0 state = self.env.reset() - if self.normalized_state: - return self.normalize_state(state) + 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 done = (done or self.steps >= self.max_steps) - if self.normalized_state: - next_state = self.normalize_state(next_state) return next_state, reward, done, info def random_action(self): @@ -70,24 +65,13 @@ class PixelAtari(BasicTask): if 'FIRE' in env.unwrapped.get_action_meanings(): env = FireResetEnv(env) env = ProcessFrame(env, frame_size) + if normalized_state: + env = NormalizeFrame(env) self.env = StackFrame(env, history_length) self.action_dim = self.env.action_space.n - self.observation_space = self.env.observation_space - self.action_space = self.env.action_space def normalize_state(self, state): - return np.asarray(state, dtype=np.float32) / 255.0 - - def step(self, action): - next_state, reward, done, info = self.env.step(action) - self.steps += 1 - done = (done or self.steps >= self.max_steps) - if done: - self.steps = 0 - next_state = self.env.reset() - if self.normalized_state: - next_state = self.normalize_state(next_state) - return next_state, reward, done, info + return np.asarray(state) / 255.0 class ContinuousMountainCar(BasicTask): name = 'MountainCarContinuous-v0' diff --git a/main.py b/main.py index 4048cb0..d156d3e 100644 --- a/main.py +++ b/main.py @@ -89,7 +89,7 @@ def a2c_cart_pole(): config.gae_tau = 1.0 config.entropy_weight = 0.01 config.rollout_length = 20 - run_episodes(A2CAgent(config)) + run_iterations(A2CAgent(config)) def dqn_pixel_atari(name): config = Config() @@ -176,7 +176,7 @@ def a2c_pixel_atari(name): config.reward_shift_fn = lambda r: np.sign(r) config.policy_fn = SamplePolicy config.discount = 0.99 - config.no_gae = False + config.use_gae = True config.gae_tau = 0.97 config.entropy_weight = 0.01 config.rollout_length = 5 @@ -184,7 +184,7 @@ def a2c_pixel_atari(name): config.iteration_log_interval = 100 config.gradient_clip = 0.5 config.logger = Logger('./log', logger, skip=True) - run_episodes(A2CAgent(config)) + run_iterations(A2CAgent(config)) def a3c_continuous(): config = Config() diff --git a/network/base_network.py b/network/base_network.py index 82ee7a5..52e0501 100644 --- a/network/base_network.py +++ b/network/base_network.py @@ -17,6 +17,7 @@ class BasicNet: gpu = -1 self.gpu = gpu self.LSTM = LSTM + self.init_weights() if self.gpu >= 0: self.cuda(self.gpu) @@ -49,6 +50,13 @@ class BasicNet: self.h = Variable(self.h.data) self.c = Variable(self.c.data) + def init_weights(self): + for layer in self.children(): + relu_gain = nn.init.calculate_gain('relu') + if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear): + nn.init.orthogonal(layer.weight.data, relu_gain) + nn.init.constant(layer.bias.data, 0) + # Base class for value based methods class VanillaNet(BasicNet): def predict(self, x, to_numpy=False): diff --git a/network/conv_network.py b/network/conv_network.py index 29b6948..2f672f3 100644 --- a/network/conv_network.py +++ b/network/conv_network.py @@ -131,16 +131,8 @@ class NatureActorCriticConvNet(nn.Module, ActorCriticNet): self.fc_actor = nn.Linear(512, n_actions) self.fc_critic = nn.Linear(512, 1) - self.init_weights() BasicNet.__init__(self, gpu=gpu) - def init_weights(self): - relu_gain = nn.init.calculate_gain('relu') - for layer in self.children(): - if isinstance(layer, nn.Conv2d) or isinstance(layer, nn.Linear): - nn.init.orthogonal(layer.weight.data, relu_gain) - nn.init.constant(layer.bias.data, 0) - def forward(self, x, _): x = self.variable(x) x = F.relu(self.conv1(x)) diff --git a/utils/config.py b/utils/config.py index 182f704..689475c 100644 --- a/utils/config.py +++ b/utils/config.py @@ -33,7 +33,7 @@ class Config: self.update_interval = 1 self.gradient_clip = 40 self.entropy_weight = 0.01 - self.no_gae = False + self.use_gae = True self.gae_tau = 1.0 self.noise_decay_interval = 0 self.target_network_mix = 0.001 diff --git a/utils/misc.py b/utils/misc.py index 26fe827..04bc96e 100644 --- a/utils/misc.py +++ b/utils/misc.py @@ -56,6 +56,20 @@ def run_episodes(agent): agent.close() return steps, rewards, avg_test_rewards +def run_iterations(agent): + config = agent.config + agent_type = agent.__class__.__name__ + iteration = 0 + while True: + agent.iteration() + if iteration % config.iteration_log_interval == 0: + config.logger.info('total steps %d, mean/max/min reward %f/%f/%f' % ( + agent.total_steps, np.mean(agent.last_episode_rewards), + np.max(agent.last_episode_rewards), + np.min(agent.last_episode_rewards) + )) + iteration += 1 + def sync_grad(target_network, src_network): for param, src_param in zip(target_network.parameters(), src_network.parameters()): param._grad = src_param.grad.clone()