Support A2C

agent/A2C_agent.py

@@ -0,0 +1,118 @@
+#######################################################################
+# 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 *
+import pickle
+import os
+import time
+import gym.monitoring
+
+class A2CAgent:
+    def __init__(self, config):
+        self.config = config
+        self.task = config.task_fn()
+        self.evaluator = self.task.task_fn()
+        self.network = config.network_fn()
+        self.optimizer = config.optimizer_fn(self.network.parameters())
+        self.policy = config.policy_fn()
+        self.total_steps = 0
+        self.states = self.task.reset()
+
+        self.episode_counts = np.zeros(config.num_workers)
+        self.episode_rewards = np.zeros(config.num_workers)
+        self.total_rewards = np.zeros(config.num_workers)
+        self.prev_episode_counts = 0.0
+        self.prev_total_rewards = 0.0
+
+    def close(self):
+        self.task.close()
+
+    def save(self, file_name):
+        with open(file_name, 'wb') as f:
+            torch.save(self.network.state_dict(), f)
+
+    def evaluate(self):
+        state = self.evaluator.reset()
+        total_rewards = 0
+        steps = 0
+        while True:
+            prob, _, _ = self.network.predict(np.stack([state]))
+            action = self.policy.sample(prob.data.numpy().flatten(), True)
+            state, reward, done, _ = self.evaluator.step(action)
+            total_rewards += reward
+            steps += 1
+            if done:
+                break
+        return total_rewards, steps
+
+    def episode(self, deterministic=False):
+        if deterministic:
+            return self.evaluate()
+
+        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.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)
+            for i, terminal in enumerate(terminals):
+                if terminals[i]:
+                    next_states[i] = self.task.reset(i)
+                    self.episode_counts[i] += 1
+                    self.total_rewards[i] += self.episode_rewards[i]
+                    self.episode_rewards[i] = 0
+
+            rollout.append([prob, log_prob, value, actions, rewards, 1 - terminals])
+            states = next_states
+
+        self.states = states
+        _, _, pending_value = self.network.predict(states)
+        rollout.append([None, None, pending_value, None, None, None])
+
+        processed_rollout = [None] * (len(rollout) - 1)
+        advantages = self.network.FloatTensor(np.zeros((config.num_workers, 1)))
+        returns = pending_value.data
+        for i in reversed(range(len(rollout) - 1)):
+            prob, log_prob, value, actions, rewards, terminals = rollout[i]
+            terminals = self.network.FloatTensor(terminals).unsqueeze(1)
+            rewards = self.network.FloatTensor(rewards).unsqueeze(1)
+            actions = self.network.LongTensor(actions).unsqueeze(1)
+            next_value = rollout[i + 1][2]
+            returns = rewards + terminals * config.discount * returns
+            td_error = rewards + config.discount * terminals * next_value.data - value.data
+            advantages = advantages * config.gae_tau * config.discount * terminals + td_error
+            processed_rollout[i] = [prob, log_prob, value, actions, returns, advantages]
+
+        prob, log_prob, value, actions, returns, advantages = map(lambda x: torch.cat(x, dim=0), zip(*processed_rollout))
+        policy_loss = -log_prob.gather(1, Variable(actions)) * Variable(advantages)
+        policy_loss += config.entropy_weight * torch.sum(prob * log_prob, dim=1, keepdim=True)
+        value_loss = 0.5 * (Variable(returns) - value).pow(2)
+
+        self.optimizer.zero_grad()
+        (policy_loss + value_loss).mean().backward()
+        nn.utils.clip_grad_norm(self.network.parameters(), config.gradient_clip)
+        self.optimizer.step()
+
+        steps = config.rollout_length * config.num_workers
+        self.total_steps += steps
+        new_episode_counts = np.sum(self.episode_counts)
+        new_total_rewards = np.sum(self.total_rewards)
+        avg_reward = (new_total_rewards - self.prev_total_rewards) / \
+                     (new_episode_counts - self.prev_episode_counts + 1e-5)
+        self.prev_total_rewards = new_total_rewards
+        self.prev_episode_counts = new_episode_counts
+        return avg_reward, steps
+
+
+

agent/__init__.py

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

component/task.py

@@ -7,15 +7,20 @@ import gym
 import sys
 import numpy as np
 from .atari_wrapper import *
+import torch.multiprocessing as mp
+import sys
 
 class BasicTask:
-    def __init__(self):
+    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)
@@ -23,6 +28,8 @@ class BasicTask:
 
     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, np.sign(reward), done, info
@@ -34,8 +41,8 @@ class MountainCar(BasicTask):
     name = 'MountainCar-v0'
     success_threshold = -110
 
