CARTPOLE - create vanilla a2c and synchronous a2c

a2c_cart.py

@@ -0,0 +1,151 @@
+"""
+Reinforcement Learning (A3C) using Pytroch + multiprocessing.
+The most simple implementation for continuous action.
+"""
+
+import torch
+import torch.nn as nn
+from utils import v_wrap, set_init, push_and_pull, record
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+from shared_adam import SharedAdam
+import gym
+import os
+
+import argparse
+import time
+from datetime import datetime
+timestr = time.strftime("%d.%m.%Y - %H:%M:%S")
+
+os.environ["OMP_NUM_THREADS"] = "1"
+
+UPDATE_GLOBAL_ITER = 500
+GAMMA = 0.9
+MAX_EP = 3000
+
+env = gym.make('CartPole-v0')
+N_S = env.observation_space.shape[0]
+N_A = env.action_space.n
+
+
+def handleArguments():
+    """Handles CLI arguments and saves them globally"""
+    parser = argparse.ArgumentParser(
+        description="Switch between modes in A2C or loading models from previous games")
+    parser.add_argument("--demo_mode", "-d", help="Renders the gym environment", action="store_true")
+    parser.add_argument("--load_model", "-l", help="Loads the model of previously gained training data", action="store_true")
+    global args
+    args = parser.parse_args()
+
+
+class Net(nn.Module):
+    def __init__(self, s_dim, a_dim):
+        super(Net, self).__init__()
+        self.s_dim = s_dim
+        self.a_dim = a_dim
+        self.pi1 = nn.Linear(s_dim, 128)
+        self.pi2 = nn.Linear(128, a_dim)
+        self.v1 = nn.Linear(s_dim, 128)
+        self.v2 = nn.Linear(128, 1)
+        set_init([self.pi1, self.pi2, self.v1, self.v2])
+        self.distribution = torch.distributions.Categorical
+
+    def forward(self, x):
+        pi1 = torch.tanh(self.pi1(x))
+        logits = self.pi2(pi1)
+        v1 = torch.tanh(self.v1(x))
+        values = self.v2(v1)
+        return logits, values
+
+    def set_init(layers):
+        for layer in layers:
+            nn.init.normal_(layer.weight, mean=0., std=0.1)
+            nn.init.constant_(layer.bias, 0.)
+
+    def choose_action(self, s):
+        self.eval()
+        logits, _ = self.forward(s)
+        prob = F.softmax(logits, dim=1).data
+        m = self.distribution(prob)
+        return m.sample().numpy()[0]
+
+    def loss_func(self, s, a, v_t):
+        self.train()
+        logits, values = self.forward(s)
+        td = v_t - values
+        c_loss = td.pow(2)
+
+        probs = F.softmax(logits, dim=1)
+        m = self.distribution(probs)
+        exp_v = m.log_prob(a) * td.detach().squeeze()
+        a_loss = -exp_v
+        total_loss = (c_loss + a_loss).mean()
+        return total_loss
+
+
+class Worker(mp.Process):
+    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name):
+        super(Worker, self).__init__()
+        self.name = 'w%02i' % name
+        self.g_ep, self.g_ep_r, self.res_queue = global_ep, global_ep_r, res_queue
+        self.gnet, self.opt = gnet, opt
+        self.lnet = Net(N_S, N_A)  # local network
+        self.env = gym.make('CartPole-v0').unwrapped
+
+    def run(self):
+        global episode
+        total_step = 1
+        while self.g_ep.value < MAX_EP:
+            s = self.env.reset()
+            buffer_s, buffer_a, buffer_r = [], [], []
+            ep_r = 0.
+            while True:
+                if self.name == 'w00':
+                    self.env.render()
+                a = self.lnet.choose_action(v_wrap(s[None, :]))
+                s_, r, done, _ = self.env.step(a)
+                if done: r = -1
+                ep_r += r
+                buffer_a.append(a)
+                buffer_s.append(s)
+                buffer_r.append(r)
+
+                if self.g_ep.value % UPDATE_GLOBAL_ITER == 0 or done:  # update global and assign to local net
+                    # sync
+                    push_and_pull(self.opt, self.lnet, self.gnet, done, s_, buffer_s, buffer_a, buffer_r, GAMMA)
+                    buffer_s, buffer_a, buffer_r = [], [], []
+
+                    if done:  # done and print information
+                        record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name)
+                        break
+                s = s_
+                total_step += 1
+        self.res_queue.put(None)
+
+
+if __name__ == "__main__":
+
+    gnet = Net(N_S, N_A)  # global network
+    gnet.share_memory()  # share the global parameters in multiprocessing
+    opt = SharedAdam(gnet.parameters(), lr=1e-4, betas=(0.92, 0.999))  # global optimizer
+    global_ep, global_ep_r, res_queue = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()
+
+    # parallel training
+    worker = Worker(gnet, opt, global_ep, global_ep_r, res_queue, 0)
+    worker.start()
+    res = []  # record episode reward to plot
+    while True:
+        r = res_queue.get()
+        if r is not None:
+            res.append(r)
+        else:
+            break
+
+
+    import matplotlib.pyplot as plt
+
+    plt.plot(res)
+    plt.ylabel('Average Reward')
+    plt.xlabel('Episode')
+    plt.show()
+

sa2c_cart.py

