Upgrade to PyTorch v0.2.0

agent/A2C_agent.py

@@ -1,99 +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                                              #
-#######################################################################
-
-from network import *
-from component import *
-from utils import *
-import numpy as np
-import time
-import os
-import pickle
-import torch
-
-class A2CAgent:
-    def __init__(self, config):
-        self.config = config
-        self.learning_network = config.network_fn()
-        self.optimizer = config.optimizer_fn(self.learning_network.parameters())
-        self.task = config.task_fn()
-        self.replay = config.replay_fn()
-        self.policy = config.policy_fn()
-        self.total_steps = 0
-
-    def episode(self, deterministic=False):
-        state = self.task.reset()
-        total_reward = 0.0
-        steps = 0
-        while True:
-            prob = self.learning_network.predict(np.stack([state]), True)
-            action = self.policy.sample(prob, deterministic=deterministic)
-            next_state, reward, done, info = self.task.step(action)
-            done = (done or (self.config.max_episode_length and steps > self.config.max_episode_length))
-            if not deterministic:
-                self.replay.feed([state, action, reward, next_state, int(done)])
-                self.total_steps += 1
-            total_reward += np.sum(reward * self.config.reward_weight)
-            steps += 1
-            state = next_state
-            if done:
-                break
-            if not deterministic and self.total_steps > self.config.min_memory_size:
-                experiences = self.replay.sample()
-                states, actions, rewards, next_states, terminals = experiences
-                prob, log_prob, value = self.learning_network.predict(states, False)
-                _, _, v_next = self.learning_network.predict(next_states, False)
-                terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
-                rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
-                actions = self.learning_network.to_torch_variable(actions, 'int64').unsqueeze(1)
-                target = rewards + self.config.discount * v_next * (1 - terminals)
-                target = target.detach()
-                advantage = target - value
-                value_loss = 0.5 * advantage.pow(2).mean()
-                policy_loss = -(log_prob.gather(1, actions) * Variable(advantage.data)).mean()
-                kl_loss = (prob * log_prob).sum(1).mean()
-
-                self.optimizer.zero_grad()
-                (value_loss + policy_loss + self.config.entropy_weight * kl_loss).backward()
-                torch.nn.utils.clip_grad_norm(self.learning_network.parameters(), self.config.gradient_clip)
-                self.optimizer.step()
-
-        return total_reward, steps
-
-    def run(self):
-        window_size = 100
-        ep = 0
-        rewards = []
-        steps = []
-        avg_test_rewards = []
-        while True:
-            ep += 1
-            reward, step = self.episode()
-            rewards.append(reward)
-            steps.append(step)
-            avg_reward = np.mean(rewards[-window_size:])
-            self.config.logger.info('episode %d, reward %f, avg reward %f, total steps %d, episode step %d' % (
-                ep, reward, avg_reward, self.total_steps, step))
-
-            if self.config.episode_limit and ep > self.config.episode_limit:
-                return rewards, steps, avg_test_rewards
-
-            if self.config.test_interval and ep % self.config.test_interval == 0:
-                self.config.logger.info('Testing...')
-                with open('data/%s-dqn-model-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
-                    pickle.dump(self.learning_network.state_dict(), f)
-                test_rewards = []
-                for _ in range(self.config.test_repetitions):
-                    reward, step = self.episode(True)
-                    test_rewards.append(reward)
-                avg_reward = np.mean(test_rewards)
-                avg_test_rewards.append(avg_reward)
-                self.config.logger.info('Avg reward %f(%f)' % (
-                    avg_reward, np.std(test_rewards) / np.sqrt(self.config.test_repetitions)))
-                with open('data/%sdqn-statistics-%s.bin' % (self.config.tag, self.task.name), 'wb') as f:
-                    pickle.dump({'rewards': rewards,
-                                 'test_rewards': avg_test_rewards}, f)
-                if avg_reward > self.task.success_threshold:
-                    break

agent/DQN_agent.py

@@ -71,10 +71,10 @@ class DQNAgent:
                             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.cat(target, dim=1).detach()
+                    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.expand_as(target))
+                    target = self.config.discount * target * (1 - terminals)
                     target.add_(rewards)
                     q = self.learning_network.predict(states, True)
                     q_action = []
@@ -87,16 +87,16 @@ class DQNAgent:
                     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)
+                        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).unsqueeze(1)
-                    rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(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)
+                    q = q.gather(1, actions).squeeze(1)
                     loss = self.learning_network.criterion(q, q_next)
                 self.learning_network.zero_grad()
                 loss.backward()

