Code cleanup

agent/A2C_agent.py

@@ -44,7 +44,7 @@ class A2CAgent:
         steps = 0
         while True:
             prob, _, _ = self.network.predict(np.stack([state]))
-            action = self.policy.sample(prob.data.numpy().flatten(), True)
+            action = self.policy.sample(prob.data.cpu().numpy().flatten(), True)
             state, reward, done, _ = self.evaluator.step(action)
             total_rewards += reward
             steps += 1
@@ -61,7 +61,7 @@ class A2CAgent:
         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.numpy()]
+            actions = [self.policy.sample(p, deterministic) 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

agent/DDPG_agent.py

@@ -40,6 +40,9 @@ class DDPGAgent:
         with open(file_name, 'wb') as f:
             torch.save(self.worker_network.state_dict(), f)
 
+    def close(self):
+        pass
+
     def episode(self, deterministic=False, video_recorder=None):
         self.random_process.reset_states()
         state = self.task.reset()
@@ -61,7 +64,6 @@ class DDPGAgent:
             next_state, reward, done, info = self.task.step(action)
             if video_recorder is not None:
                 video_recorder.capture_frame()
-            done = (done or (config.max_episode_length and steps >= config.max_episode_length))
             next_state = self.state_normalizer(next_state)
             total_reward += reward
             reward = self.reward_normalizer(reward)

agent/DQN_agent.py

@@ -41,8 +41,7 @@ class DQNAgent:
                 action = np.random.randint(0, len(value))
             else:
                 action = self.policy.sample(value)
-            next_state, reward, done, info = self.task.step(action)
-            done = (done or (self.config.max_episode_length and steps > self.config.max_episode_length))
+            next_state, reward, done, _ = self.task.step(action)
             self.history_buffer.pop(0)
             self.history_buffer.append(next_state)
             next_state = np.vstack(self.history_buffer)
@@ -60,41 +59,20 @@ class DQNAgent:
                 states, actions, rewards, next_states, terminals = experiences
                 states = self.task.normalize_state(states)
                 next_states = self.task.normalize_state(next_states)
-                if self.config.hybrid_reward:
-                    q_next = self.target_network.predict(next_states, True)
-                    target = []
-                    for q_next_ in q_next:
-                        if self.config.target_type == self.config.q_target:
-                            target.append(q_next_.detach().max(1)[0])
-                        elif self.config.target_type == self.config.expected_sarsa_target:
-                            target.append(q_next_.detach().mean(1))
-                    target = torch.stack(target, dim=1).detach()
-                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
-                    rewards = self.learning_network.to_torch_variable(rewards)
-                    target = self.config.discount * target * (1 - terminals)
-                    target.add_(rewards)
-                    q = self.learning_network.predict(states, True)
-                    q_action = []
-                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
-                    for q_ in q:
-                        q_action.append(q_.gather(1, actions))
-                    q_action = torch.cat(q_action, dim=1)
-                    loss = self.learning_network.criterion(q_action, target)
+                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)
+                    q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
                 else:
-                    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)
-                        q_next = q_next.gather(1, best_actions.unsqueeze(1)).squeeze(1)
-                    else:
-                        q_next, _ = q_next.max(1)
-                    terminals = self.learning_network.to_torch_variable(terminals)
-                    rewards = self.learning_network.to_torch_variable(rewards)
-                    q_next = self.config.discount * q_next * (1 - terminals)
-                    q_next.add_(rewards)
-                    actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
-                    q = self.learning_network.predict(states, False)
-                    q = q.gather(1, actions).squeeze(1)
-                    loss = self.criterion(q, q_next)
+                    q_next, _ = q_next.max(1)
+                terminals = self.learning_network.to_torch_variable(terminals)
+                rewards = self.learning_network.to_torch_variable(rewards)
+                q_next = self.config.discount * q_next * (1 - terminals)
+                q_next.add_(rewards)
+                actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
+                q = self.learning_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()
@@ -110,3 +88,6 @@ class DQNAgent:
     def save(self, file_name):
         with open(file_name, 'wb') as f:
             torch.save(self.learning_network.state_dict(), f)
+
+    def close(self):
+        pass