Code cleanup

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

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'

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()

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):

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))

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

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()