agent/MSDQN_agent.py

@@ -1,101 +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                                              #
-#######################################################################
-
-from network import *
-from component import *
-from utils import *
-import numpy as np
-import time
-import os
-import pickle
-import torch
-
-# This HRA DQN with removing irrelevant features
-class MSDQNAgent:
-    def __init__(self, config):
-        self.config = config
-        self.learning_network = config.network_fn(config.optimizer_fn)
-        self.target_network = config.network_fn(config.optimizer_fn)
-        self.target_network.load_state_dict(self.learning_network.state_dict())
-        self.task = config.task_fn()
-        self.replay = config.replay_fn()
-        self.policy = config.policy_fn()
-        self.total_steps = 0
-
-    def episode(self, deterministic=False):
-        episode_start_time = time.time()
-        state = self.task.reset()
-        total_reward = 0.0
-        steps = 0
-        while True:
-            value = self.learning_network.predict(np.stack([state]), True)
-            value = value.cpu().data.numpy().flatten()
-            if deterministic:
-                action = np.argmax(value)
-            elif self.total_steps < self.config.exploration_steps:
-                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))
-            if not deterministic:
-                self.replay.feed([state, action, reward, next_state, int(done)])
-                self.total_steps += 1
-            total_reward += np.sum(reward * self.config.reward_weight)
-            steps += 1
-            state = next_state
-            if done:
-                break
-            if not deterministic and self.total_steps > self.config.exploration_steps:
-                experiences = self.replay.sample()
-                states, actions, rewards, next_states, terminals = experiences
-                if self.config.hybrid_reward:
-                    q_next = self.target_network.predict(next_states, False)
-                    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.cat(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.expand_as(target))
-                    target.add_(rewards)
-                    q = self.learning_network.predict(states, False)
-                    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)
-                else:
-                    q_next = self.target_network.predict(next_states, True).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)
-                    else:
-                        q_next, _ = q_next.max(1)
-                    terminals = self.learning_network.to_torch_variable(terminals).unsqueeze(1)
-                    rewards = np.sum(rewards * self.config.reward_weight, axis=1)
-                    rewards = self.learning_network.to_torch_variable(rewards).unsqueeze(1)
-                    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, True)
-                    q = q.gather(1, actions)
-                    loss = self.learning_network.criterion(q, q_next)
-                self.learning_network.zero_grad()
-                loss.backward()
-                self.learning_network.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())
-            if not deterministic and self.total_steps > self.config.exploration_steps:
-                self.policy.update_epsilon()
-        episode_time = time.time() - episode_start_time
-        self.config.logger.debug('episode steps %d, episode time %f, time per step %f' %
-                          (steps, episode_time, episode_time / float(steps)))
-        return total_reward, steps

agent/__init__.py

@@ -1,5 +1,3 @@
 from .async_agent import *
 from .DDPG_agent import *
 from .DQN_agent import *
-from .A2C_agent import *
-from .MSDQN_agent import *

async_worker/n_step_q.py

@@ -47,7 +47,7 @@ class NStepQLearning:
             if terminal or len(pending) >= config.update_interval:
                 loss = 0
                 if terminal:
-                    R = torch.FloatTensor([[0]])
+                    R = torch.FloatTensor([0])
                 else:
                     R, _ = self.target_network.predict(
                         np.stack([next_state])).data.max(1)
@@ -55,7 +55,8 @@ class NStepQLearning:
                 for i in reversed(range(len(pending))):
                     q, action, reward = pending[i]
                     R = reward + config.discount * R
-                    loss += 0.5 * (Variable(R) - q.gather(1, Variable(torch.LongTensor([[action]])))).pow(2)
+                    q = q.gather(1, Variable(torch.LongTensor([[action]]))).unsqueeze(1)
+                    loss += 0.5 * (Variable(R) - q).pow(2)
 
                 pending = []
                 self.worker_network.zero_grad()

async_worker/one_step_q.py

@@ -52,7 +52,7 @@ class OneStepQLearning:
                     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]])))
+                    q = q.gather(1, Variable(torch.LongTensor([[action]]))).unsqueeze(1)
                     loss += 0.5 * (q - Variable(q_next)).pow(2)
 
                 pending = []

async_worker/ppo.py

