Major refactor

agent/DDPG_agent.py

@@ -5,7 +5,8 @@
 #######################################################################
 
 from network import *
-from replay import *
+from component import *
+from utils import *
 import pickle
 
 class DDPGAgent:

agent/DQN_agent.py

@@ -5,8 +5,8 @@
 #######################################################################
 
 from network import *
-from replay import *
+from component import *
-from policy import *
+from utils import *
 import numpy as np
 import time
 import os

agent/__init__.py

@@ -0,0 +1,3 @@
+from async_agent import *
+from DDPG_agent import *
+from DQN_agent import *

agent/async_agent.py

@@ -4,16 +4,12 @@
 # declaration at the top                                              #
 #######################################################################
 
-from network import *
-from policy import *
 import numpy as np
 import torch.multiprocessing as mp
-from task import *
 from network import *
-from async_workers.one_step_sarsa import *
+from utils import *
-from async_workers.n_step_q import *
+from component import *
-from async_workers.actor_critic import *
+from async_worker import *
-from async_workers.one_step_sarsa import *
 import pickle
 import os
 import time

async_worker/__init__.py

@@ -0,0 +1,5 @@
+from actor_critic import *
+from continuous_actor_critic import *
+from n_step_q import *
+from one_step_sarsa import *
+from one_step_q import *

async_worker/actor_critic.py

@@ -60,12 +60,14 @@ class AdvantageActorCritic:
 
                 pending = []
                 self.worker_network.zero_grad()