-    def __init__(self):
-        BasicTask.__init__(self)
+    def __init__(self, max_steps=200):
+        BasicTask.__init__(self, max_steps)
         self.env = gym.make(self.name)
         self.env._max_episode_steps = sys.maxsize
 
@@ -43,8 +50,8 @@ class CartPole(BasicTask):
     name = 'CartPole-v0'
     success_threshold = 195
 
-    def __init__(self):
-        BasicTask.__init__(self)
+    def __init__(self, max_steps=200):
+        BasicTask.__init__(self, max_steps)
         self.env = gym.make(self.name)
         self.env._max_episode_steps = sys.maxsize
 
@@ -182,121 +189,50 @@ class Roboschool(BasicTask):
         next_state, reward, done, info = self.env.step(action)
         return next_state, reward, done, info
 
-class Fruit(BasicTask):
-    def __init__(self, hybrid_reward=False, pseudo_reward=False, atomic_state=True):
-        self.hybrid_reward = hybrid_reward
-        self.atomic_state = atomic_state
-        self.pseudo_reward = pseudo_reward
-        self.name = "Fruit"
-        self.success_threshold = 5
-        self.width = 10
-        self.height = 10
-        self.possible_fruits = 10
-        self.actual_fruits = 5
-        xs = np.random.randint(0, self.width, size=self.possible_fruits)
-        ys = np.random.randint(0, self.height, size=self.possible_fruits)
-        self.possible_locations = list(zip(xs, ys))
-        self.x = 0
-        self.y = 0
-        self.indices = np.arange(self.possible_fruits)
-        self.taken = []
-        self.remaining_fruits = 0
-
-    def get_nearest(self):
-        def distance(i):
-            x, y = self.possible_locations[i]
-            return np.abs(self.x - x) + np.abs(self.y - y)
-        pool = []
-        for i in range(self.possible_fruits):
-            if not self.taken[i]:
-                pool.append([i, distance(i)])
-        pool = sorted(pool, key=lambda x:x[1])
-        return pool[0][0]
-
-    def encode_pos(self, x, y):
-        return '{:04b}'.format(x) + '{:04b}'.format(y)
-
-    def encode_atomic_state(self):
-        offset = 8 * self.possible_fruits
-        state = np.copy(self.base_state)
-        str = self.encode_pos(self.x, self.y)
-        for i in range(len(str)):
-            state[offset + i] = int(str[i])
-        offset += 8
-        for i in range(len(self.taken)):
-            state[offset + i] = self.taken[i]
-        return state
-
-    def encode_decomposed_state(self):
-        state_size = (4 + 4) * 2 + 1
-        base_state = np.zeros(state_size)
-        str = self.encode_pos(self.x, self.y)
-        for i in range(len(str)):
-            base_state[i] = int(str[i])
-        states = []
-        for i in range(self.possible_fruits):
-            states.append(np.copy(base_state))
-            str = self.encode_pos(*self.possible_locations[i])
-            for j in range(len(str)):
-                states[-1][8 + j] = int(str[j])
-            states[-1][-1] = self.taken[i]
-        return np.asarray(states)
-
-    def encode_state(self):
-        if self.atomic_state:
-            return self.encode_atomic_state()
-        return self.encode_decomposed_state()
-
-    def reset(self):
-        self.x = np.random.randint(0, self.width)
-        self.y = np.random.randint(0, self.height)
-        np.random.shuffle(self.indices)
-        self.taken = np.ones(self.possible_fruits, dtype=np.bool)
-        self.taken[self.indices[: self.actual_fruits]] = False
-        self.remaining_fruits = self.actual_fruits
-        state_size = (4 + 4) * (self.possible_fruits + 1) + self.possible_fruits
-        self.base_state = np.zeros(state_size)
-        offset = 0
-        for x, y in self.possible_locations:
-            str = self.encode_pos(x, y)
-            for i in range(len(str)):
-                self.base_state[offset + i] = int(str[i])
-            offset += 8
-        return self.encode_state()
-
-    def step(self, action):
-        # action = action[0]
-        if action == 0:
-            self.x -= 1
-        elif action == 1:
-            self.x += 1
-        elif action == 2:
-            self.y -= 1
-        elif action == 3:
-            self.y += 1
+def sub_task(parent_pipe, pipe, task_fn):
+    parent_pipe.close()
+    task = task_fn()
+    while True:
+        op, data = pipe.recv()
+        if op == 'step':
+            pipe.send(task.step(data))
+        elif op == 'reset':
+            pipe.send(task.reset())
+        elif op == 'exit':
+            pipe.close()
+            return
         else:
-            assert False
-        self.x = min(max(self.x, 0), self.width - 1)
-        self.y = min(max(self.y, 0), self.height - 1)
-        try:
-            pos = self.possible_locations.index((self.x, self.y))
-        except ValueError:
-            pos = -1
-        if self.hybrid_reward:
-            reward = np.zeros(self.possible_fruits)
-            if pos >= 0 and not self.taken[pos]:
-                reward[pos] = 10
-                self.taken[pos] = True
-                self.remaining_fruits -= 1
-            if self.pseudo_reward:
-                pseudo_reward = np.zeros(self.possible_fruits)
-                if pos >= 0:
-                    pseudo_reward[pos] = 1
-                reward = (reward, pseudo_reward)
+            assert False, 'Unknown Operation'
+
+class ParallelizedTask:
+    def __init__(self, task_fn, num_workers):
+        self.task_fn = task_fn
+        self.task = task_fn()
+        self.name = self.task.name
+        self.pipes, worker_pipes = zip(*[mp.Pipe() for _ in range(num_workers)])
+        args = [(p, wp, task_fn) for p, wp in zip(self.pipes, worker_pipes)]
+        self.workers = [mp.Process(target=sub_task, args=arg) for arg in args]
+        for p in self.workers: p.start()
+        for p in worker_pipes: p.close()
+
+    def step(self, actions):
+        for pipe, action in zip(self.pipes, actions):
+            pipe.send(('step', action))
+        results = [p.recv() for p in self.pipes]
+        results = map(lambda x: np.stack(x), zip(*results))
+        return results
+
+    def reset(self, i=None):
+        if i is None:
+            for pipe in self.pipes:
+                pipe.send(('reset', None))
+            results = [p.recv() for p in self.pipes]
         else:
-            reward = 0.0
-            if pos >= 0 and not self.taken[pos]:
-                reward = 1.0
-                self.taken[pos] = True
-                self.remaining_fruits -= 1
-        return self.encode_state(), reward, not self.remaining_fruits, self.taken
+            self.pipes[i].send(('reset', None))
+            results = self.pipes[i].recv()
+        return np.stack(results)
+
+    def close(self):
+        for pipe in self.pipes:
+            pipe.send(('exit', None))
+        for p in self.workers: p.join()

main.py

@@ -69,6 +69,24 @@ def a3c_cart_pole():
     agent = AsyncAgent(config)
     agent.run()
 
+def a2c_cart_pole():
+    config = Config()
+    task_fn = lambda: CartPole(max_steps=200)
+    config.num_workers = 3
+    config.task_fn = lambda: ParallelizedTask(task_fn, config.num_workers)
+    config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
+    config.network_fn = lambda: ActorCriticFCNet(4, 2)
+    config.policy_fn = SamplePolicy
+    config.discount = 0.99
+    config.test_interval = 20
+    config.test_repetitions = 10
+    config.logger = Logger('./log', logger)
+    config.gae_tau = 1.0
+    config.entropy_weight = 0.01
+    config.rollout_length = 50
+    config.success_threshold = 195
+    run_episodes(A2CAgent(config))
+
 def dqn_pixel_atari(name):
     config = Config()
     config.history_length = 4
@@ -317,9 +335,10 @@ if __name__ == '__main__':
     # logger.setLevel(logging.DEBUG)
     logger.setLevel(logging.INFO)
 
-    dqn_cart_pole()
+    # dqn_cart_pole()
     # async_cart_pole()
     # a3c_cart_pole()
+    a2c_cart_pole()
     # a3c_continuous()
     # p3o_continuous()
     # d3pg_continuous()

network/base_network.py

@@ -18,8 +18,10 @@ class BasicNet:
         if self.gpu:
             self.cuda()
             self.FloatTensor = torch.cuda.FloatTensor
+            self.LongTensor = torch.cuda.LongTensor
         else:
             self.FloatTensor = torch.FloatTensor
+            self.LongTensor = torch.LongTensor
 
     def to_torch_variable(self, x, dtype='float32'):
         if isinstance(x, Variable):

utils/config.py

@@ -51,3 +51,4 @@ class Config:
         self.max_steps = 0
         self.success_threshold = float('inf')
         self.render_episode_freq = 0
+        self.rollout_length = None

utils/misc.py

@@ -60,6 +60,7 @@ def run_episodes(agent):
             if avg_reward > config.success_threshold:
                 break
 
+    agent.close()
     return steps, rewards, avg_test_rewards
 
 def sync_grad(target_network, src_network):