[WIP]: ACER

choru-k · choru-k · commit 6fea9ea51082 · 2018-12-13T20:52:19.000+09:00
diff --git a/README.md b/README.md
@@ -27,10 +27,10 @@ So you can run this example in your computer(maybe it take just only 1~2 minitue
 ## Parallel
 - [x] Asynchronous Q-learning [[11]](#reference)
 - [x] A3C (Asynchronous Advange Actor Critice) [[11]](#reference)
-
-## Will
 - [ ] ACER [[21]](#reference)
 - [ ] APE-X [[15]](#reference)
+
+## Will
 - [ ] R2D2 [[16]](#reference)
 - [ ] RND [[17]](#reference)
 - [ ] QRDQN [[18]](#reference)
@@ -67,3 +67,5 @@ So you can run this example in your computer(maybe it take just only 1~2 minitue
 - https://github.com/reinforcement-learning-kr/pg_travel
 - https://github.com/reinforcement-learning-kr/distributional_rl
 - https://github.com/Kaixhin/Rainbow
+- https://github.com/Kaixhin/ACER
+- https://github.com/higgsfield/RL-Adventure-2
diff --git a/parallel/2-A3C/model.py b/parallel/2-A3C/model.py
@@ -4,11 +4,6 @@
 
 from config import gamma
 
-def set_init(layers):
-    for layer in layers:
-        nn.init.normal_(layer.weight, mean=0., std=0.1)
-        nn.init.constant_(layer.bias, 0.1)
-
 class Model(nn.Module):
     def __init__(self, num_inputs, num_outputs):
         super(Model, self).__init__()
diff --git a/parallel/3-ACER/config.py b/parallel/3-ACER/config.py
@@ -0,0 +1,17 @@
+import torch
+
+env_name = 'CartPole-v1'
+gamma = 0.99
+lr = 0.001
+goal_score = 200
+log_interval = 10
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+max_episode = 30000
+
+
+replay_memory_capacity = 1000
+truncation_clip = 10
+delta = 1
+trust_region_decay = 0.99
+replay_ratio = 4
+max_gradient_norm = 40
diff --git a/parallel/3-ACER/memory.py b/parallel/3-ACER/memory.py
@@ -0,0 +1,34 @@
+import random
+from collections import namedtuple, deque
+
+Transition = namedtuple('Transition', ('state', 'next_state', 'action', 'reward', 'mask', 'policy'))
+
+
+class Memory(object):
+    def __init__(self, capacity):
+        self.memory = deque(maxlen=capacity)
+        self.capacity = capacity
+
+    def push(self, trajectory):
+        self.memory.append(trajectory.trajectory)
+
+    def sample(self):
+        trajectory = self.memory[random.randrange(len(self.memory))]
+        return Transition(*zip(*trajectory))
+
+    def __len__(self):
+        return len(self.memory)
+
+class Trajectory(object):
+    def __init__(self):
+        self.trajectory = []
+
+    def push(self, state, next_state, action, reward, mask, policy):
+        self.trajectory.append(Transition(state, next_state, action, reward, mask, policy))
+
+    def sample(self):
+        trajectory = self.trajectory
+        return Transition(*zip(*trajectory))
+
+    def __len__(self):
+        return len(self.trajectory)
diff --git a/parallel/3-ACER/model.py b/parallel/3-ACER/model.py
@@ -0,0 +1,134 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from config import gamma, truncation_clip, delta, max_gradient_norm, trust_region_decay
+
+class Model(nn.Module):
+    def __init__(self, num_inputs, num_outputs):
+        super(Model, self).__init__()
+        self.num_inputs = num_inputs
+        self.num_outputs = num_outputs
+
+        self.fc = nn.Linear(num_inputs, 128)
+        self.fc_actor = nn.Linear(128, num_outputs)
+        self.fc_critic = nn.Linear(128, num_outputs)
+
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform(m.weight)
+
+    def forward(self, input):
+        x = F.relu(self.fc(input))
+        policy = F.softmax(self.fc_actor(x), dim=1)
+        q_value = self.fc_critic(x)
+        value = (policy * q_value).sum(-1, keepdim=True).view(-1)
+        return policy, q_value, value
+
+class LocalModel(Model):
+    def __init__(self, num_inputs, num_outputs):
+        super(LocalModel, self).__init__(num_inputs, num_outputs)
+
+    def pull_from_global_model(self, global_model):
+        self.load_state_dict(global_model.state_dict())
+
+    def update_model(self, loss, global_optimizer, global_model, global_average_model):
+        global_optimizer.zero_grad()
+        loss.backward()
+        # nn.utils.clip_grad_norm_(self.parameters(), max_gradient_norm)
+
+        for lp, gp in zip(self.parameters(), global_model.parameters()):
+            if gp.grad is not None:
+                return
+            gp.grad = lp.grad
+
+        global_optimizer.step()
+
+        for gp, gap in zip(global_model.parameters(), global_average_model.parameters()):
+            gap = trust_region_decay * gap + (1 - trust_region_decay) * gp
+
+
+    def compute_q_retraces(self, rewards, masks, values, q_actions, rho_actions, next_value):
+        q_retraces = torch.zeros(rewards.size())
+        q_retraces[-1] = next_value
+
+        q_ret = q_retraces[-1]
+        for step in reversed(range(len(rewards) - 1)):
+            q_ret = rewards[step] + gamma * q_ret
+            q_retraces[step] = q_ret
+            q_ret = rho_actions[step] * (q_ret - q_actions[step]) + values[step]
+
+        return q_retraces
+
+
+    def get_loss(self, on_policy, trajectory, average_model):
+        states, next_states, actions, rewards, masks, old_policies = trajectory
+        states = torch.stack(states)
+        next_states = torch.stack(next_states)
+        actions = torch.Tensor(actions).long().view(-1,1)
+        rewards = torch.Tensor(rewards)
+        masks = torch.Tensor(masks)
+        old_policies = torch.stack(old_policies)
+
+        states = states.view(-1, self.num_inputs)
+        next_states = next_states.view(-1, self.num_inputs)
+        policies, Qs, Vs = self.forward(states)
+
+        Q_actions = Qs.gather(1, actions).view(-1)
+
+        if not on_policy:
+            rhos = policies / old_policies
+        else:
+            rhos = torch.zeros(policies.size()).fill_(1)
+
+        rho_actions = rhos.gather(1, actions).view(-1)
+
+        if masks[-1] == 0:
+            Qret = 0
+        else:
+            Qret = Vs[-1]
+        Qrets = self.compute_q_retraces(rewards, masks, Vs, Q_actions, rho_actions, Qret)
+        log_policy = torch.log(policies)
+        log_policy_action = log_policy.gather(1, actions).view(-1)
+
+        actor_loss_1 = - (log_policy_action * (
+            rho_actions.clamp(max=truncation_clip) * (Qrets - Vs)
+        ).detach()).mean()
+        actor_loss_2 = - (log_policy * (
+            (1 - truncation_clip / rhos).clamp(min=0) * policies * (Qs - Vs.view(-1,1).expand_as(Qs))
+        ).detach()).sum(1).mean()
+        actor_loss = actor_loss_1 + actor_loss_2
+
+        value_loss = ((Qret - Q_actions) ** 2).mean()
+
+
+        g_1 = ((1 / log_policy_action) * (
+            rho_actions.clamp(max=truncation_clip) * (Qrets - Vs)
+        ))
+        g_2 = ((1 / log_policy) * (
+            (1 - truncation_clip / rhos).clamp(min=0) * policies * (Qs - Vs.view(-1,1).expand_as(Qs))
+        )).sum(1)
+        g = (g_1 + g_2).detach()
+        average_policies, _, _ = average_model(states)
+        k = (average_policies / policies).gather(1, actions).view(-1)
+
+        kl = (average_policies * torch.log(average_policies / policies)).sum(1).mean(0)
+
+
+        k_dot_g = (k * g).sum()
+        k_dot_k = (k * k).sum()
+
+        adj = ((k_dot_g - delta) / k_dot_k).clamp(min=0).detach()
+        trust_region_actor_loss = actor_loss + adj * kl
+
+        loss = trust_region_actor_loss + value_loss
+
+        return loss
+
+    def get_action(self, input):
+        policy, _, _ = self.forward(input)
+        policy = policy[0].data.numpy()
+
+        action = np.random.choice(self.num_outputs, 1, p=policy)[0]
+        return action, policy
diff --git a/parallel/3-ACER/shared_adam.py b/parallel/3-ACER/shared_adam.py
@@ -0,0 +1,34 @@
+import torch
+class SharedAdam(torch.optim.Adam): # extend a pytorch optimizer so it shares grads across processes
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        super(SharedAdam, self).__init__(params, lr, betas, eps, weight_decay)
+        for group in self.param_groups:
+            for p in group['params']:
+                state = self.state[p]
+                state['shared_steps'], state['step'] = torch.zeros(1).share_memory_(), 0
+                state['exp_avg'] = p.data.new().resize_as_(p.data).zero_().share_memory_()
+                state['exp_avg_sq'] = p.data.new().resize_as_(p.data).zero_().share_memory_()
+
+        def step(self, closure=None):
+            for group in self.param_groups:
+                for p in group['params']:
+                    if p.grad is None: continue
+                    self.state[p]['shared_steps'] += 1
+                    self.state[p]['step'] = self.state[p]['shared_steps'][0] - 1 # a "step += 1"  comes later
+            super.step(closure)
+
+# class SharedAdam(torch.optim.Adam):
+#     def __init__(self, params, lr=1e-3, betas=(0.9, 0.9), eps=1e-8,
+#                  weight_decay=0):
+#         super(SharedAdam, self).__init__(params, lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+#         # State initialization
+#         for group in self.param_groups:
+#             for p in group['params']:
+#                 state = self.state[p]
+#                 state['step'] = 0
+#                 state['exp_avg'] = torch.zeros_like(p.data)
+#                 state['exp_avg_sq'] = torch.zeros_like(p.data)
+#
+#                 # share in memory
+#                 state['exp_avg'].share_memory_()
+#                 state['exp_avg_sq'].share_memory_()
diff --git a/parallel/3-ACER/train.py b/parallel/3-ACER/train.py
@@ -0,0 +1,47 @@
+import gym
+import torch
+
+from model import Model
+from worker import Worker
+from shared_adam import SharedAdam
+from tensorboardX import SummaryWriter
+import torch.multiprocessing as mp
+
+from config import env_name, lr
+
+def main():
+    env = gym.make(env_name)
+    env.seed(500)
+    torch.manual_seed(500)
+
+    num_inputs = env.observation_space.shape[0]
+    num_actions = env.action_space.n
+    env.close()
+
+    global_model = Model(num_inputs, num_actions)
+    global_average_model = Model(num_inputs, num_actions)
+    global_model.share_memory()
+    global_average_model.share_memory()
+    global_optimizer = SharedAdam(global_model.parameters(), lr=lr)
+    global_ep, global_ep_r, res_queue = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()
+
+    writer = SummaryWriter('logs')
+
+    # n = mp.cpu_count()
+    n = 1
+    workers = [Worker(global_model, global_average_model, global_optimizer, global_ep, global_ep_r, res_queue, i) for i in range(n)]
+    [w.start() for w in workers]
+    res = []
+    while True:
+        r = res_queue.get()
+        if r is not None:
+            res.append(r)
+            [ep, ep_r, loss] = r
+            writer.add_scalar('log/score', float(ep_r), ep)
+            writer.add_scalar('log/loss', float(loss), ep)
+        else:
+            break
+    [w.join() for w in workers]
+
+if __name__=="__main__":
+    main()
diff --git a/parallel/3-ACER/worker.py b/parallel/3-ACER/worker.py
@@ -0,0 +1,87 @@
+import gym
+import torch
+import torch.multiprocessing as mp
+import numpy as np
+from model import LocalModel
+from memory import Memory, Trajectory
+from config import env_name, max_episode, log_interval, replay_memory_capacity, replay_ratio
+
+class Worker(mp.Process):
+    def __init__(self, global_model, global_average_model, global_optimizer, global_ep, global_ep_r, res_queue, name):
+        super(Worker, self).__init__()
+
+        self.env = gym.make(env_name)
+        self.env.seed(500)
+
+        self.name = 'w%i' % name
+        self.global_ep, self.global_ep_r, self.res_queue = global_ep, global_ep_r, res_queue
+        self.global_model, self.global_average_model, self.global_optimizer = global_model, global_average_model, global_optimizer
+        self.local_model = LocalModel(self.env.observation_space.shape[0], self.env.action_space.n)
+        self.num_actions = self.env.action_space.n
+
+        self.memory = Memory(replay_memory_capacity)
+
+    def record(self, score, loss):
+        with self.global_ep.get_lock():
+            self.global_ep.value += 1
+        with self.global_ep_r.get_lock():
+            if self.global_ep_r.value == 0.:
+                self.global_ep_r.value = score
+            else:
+                self.global_ep_r.value = 0.99 * self.global_ep_r.value + 0.01 * score
+        if self.global_ep.value % log_interval == 0:
+            print('{} , {} episode | score: {:.2f}'.format(
+                self.name, self.global_ep.value, self.global_ep_r.value))
+
+        self.res_queue.put([self.global_ep.value, self.global_ep_r.value, loss])
+
+    def run(self):
+        while self.global_ep.value < max_episode:
+            self.algorithm(True)
+            n = np.random.poisson(replay_ratio)
+            for _ in range(n):
+                self.algorithm(False)
+
+    def algorithm(self, on_policy):
+        self.local_model.pull_from_global_model(self.global_model)
+        if not on_policy and len(self.memory) > 100:
+            trajectory = self.memory.sample()
+        else:
+            trajectory, score = self.run_env()
+        loss = self.local_model.get_loss(on_policy, trajectory, self.global_average_model)
+        self.local_model.update_model(loss, self.global_optimizer, self.global_model, self.global_average_model)
+        if on_policy:
+            self.record(score, loss)
+
+
+    def run_env(self):
+        done = False
+        score = 0
+        steps = 0
+
+        state = self.env.reset()
+        state = torch.Tensor(state)
+        state = state.unsqueeze(0)
+        trajectory = Trajectory()
+
+        while True:
+            action, policy = self.local_model.get_action(state)
+            policy = torch.Tensor(policy)
+
+            next_state, reward, done, _ = self.env.step(action)
+            next_state = torch.Tensor(next_state)
+            next_state = next_state.unsqueeze(0)
+
+            mask = 0 if done else 1
+            reward = reward if not done or score == 499 else -1
+            trajectory.push(state, next_state, action, reward, mask, policy)
+
+            score += reward
+            state = next_state
+
+            if done:
+                break
+
+        self.memory.push(trajectory)
+        trajectory = trajectory.sample()
+        return trajectory, score