@@ -0,0 +1,162 @@
+"""
+Reinforcement Learning (A3C) using Pytroch + multiprocessing.
+The most simple implementation for continuous action.
+View more on my Chinese tutorial page [莫烦Python](https://morvanzhou.github.io/).
+"""
+
+import torch
+import torch.nn as nn
+from utils import v_wrap, set_init, push_and_pull, record
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+from shared_adam import SharedAdam
+import gym
+import os
+
+import argparse
+import time
+from datetime import datetime
+timestr = time.strftime("%d.%m.%Y - %H:%M:%S")
+
+os.environ["OMP_NUM_THREADS"] = "1"
+
+UPDATE_GLOBAL_ITER = 500
+GAMMA = 0.9
+MAX_EP = 3000
+episode = 0
+
+env = gym.make('CartPole-v0')
+N_S = env.observation_space.shape[0]
+N_A = env.action_space.n
+
+
+def handleArguments():
+    """Handles CLI arguments and saves them globally"""
+    parser = argparse.ArgumentParser(
+        description="Switch between modes in A2C or loading models from previous games")
+    parser.add_argument("--demo_mode", "-d", help="Renders the gym environment", action="store_true")
+    parser.add_argument("--load_model", "-l", help="Loads the model of previously gained training data", action="store_true")
+    global args
+    args = parser.parse_args()
+
+
+class Net(nn.Module):
+    def __init__(self, s_dim, a_dim):
+        super(Net, self).__init__()
+        self.s_dim = s_dim
+        self.a_dim = a_dim
+        self.pi1 = nn.Linear(s_dim, 128)
+        self.pi2 = nn.Linear(128, a_dim)
+        self.v1 = nn.Linear(s_dim, 128)
+        self.v2 = nn.Linear(128, 1)
+        set_init([self.pi1, self.pi2, self.v1, self.v2])
+        self.distribution = torch.distributions.Categorical
+
+    def forward(self, x):
+        pi1 = torch.tanh(self.pi1(x))
+        logits = self.pi2(pi1)
+        v1 = torch.tanh(self.v1(x))
+        values = self.v2(v1)
+        return logits, values
+
+    def set_init(layers):
+        for layer in layers:
+            nn.init.normal_(layer.weight, mean=0., std=0.1)
+            nn.init.constant_(layer.bias, 0.)
+
+    def choose_action(self, s):
+        self.eval()
+        logits, _ = self.forward(s)
+        prob = F.softmax(logits, dim=1).data
+        m = self.distribution(prob)
+        return m.sample().numpy()[0]
+
+    def loss_func(self, s, a, v_t):
+        self.train()
+        logits, values = self.forward(s)
+        td = v_t - values
+        c_loss = td.pow(2)
+
+        probs = F.softmax(logits, dim=1)
+        m = self.distribution(probs)
+        exp_v = m.log_prob(a) * td.detach().squeeze()
+        a_loss = -exp_v
+        total_loss = (c_loss + a_loss).mean()
+        return total_loss
+
+
+class Worker(mp.Process):
+    def __init__(self, gnet, opt, global_ep, global_ep_r, res_queue, name):
+        super(Worker, self).__init__()
+        self.name = 'w%02i' % name
+        self.g_ep, self.g_ep_r, self.res_queue = global_ep, global_ep_r, res_queue
+        self.gnet, self.opt = gnet, opt
+        self.lnet = Net(N_S, N_A)  # local network
+        self.env = gym.make('CartPole-v0').unwrapped
+
+    def run(self):
+        global episode
+        total_step = 1
+        while self.g_ep.value < MAX_EP:
+            s = self.env.reset()
+            buffer_s, buffer_a, buffer_r = [], [], []
+            ep_r = 0.
+            while True:
+                if self.name == 'w00':
+                    self.env.render()
+                a = self.lnet.choose_action(v_wrap(s[None, :]))
+                s_, r, done, _ = self.env.step(a)
+                if done: r = -1
+                ep_r += r
+                buffer_a.append(a)
+                buffer_s.append(s)
+                buffer_r.append(r)
+
+                if self.g_ep.value % UPDATE_GLOBAL_ITER == 0 or done:  # update global and assign to local net
+                    # sync
+                    if self.g_ep.value % UPDATE_GLOBAL_ITER == 0 and self.g_ep.value != 0:
+                        print (self.g_ep.value)
+                        print("sleep...")
+                        self.res_queue.put(time.sleep(10))
+                        push_and_pull(self.opt, self.lnet, self.gnet, done, s_, buffer_s, buffer_a, buffer_r, GAMMA)
+                        buffer_s, buffer_a, buffer_r = [], [], []
+
+                    if done:  # done and print information
+                        record(self.g_ep, self.g_ep_r, ep_r, self.res_queue, self.name)
+                        break
+                s = s_
+                total_step += 1
+        self.res_queue.put(None)
+
+
+if __name__ == "__main__":
+
+    gnet = Net(N_S, N_A)  # global network
+    gnet.share_memory()  # share the global parameters in multiprocessing
+    opt = SharedAdam(gnet.parameters(), lr=1e-4, betas=(0.92, 0.999))  # global optimizer
+    global_ep, global_ep_r, res_queue = mp.Value('i', 0), mp.Value('d', 0.), mp.Queue()
+
+    # parallel training
+    workers = [Worker(gnet, opt, global_ep, global_ep_r, res_queue, i) for i in range(mp.cpu_count())]
+    [w.start() for w in workers]
+    res = []  # record episode reward to plot
+    while True:
+        #print ("Episode: ", global_ep.value)
+        r = res_queue.get()
+        if r is not None:
+            res.append(r)
+            #if global_ep.value % 500 == 0:
+                #res.put(time.sleep(5))
+                #[w.join() for w in workers]
+        elif global_ep.value == MAX_EP:
+            break
+        else:
+            True
+
+    import matplotlib.pyplot as plt
+
+    plt.plot(res)
+    plt.ylabel('Average Reward')
+    plt.xlabel('Episode')
+    plt.show()
+