+                self.optimizer.zero_grad()
                 loss.backward()
                 nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
                 for param, worker_param in zip(
                         config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
+                    if param.grad is not None:
+                        break
+                    param._grad = worker_param.grad
                 self.optimizer.step()
                 self.worker_network.load_state_dict(config.learning_network.state_dict())
                 self.worker_network.reset(terminal)

async_worker/continuous_actor_critic.py

@@ -0,0 +1,86 @@
+#######################################################################
+# 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
+from torch.autograd import Variable
+import torch.nn as nn
+
+class ContinuousAdvantageActorCritic:
+    def __init__(self, config):
+        self.config = config
+        self.optimizer = config.optimizer_fn(config.learning_network.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(config.learning_network.state_dict())
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+    def episode(self, deterministic=False):
+        config = self.config
+        state = self.task.reset()
+        steps = 0
+        total_reward = 0
+        pending = []
+        pi = Variable(torch.FloatTensor([np.pi]))
+        while not config.stop_signal.value and \
+                (not config.max_episode_length or steps < config.max_episode_length):
+            mean, var, value = self.worker_network.predict(np.stack([state]))
+            action = self.policy.sample(mean.data.numpy().flatten(),
+                                        var.data.numpy().flatten(),
+                                        deterministic)
+            next_state, reward, terminal, _ = self.task.step(action)
+
+            steps += 1
+            total_reward += reward
+
+            if deterministic:
+                if terminal:
+                    break
+                state = next_state
+                continue
+
+            pending.append([mean, var, value, action, reward])
+            with config.steps_lock:
+                config.total_steps.value += 1
+
+            if terminal or len(pending) >= config.update_interval:
+                loss = 0
+                if terminal:
+                    R = torch.FloatTensor([[0]])
+                else:
+                    R = self.worker_network.critic(np.stack([next_state])).data
+                GAE = torch.FloatTensor([[0]])
+                for i in reversed(range(len(pending))):
+                    mean, var, value, action, reward = pending[i]
+                    R = reward + config.discount * R
+                    advantage = Variable(R) - value
+                    GAE = config.discount * config.gae_tau * GAE + advantage.data
+                    loss += 0.5 * advantage.pow(2)
+                    action = Variable(torch.FloatTensor([action]))
+                    prob_part1 = (-(action - mean).pow(2) / (2 * var)).exp()
+                    prob_part2 = 1 / (2 * var * pi.expand_as(var)).sqrt()
+                    prob = prob_part1 * prob_part2
+                    log_prob = prob.log()
+                    loss += -torch.sum(log_prob) * Variable(GAE)
+                    entropy = 0.5 * (1.0 + (var + 2 * pi.expand_as(var)).log()).sum()
+                    loss += config.entropy_weight * entropy
+
+                pending = []
+                self.worker_network.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+                self.optimizer.zero_grad()
+                for param, worker_param in zip(
+                        config.learning_network.parameters(), self.worker_network.parameters()):
+                    param._grad = worker_param.grad.clone()
+                self.optimizer.step()
+                self.worker_network.load_state_dict(config.learning_network.state_dict())
+                self.worker_network.reset(terminal)
+
+            if terminal:
+                break
+            state = next_state
+
+        return steps, total_reward

async_worker/n_step_q.py

@@ -57,12 +57,14 @@ class NStepQLearning:
 
                 pending = []
                 self.worker_network.zero_grad()
+                self.optimizer.zero_grad()
                 loss.backward()
                 nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
                 for param, worker_param in zip(
                         config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
+                    if param.grad is not None:
+                        break
+                    param._grad = worker_param.grad
                 self.optimizer.step()
                 self.worker_network.load_state_dict(config.learning_network.state_dict())
                 self.worker_network.reset(terminal)

async_worker/one_step_q.py

@@ -55,12 +55,14 @@ class OneStepQLearning:
 
                 pending = []
                 self.worker_network.zero_grad()
+                self.optimizer.zero_grad()
                 loss.backward()
                 nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
                 for param, worker_param in zip(
                         config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
+                    if param.grad is not None:
+                        break
+                    param._grad = worker_param.grad
                 self.optimizer.step()
                 self.worker_network.load_state_dict(config.learning_network.state_dict())
                 self.worker_network.reset(terminal)

async_worker/one_step_sarsa.py

@@ -60,12 +60,14 @@ class OneStepSarsa:
 
                 pending = []
                 self.worker_network.zero_grad()
+                self.optimizer.zero_grad()
                 loss.backward()
                 nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
-                self.optimizer.zero_grad()
                 for param, worker_param in zip(
                         config.learning_network.parameters(), self.worker_network.parameters()):
-                    param._grad = worker_param.grad.clone()
+                    if param.grad is not None:
+                        break
+                    param._grad = worker_param.grad
                 self.optimizer.step()
                 self.worker_network.load_state_dict(config.learning_network.state_dict())
                 self.worker_network.reset(terminal)

component/__init__.py

@@ -0,0 +1,5 @@
+from atari_wrapper import *
+from policy import *
+from replay import *
+from task import *
+from random_process import *

component/policy.py

@@ -45,4 +45,13 @@ class SamplePolicy:
             return np.argmax(action_value)
         return np.random.choice(np.arange(len(action_value)), p=action_value)
     def update_epsilon(self):
-        pass
+        pass
+
+class GaussianPolicy:
+    def sample(self, mean, var, deterministic=False):
+        if deterministic:
+            return mean
+        return mean + np.sqrt(var) * np.random.randn(*mean.shape)
+
+    def update_epsilon(self):
+        pass

component/task.py

@@ -88,7 +88,8 @@ class Pendulum(BasicTask):
         self.state_dim = self.env.observation_space.shape[0]
 
     def step(self, action):
-        action = 2 * np.clip(action, -1, 1)
+        # action = 2 * np.clip(action, -1, 1)
+        action = np.clip(action, -2, 2)
         next_state, reward, done, info = self.env.step(action)
         return next_state, reward, done, info
 

main.py

@@ -1,10 +1,7 @@
-from async_agent import *
-from DQN_agent import *
-from DDPG_agent import *
-from logger import *
 import logging
-from random_process import *
+from agent import *
-from config import Config
+from component import *
+from utils import *
 
 def dqn_cart_pole():
     config = dict()
@@ -35,8 +32,8 @@ def async_cart_pole():
     config.network_fn = lambda: FCNet([4, 50, 200, 2])
     config.policy_fn = lambda: GreedyPolicy(epsilon=0.5, final_step=5000, min_epsilon=0.1)
     # config.worker = OneStepQLearning
-    # config.worker = NStepQLearning
+    config.worker = NStepQLearning
-    config.worker = OneStepSarsa
+    # config.worker = OneStepSarsa
     config.discount = 0.99
     config.target_network_update_freq = 200
     config.max_episode_length = 200
@@ -52,15 +49,37 @@ def a3c_cart_pole():
     config = Config()
     config.task_fn = lambda: CartPole()
     config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
-    config.network_fn = lambda: ActorCriticFCNet([4, 200, 2])
+    config.network_fn = lambda: ActorCriticFCNet(4, 2)
     config.policy_fn = SamplePolicy
     config.worker = AdvantageActorCritic
     config.discount = 0.99
     config.max_episode_length = 200
     config.num_workers = 16
     config.update_interval = 6
+    config.test_interval = 100
+    config.test_repetitions = 30
+    config.logger = Logger('./log', gym.logger)
+    config.gae_tau = 1.0
+    config.entropy_weight = 0.01
+    agent = AsyncAgent(config)
+    agent.run()
+
+def a3c_pendulum():
+    config = Config()
+    config.task_fn = lambda: Pendulum()
+    task = config.task_fn()
+    config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
+    config.network_fn = lambda: ContinuousActorCriticNet(
+        task.env.observation_space.shape[0], 64, task.env.action_space.shape[0])
+    config.policy_fn = lambda: GaussianPolicy()
+    config.worker = ContinuousAdvantageActorCritic
+    config.discount = 0.99
+    config.max_episode_length = 200
+    config.num_workers = 16
+    config.update_interval = 20
     config.test_interval = 1
     config.test_repetitions = 50
+    config.entropy_weight = 0.0001
     config.logger = Logger('./log', gym.logger)
     agent = AsyncAgent(config)
     agent.run()
@@ -190,6 +209,7 @@ if __name__ == '__main__':
     # dqn_cart_pole()
     # async_cart_pole()
     a3c_cart_pole()
+    # a3c_pendulum()
 
     # dqn_pixel_atari('PongNoFrameskip-v3')
     # async_pixel_atari('PongNoFrameskip-v3')

network.py

@@ -1,346 +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                                              #
-#######################################################################
-
-import torch
-from torch.autograd import Variable
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-
-# Base class for all kinds of network
-class BasicNet:
-    def __init__(self, optimizer_fn, gpu, LSTM=False):
-        if optimizer_fn is not None:
-            self.optimizer = optimizer_fn(self.parameters())
-        self.gpu = gpu and torch.cuda.is_available()
-        self.LSTM = LSTM
-        if self.gpu:
-            self.cuda()
-
-    def to_torch_variable(self, x, dtype='float32'):
-        if isinstance(x, Variable):
-            return x
-        if not isinstance(x, torch.FloatTensor):
-            x = torch.from_numpy(np.asarray(x, dtype=dtype))
-        if self.gpu:
-            x = x.cuda()
-        return Variable(x)
-
-    def reset(self, terminal):
-        if not self.LSTM:
-            return
-        if terminal:
-            self.h.data.zero_()
-            self.c.data.zero_()
-        self.h = Variable(self.h.data)
-        self.c = Variable(self.c.data)
-
-# Base class for value based methods
-class VanillaNet(BasicNet):
-    def predict(self, x, to_numpy=False):
-        y = self.forward(x)
-        if to_numpy:
-            y = y.cpu().data.numpy()
-        return y
-
-# Base class for actor critic method
-class ActorCriticNet(BasicNet):
-    def predict(self, x):
-        phi = self.forward(x, True)
-        pre_prob = self.fc_actor(phi)
-        prob = F.softmax(pre_prob)
-        log_prob = F.log_softmax(pre_prob)
-        value = self.fc_critic(phi)
-        return prob, log_prob, value
-
-    def critic(self, x):
-        phi = self.forward(x, False)
-        return self.fc_critic(phi)
-
-# Base class for dueling architecture
-class DuelingNet(BasicNet):
-    def predict(self, x, to_numpy=False):
-        phi = self.forward(x)
-        value = self.fc_value(phi)
-        advantange = self.fc_advantage(phi)
-        q = value.expand_as(advantange) + (advantange - advantange.mean(1).expand_as(advantange))
-        if to_numpy:
-            return q.cpu().data.numpy()
-        return q
-
-# Starting of several network instances
-
-# Network for CartPole with value based methods
-class FCNet(nn.Module, VanillaNet):
-    def __init__(self, dims, optimizer_fn=None, gpu=True):
-        super(FCNet, self).__init__()
-        self.fc1 = nn.Linear(dims[0], dims[1])
-        self.fc2 = nn.Linear(dims[1], dims[2])
-        self.fc3 = nn.Linear(dims[2], dims[3])
-        self.criterion = nn.MSELoss()
-        BasicNet.__init__(self, optimizer_fn, gpu)
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        x = x.view(x.size(0), -1)
-        y = F.relu(self.fc1(x))
-        y = F.relu(self.fc2(y))
-        y = self.fc3(y)
-        return y
-
-# Network for CartPole with dueling architecture
-class DuelingFCNet(nn.Module, DuelingNet):
-    def __init__(self, dims, optimizer_fn=None, gpu=True):
-        super(DuelingFCNet, self).__init__()
-        self.fc1 = nn.Linear(dims[0], dims[1])
-        self.fc2 = nn.Linear(dims[1], dims[2])
-        self.fc_value = nn.Linear(dims[2], 1)
-        self.fc_advantage = nn.Linear(dims[2], dims[3])
-        self.criterion = nn.MSELoss()
-        BasicNet.__init__(self, optimizer_fn, gpu)
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        x = x.view(x.size(0), -1)
-        y = F.relu(self.fc1(x))
-        phi = F.relu(self.fc2(y))
-        return phi
-
-    # Network for CartPole with actor critic
-class ActorCriticFCNet(nn.Module, ActorCriticNet):
-    def __init__(self,
-                 dims):
-        super(ActorCriticFCNet, self).__init__()
-        self.layer1 = nn.Linear(dims[0], dims[1])
-        self.fc_actor = nn.Linear(dims[1], dims[2])
-        self.fc_critic = nn.Linear(dims[1], 1)
-        BasicNet.__init__(self, None, False)
-
-    def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
-        x = x.view(x.size(0), -1)
-        phi = self.layer1(x)
-        return phi
-
-# Network for pixel Atari game with value based methods
-class NatureConvNet(nn.Module, VanillaNet):
-    def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True):
-        super(NatureConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
-        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
-        self.fc4 = nn.Linear(7 * 7 * 64, 512)
-        self.fc5 = nn.Linear(512, n_actions)
-        self.criterion = nn.MSELoss()
-        BasicNet.__init__(self, optimizer_fn, gpu)
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        y = F.relu(self.conv1(x))
-        y = F.relu(self.conv2(y))
-        y = F.relu(self.conv3(y))
-        y = y.view(y.size(0), -1)
-        y = F.relu(self.fc4(y))
-        return self.fc5(y)
-
-# Network for pixel Atari game with dueling architecture
-class DuelingNatureConvNet(nn.Module, DuelingNet):
-    def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True):
-        super(DuelingNatureConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
-        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
-        self.fc4 = nn.Linear(7 * 7 * 64, 512)
-        self.fc_advantage = nn.Linear(512, n_actions)
-        self.fc_value = nn.Linear(512, 1)
-        self.criterion = nn.MSELoss()
-        BasicNet.__init__(self, optimizer_fn, gpu)
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        y = F.relu(self.conv1(x))
-        y = F.relu(self.conv2(y))
-        y = F.relu(self.conv3(y))
-        y = y.view(y.size(0), -1)
-        phi = F.relu(self.fc4(y))
-        return phi
-
-
-
-# Network for pixel Atari game with actor critic
-class ActorCriticNatureConvNet(nn.Module, ActorCriticNet):
-    def __init__(self,
-                 in_channels,
-                 n_actions,
-                 xentropy_weight=0.01,
-                 grad_threshold=40,
-                 gpu=True):
-        super(ActorCriticNatureConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
-        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
-        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
-        self.fc4 = nn.Linear(7 * 7 * 64, 512)
-        self.fc_actor = nn.Linear(512, n_actions)
-        self.fc_critic = nn.Linear(512, 1)
-        self.xentropy_weight = xentropy_weight
-        self.grad_threshold = grad_threshold
-        BasicNet.__init__(self, optimizer_fn=None, gpu=gpu)
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        y = F.elu(self.conv1(x))
-        y = F.elu(self.conv2(y))
-        y = F.elu(self.conv3(y))
-        y = y.view(y.size(0), -1)
-        return F.elu(self.fc4(y))
-
-class OpenAIActorCriticConvNet(nn.Module, ActorCriticNet):
-    def __init__(self,
-                 in_channels,
-                 n_actions,
-                 LSTM=False):
-        super(OpenAIActorCriticConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
-        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-
-        self.LSTM = LSTM
-        hidden_units = 256
-
-        if LSTM:
-            self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
-        else:
-            self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
-
-        self.fc_actor = nn.Linear(hidden_units, n_actions)
-        self.fc_critic = nn.Linear(hidden_units, 1)
-        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
-        if LSTM:
-            self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
-            self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
-
-    def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
-        y = F.elu(self.conv1(x))
-        y = F.elu(self.conv2(y))
-        y = F.elu(self.conv3(y))
-        y = F.elu(self.conv4(y))
-        y = y.view(y.size(0), -1)
-        if self.LSTM:
-            h, c = self.layer5(y, (self.h, self.c))
-            if update_LSTM:
-                self.h = h
-                self.c = c
-            phi = h
-        else:
-            phi = F.elu(self.layer5(y))
-        return phi
-
-class OpenAIConvNet(nn.Module, VanillaNet):
-    def __init__(self,
-                 in_channels,
-                 n_actions):
-        super(OpenAIConvNet, self).__init__()
-        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
-        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
-
-        hidden_units = 256
-        self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
-        self.fc6 = nn.Linear(hidden_units, n_actions)
-
-        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=False)
-
-    def forward(self, x, update_LSTM=True):
-        x = self.to_torch_variable(x)
-        y = F.elu(self.conv1(x))
-        y = F.elu(self.conv2(y))
-        y = F.elu(self.conv3(y))
-        y = F.elu(self.conv4(y))
-        y = y.view(y.size(0), -1)
-        phi = F.elu(self.layer5(y))
-        return self.fc6(phi)
-
-class DDPGActorNet(nn.Module, BasicNet):
-    def __init__(self,
-                 state_dim,
-                 action_dim,
-                 output_gate,
-                 gpu=False):
-        super(DDPGActorNet, self).__init__()
-        self.layer1 = nn.Linear(state_dim, 400)
-        self.layer2 = nn.Linear(400, 300)
-        self.layer3 = nn.Linear(300, action_dim)
-        self.output_gate = output_gate
-        BasicNet.__init__(self, None, False, False)
-        self.init_weights()
-
-    def init_weights(self):
-        bound = 3e-3
-        self.layer3.weight.data.uniform_(-bound, bound)
-        # self.layer3.bias.data.uniform_(-bound, bound)
-
-        def fanin(size):
-            v = 1.0 / np.sqrt(size[1])
-            return torch.FloatTensor(size).uniform_(-v, v)
-
-        self.layer1.weight.data = fanin(self.layer1.weight.data.size())
-        # self.layer1.bias.data = fanin(self.layer1.bias.data.size())
-        self.layer2.weight.data = fanin(self.layer2.weight.data.size())
-        # self.layer2.bias.data = fanin(self.layer2.bias.data.size())
-
-    def forward(self, x):
-        x = self.to_torch_variable(x)
-        x = F.relu(self.layer1(x))
-        x = F.relu(self.layer2(x))
-        x = self.layer3(x)
-        # x = self.output_gate(self.layer3(x))
-        return x
-
-    def predict(self, x, to_numpy=True):
-        y = self.forward(x)
-        if to_numpy:
-            y = y.cpu().data.numpy()
-        return y
-
-class DDPGCriticNet(nn.Module, BasicNet):
-    def __init__(self,
-                 state_dim,
-                 action_dim,
-                 gpu=False):
-        super(DDPGCriticNet, self).__init__()
-        self.layer1 = nn.Linear(state_dim, 400)
-        self.layer2 = nn.Linear(400 + action_dim, 300)
-        self.layer3 = nn.Linear(300, 1)
-        BasicNet.__init__(self, None, False, False)
-        self.init_weights()
-
-    def init_weights(self):
-        bound = 3e-3
-        self.layer3.weight.data.uniform_(-bound, bound)
-        # self.layer3.bias.data.uniform_(-bound, bound)
-
-        def fanin(size):
-            v = 1.0 / np.sqrt(size[1])
-            return torch.FloatTensor(size).uniform_(-v, v)
-
-        self.layer1.weight.data = fanin(self.layer1.weight.data.size())
-        # self.layer1.bias.data = fanin(self.layer1.bias.data.size())
-        self.layer2.weight.data = fanin(self.layer2.weight.data.size())
-        # self.layer2.bias.data = fanin(self.layer2.bias.data.size())
-
-    def forward(self, x, action):
-        x = self.to_torch_variable(x)
-        action = self.to_torch_variable(action)
-        x = F.relu(self.layer1(x))
-        x = F.relu(self.layer2(torch.cat([x, action], dim=1)))
-        x = self.layer3(x)
-        return x
-
-    def predict(self, x, action):
-        return self.forward(x, action)

network/__init__.py

@@ -0,0 +1,3 @@
+from conv_network import *
+from shallow_network import *
+from continuous_action_network import *

network/continuous_action_network.py

@@ -0,0 +1,116 @@
+#######################################################################
+# 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 *
+
+class ContinuousActorCriticNet(nn.Module, BasicNet):
+    def __init__(self, state_dim, hidden_dim, action_dim):
+        super(ContinuousActorCriticNet, self).__init__()
+        hidden_size1 = 64
+        hidden_size2 = 64
+        self.fc1 = nn.Linear(state_dim, hidden_size1)
+        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
+        self.fc_mean = nn.Linear(hidden_size2, action_dim)
+        self.fc_var = nn.Linear(hidden_size2, action_dim)
+        self.fc_critic = nn.Linear(hidden_size2, 1)
+        BasicNet.__init__(self, None, False)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        phi = F.relu(self.fc2(x))
+        return phi
+
+    def predict(self, x):
+        phi = self.forward(x)
+        mean = self.fc_mean(phi)
+        var = F.softplus(self.fc_var(phi) + 1e-5)
+        value = self.fc_critic(phi)
+        return mean, var, value
+
+    def critic(self, x):
+        phi = self.forward(x)
+        return self.fc_critic(phi)
+
+class DDPGActorNet(nn.Module, BasicNet):
+    def __init__(self,
+                 state_dim,
+                 action_dim,
+                 output_gate,
+                 gpu=False):
+        super(DDPGActorNet, self).__init__()
+        self.layer1 = nn.Linear(state_dim, 400)
+        self.layer2 = nn.Linear(400, 300)
+        self.layer3 = nn.Linear(300, action_dim)
+        self.output_gate = output_gate
+        BasicNet.__init__(self, None, False, False)
+        self.init_weights()
+
+    def init_weights(self):
+        bound = 3e-3
+        self.layer3.weight.data.uniform_(-bound, bound)
+        # self.layer3.bias.data.uniform_(-bound, bound)
+
+        def fanin(size):
+            v = 1.0 / np.sqrt(size[1])
+            return torch.FloatTensor(size).uniform_(-v, v)
+
+        self.layer1.weight.data = fanin(self.layer1.weight.data.size())
+        # self.layer1.bias.data = fanin(self.layer1.bias.data.size())
+        self.layer2.weight.data = fanin(self.layer2.weight.data.size())
+        # self.layer2.bias.data = fanin(self.layer2.bias.data.size())
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        x = F.relu(self.layer1(x))
+        x = F.relu(self.layer2(x))
+        x = self.layer3(x)
+        # x = self.output_gate(self.layer3(x))
+        return x
+
+    def predict(self, x, to_numpy=True):
+        y = self.forward(x)
+        if to_numpy:
+            y = y.cpu().data.numpy()
+        return y
+
+class DDPGCriticNet(nn.Module, BasicNet):
+    def __init__(self,
+                 state_dim,
+                 action_dim,
+                 gpu=False):
+        super(DDPGCriticNet, self).__init__()
+        self.layer1 = nn.Linear(state_dim, 400)
+        self.layer2 = nn.Linear(400 + action_dim, 300)
+        self.layer3 = nn.Linear(300, 1)
+        BasicNet.__init__(self, None, False, False)
+        self.init_weights()
+
+    def init_weights(self):
+        bound = 3e-3
+        self.layer3.weight.data.uniform_(-bound, bound)
+        # self.layer3.bias.data.uniform_(-bound, bound)
+
+        def fanin(size):
+            v = 1.0 / np.sqrt(size[1])
+            return torch.FloatTensor(size).uniform_(-v, v)
+
+        self.layer1.weight.data = fanin(self.layer1.weight.data.size())
+        # self.layer1.bias.data = fanin(self.layer1.bias.data.size())
+        self.layer2.weight.data = fanin(self.layer2.weight.data.size())
+        # self.layer2.bias.data = fanin(self.layer2.bias.data.size())
+
+    def forward(self, x, action):
+        x = self.to_torch_variable(x)
+        action = self.to_torch_variable(action)
+        x = F.relu(self.layer1(x))
+        x = F.relu(self.layer2(torch.cat([x, action], dim=1)))
+        x = self.layer3(x)
+        return x
+
+    def predict(self, x, action):
+        return self.forward(x, action)

network/conv_network.py

@@ -0,0 +1,147 @@
+#######################################################################
+# 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 *
+
+# Network for pixel Atari game with value based methods
+class NatureConvNet(nn.Module, VanillaNet):
+    def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True):
+        super(NatureConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
+        self.fc4 = nn.Linear(7 * 7 * 64, 512)
+        self.fc5 = nn.Linear(512, n_actions)
+        self.criterion = nn.MSELoss()
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        y = F.relu(self.conv1(x))
+        y = F.relu(self.conv2(y))
+        y = F.relu(self.conv3(y))
+        y = y.view(y.size(0), -1)
+        y = F.relu(self.fc4(y))
+        return self.fc5(y)
+
+# Network for pixel Atari game with dueling architecture
+class DuelingNatureConvNet(nn.Module, DuelingNet):
+    def __init__(self, in_channels, n_actions, optimizer_fn=None, gpu=True):
+        super(DuelingNatureConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
+        self.fc4 = nn.Linear(7 * 7 * 64, 512)
+        self.fc_advantage = nn.Linear(512, n_actions)
+        self.fc_value = nn.Linear(512, 1)
+        self.criterion = nn.MSELoss()
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        y = F.relu(self.conv1(x))
+        y = F.relu(self.conv2(y))
+        y = F.relu(self.conv3(y))
+        y = y.view(y.size(0), -1)
+        phi = F.relu(self.fc4(y))
+        return phi
+
+
+# Network for pixel Atari game with actor critic
+class ActorCriticNatureConvNet(nn.Module, ActorCriticNet):
+    def __init__(self,
+                 in_channels,
+                 n_actions,
+                 xentropy_weight=0.01,
+                 grad_threshold=40,
+                 gpu=True):
+        super(ActorCriticNatureConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=8, stride=4)
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=4, stride=2)
+        self.conv3 = nn.Conv2d(64, 64, kernel_size=3, stride=1)
+        self.fc4 = nn.Linear(7 * 7 * 64, 512)
+        self.fc_actor = nn.Linear(512, n_actions)
+        self.fc_critic = nn.Linear(512, 1)
+        self.xentropy_weight = xentropy_weight
+        self.grad_threshold = grad_threshold
+        BasicNet.__init__(self, optimizer_fn=None, gpu=gpu)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        y = F.elu(self.conv1(x))
+        y = F.elu(self.conv2(y))
+        y = F.elu(self.conv3(y))
+        y = y.view(y.size(0), -1)
+        return F.elu(self.fc4(y))
+
+class OpenAIActorCriticConvNet(nn.Module, ActorCriticNet):
+    def __init__(self,
+                 in_channels,
+                 n_actions,
+                 LSTM=False):
+        super(OpenAIActorCriticConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+
+        self.LSTM = LSTM
+        hidden_units = 256
+
+        if LSTM:
+            self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
+        else:
+            self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
+
+        self.fc_actor = nn.Linear(hidden_units, n_actions)
+        self.fc_critic = nn.Linear(hidden_units, 1)
+        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
+        if LSTM:
+            self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
+            self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
+
+    def forward(self, x, update_LSTM=True):
+        x = self.to_torch_variable(x)
+        y = F.elu(self.conv1(x))
+        y = F.elu(self.conv2(y))
+        y = F.elu(self.conv3(y))
+        y = F.elu(self.conv4(y))
+        y = y.view(y.size(0), -1)
+        if self.LSTM:
+            h, c = self.layer5(y, (self.h, self.c))
+            if update_LSTM:
+                self.h = h
+                self.c = c
+            phi = h
+        else:
+            phi = F.elu(self.layer5(y))
+        return phi
+
+class OpenAIConvNet(nn.Module, VanillaNet):
+    def __init__(self,
+                 in_channels,
+                 n_actions):
+        super(OpenAIConvNet, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+        self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
+
+        hidden_units = 256
+        self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
+        self.fc6 = nn.Linear(hidden_units, n_actions)
+
+        BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=False)
+
+    def forward(self, x, update_LSTM=True):
+        x = self.to_torch_variable(x)
+        y = F.elu(self.conv1(x))
+        y = F.elu(self.conv2(y))
+        y = F.elu(self.conv3(y))
+        y = F.elu(self.conv4(y))
+        y = y.view(y.size(0), -1)
+        phi = F.elu(self.layer5(y))
+        return self.fc6(phi)

network/network.py

@@ -0,0 +1,72 @@
+#######################################################################
+# 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 torch
+from torch.autograd import Variable
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+# Base class for all kinds of network
+class BasicNet:
+    def __init__(self, optimizer_fn, gpu, LSTM=False):
+        if optimizer_fn is not None:
+            self.optimizer = optimizer_fn(self.parameters())
+        self.gpu = gpu and torch.cuda.is_available()
+        self.LSTM = LSTM
+        if self.gpu:
+            self.cuda()
+
+    def to_torch_variable(self, x, dtype='float32'):
+        if isinstance(x, Variable):
+            return x
+        if not isinstance(x, torch.FloatTensor):
+            x = torch.from_numpy(np.asarray(x, dtype=dtype))
+        if self.gpu:
+            x = x.cuda()
+        return Variable(x)
+
+    def reset(self, terminal):
+        if not self.LSTM:
+            return
+        if terminal:
+            self.h.data.zero_()
+            self.c.data.zero_()
+        self.h = Variable(self.h.data)
+        self.c = Variable(self.c.data)
+
+# Base class for value based methods
+class VanillaNet(BasicNet):
+    def predict(self, x, to_numpy=False):
+        y = self.forward(x)
+        if to_numpy:
+            y = y.cpu().data.numpy()
+        return y
+
+# Base class for actor critic method
+class ActorCriticNet(BasicNet):
+    def predict(self, x):
+        phi = self.forward(x, True)
+        pre_prob = self.fc_actor(phi)
+        prob = F.softmax(pre_prob)
+        log_prob = F.log_softmax(pre_prob)
+        value = self.fc_critic(phi)
+        return prob, log_prob, value
+
+    def critic(self, x):
+        phi = self.forward(x, False)
+        return self.fc_critic(phi)
+
+# Base class for dueling architecture
+class DuelingNet(BasicNet):
+    def predict(self, x, to_numpy=False):
+        phi = self.forward(x)
+        value = self.fc_value(phi)
+        advantange = self.fc_advantage(phi)
+        q = value.expand_as(advantange) + (advantange - advantange.mean(1).expand_as(advantange))
+        if to_numpy:
+            return q.cpu().data.numpy()
+        return q

network/shallow_network.py

@@ -0,0 +1,64 @@
+#######################################################################
+# 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 *
+
+# Network for CartPole with value based methods
+class FCNet(nn.Module, VanillaNet):
+    def __init__(self, dims, optimizer_fn=None, gpu=True):
+        super(FCNet, self).__init__()
+        self.fc1 = nn.Linear(dims[0], dims[1])
+        self.fc2 = nn.Linear(dims[1], dims[2])
+        self.fc3 = nn.Linear(dims[2], dims[3])
+        self.criterion = nn.MSELoss()
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        y = F.relu(self.fc1(x))
+        y = F.relu(self.fc2(y))
+        y = self.fc3(y)
+        return y
+
+# Network for CartPole with dueling architecture
+class DuelingFCNet(nn.Module, DuelingNet):
+    def __init__(self, dims, optimizer_fn=None, gpu=True):
+        super(DuelingFCNet, self).__init__()
+        self.fc1 = nn.Linear(dims[0], dims[1])
+        self.fc2 = nn.Linear(dims[1], dims[2])
+        self.fc_value = nn.Linear(dims[2], 1)
+        self.fc_advantage = nn.Linear(dims[2], dims[3])
+        self.criterion = nn.MSELoss()
+        BasicNet.__init__(self, optimizer_fn, gpu)
+
+    def forward(self, x):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        y = F.relu(self.fc1(x))
+        phi = F.relu(self.fc2(y))
+        return phi
+
+# Network for CartPole with actor critic
+class ActorCriticFCNet(nn.Module, ActorCriticNet):
+    def __init__(self, state_dim, action_dim):
+        super(ActorCriticFCNet, self).__init__()
+        hidden_size1 = 50
+        hidden_size2 = 200
+        self.fc1 = nn.Linear(state_dim, hidden_size1)
+        self.fc2 = nn.Linear(hidden_size1, hidden_size2)
+        self.fc_actor = nn.Linear(hidden_size2, action_dim)
+        self.fc_critic = nn.Linear(hidden_size2, 1)
+        BasicNet.__init__(self, None, False)
+
+    def forward(self, x, update_LSTM=True):
+        x = self.to_torch_variable(x)
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        phi = self.fc2(x)
+        return phi
+
+

utils/__init__.py

@@ -0,0 +1,5 @@
+from config import *
+try:
+    from tf_logger import Logger
+except:
+    from vanilla_logger import Logger

utils/config.py

@@ -25,3 +25,5 @@ class Config:
         self.worker = None
         self.update_interval = 1
         self.gradient_clip = 40
+        self.entropy_weight = 0.01
+        self.gae_tau = 1.0