@@ -127,8 +127,8 @@ class ProximalPolicyOptimization:
             states = actor_net.to_torch_variable(np.stack(states))
             actions = actor_net.to_torch_variable(np.stack(actions))
             returns = torch.cat(returns, 0)
-            advantages = torch.cat(advantages, 0)
-            advantages = (advantages - advantages.mean().expand_as(advantages)) / advantages.std().expand_as(advantages)
+            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)

main.py

@@ -19,8 +19,8 @@ def dqn_cart_pole():
     config.history_length = 2
     config.test_interval = 100
     config.test_repetitions = 50
-    # config.double_q = True
-    config.double_q = False
+    config.double_q = True
+    # config.double_q = False
     run_episodes(DQNAgent(config))
 
 def async_cart_pole():
@@ -174,35 +174,10 @@ def hrdqn_fruit():
     config.episode_limit = 5000
     run_episodes(DQNAgent(config))
 
-def hrmsdqn_fruit():
-    config = Config()
-    config.task_fn = lambda: Fruit(hybrid_reward=True, atomic_state=False)
-    config.hybrid_reward = True
-    config.reward_weight = np.ones(10) / 10
-    # config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
-    config.optimizer_fn = lambda params: torch.optim.SGD(params, 0.1, momentum=0.9)
-    config.network_fn = lambda optimizer_fn: FruitMultiStatesFCNet(
-        17, 4, config.reward_weight, optimizer_fn)
-    config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=10000, min_epsilon=0.1)
-    config.replay_fn = lambda: HybridRewardReplay(memory_size=10000, batch_size=15)
-    config.discount = 0.95
-    config.target_network_update_freq = 200
-    config.max_episode_length = 100
-    config.exploration_steps = 200
-    config.logger = Logger('./log', gym.logger)
-    config.history_length = 1
-    config.test_interval = 0
-    config.test_repetitions = 10
-    config.target_type = config.expected_sarsa_target
-    # config.target_type = config.q_target
-    config.double_q = False
-    config.episode_limit = 5000
-    run_episodes(MSDQNAgent(config))
-
 def a3c_continuous():
     config = Config()
-    # config.task_fn = lambda: Pendulum()
-    config.task_fn = lambda: BipedalWalkerHardcore()
+    config.task_fn = lambda: Pendulum()
+    # config.task_fn = lambda: BipedalWalkerHardcore()
     task = config.task_fn()
     config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 0.0001)
     config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
@@ -257,8 +232,8 @@ def dppo_continuous():
 
 def ddpg_continuous():
     config = Config()
-    # config.task_fn = lambda: Pendulum()
-    config.task_fn = lambda: Roboschool('RoboschoolInvertedPendulum-v1')
+    config.task_fn = lambda: Pendulum()
+    # config.task_fn = lambda: Roboschool('RoboschoolInvertedPendulum-v1')
     # config.task_fn = lambda: Roboschool('RoboschoolReacher-v1')
     task = config.task_fn()
     config.actor_network_fn = lambda: DeterministicActorNet(
@@ -288,22 +263,21 @@ if __name__ == '__main__':
     # gym.logger.setLevel(logging.DEBUG)
     gym.logger.setLevel(logging.INFO)
 
-    # dqn_cart_pole()
+    dqn_cart_pole()
     # async_cart_pole()
     # a3c_cart_pole()
     # a3c_continuous()
     # dppo_continuous()
-    ddpg_continuous()
+    # ddpg_continuous()
 
     # dqn_fruit()
     # hrdqn_fruit()
-    # hrmsdqn_fruit()
 
-    # dqn_pixel_atari('PongNoFrameskip-v3')
-    # async_pixel_atari('PongNoFrameskip-v3')
-    # a3c_pixel_atari('PongNoFrameskip-v3')
+    # dqn_pixel_atari('PongNoFrameskip-v4')
+    # async_pixel_atari('PongNoFrameskip-v4')
+    # a3c_pixel_atari('PongNoFrameskip-v4')
 
-    # dqn_pixel_atari('BreakoutNoFrameskip-v3')
-    # async_pixel_atari('BreakoutNoFrameskip-v3')
-    # a3c_pixel_atari('BreakoutNoFrameskip-v3')
+    # dqn_pixel_atari('BreakoutNoFrameskip-v4')
+    # async_pixel_atari('BreakoutNoFrameskip-v4')
+    # a3c_pixel_atari('BreakoutNoFrameskip-v4')
 

utils/run.py

@@ -38,7 +38,7 @@ def run_episodes(agent):
             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))
+                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)' % (