Continuous A3C

.gitignore

@@ -10,6 +10,7 @@ upload.py
 data
 draw_*
 log
+evaluation_log
 figure
 to_plot
 

README.md

@@ -11,6 +11,7 @@ Implemented algorithms:
 * Async One-Step Q-Learning
 * Async One-Step Sarsa 
 * Async N-Step Q-Learning
+* Continuous A3C
 * Deep Deterministic Policy Gradient (DDPG)
 
 # Curves
@@ -25,7 +26,7 @@ Xeon E5-2620 v3 and Titan X. For Breakout, test is triggered every 1000 episodes
 In total, 16M frames cost about 4 days and 10 hours. For Pong, test is triggered 
 every 10 episodes with no repetition. In total, 4M frames cost about 18 hours.
 
-## A3C, etc.
+## Discrete A3C
 
 ![Loading...](https://raw.githubusercontent.com/ShangtongZhang/DeepRL/master/images/A3C-Pong.png)
 ![Loading...](https://raw.githubusercontent.com/ShangtongZhang/DeepRL/master/images/Async-Pong.png)
@@ -37,6 +38,11 @@ Training of A3C took about 2 hours (16 processes) in a server with two Xeon E5-2
 Those value based async methods do work but I don't know how to make them stable.
 This is the test curve. Test is triggered in a separate deterministic test process every 50K frames.
 
+## Continuous A3C
+![Loading...](https://raw.githubusercontent.com/ShangtongZhang/DeepRL/master/images/Continuous-A3C.png)
+
+Sometimes _Bipedal Walker_ may run into _NAN_, I'm still not able to totally solve it. And continuous A3C is very sensible to hyper parameters.
+
 # Dependency
 * Open AI gym
 * PyTorch
@@ -56,6 +62,8 @@ Detailed usage and all training details can be found in ```main.py```
 * [HOGWILD!: A Lock-Free Approach to Parallelizing Stochastic Gradient Descent](https://arxiv.org/abs/1106.5730)
 * [Deterministic Policy Gradient Algorithms](http://proceedings.mlr.press/v32/silver14.pdf)
 * [Continuous control with deep reinforcement learning](https://arxiv.org/abs/1509.02971)
+* [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438)
 * [transedward/pytorch-dqn](https://github.com/transedward/pytorch-dqn)
 * [ikostrikov/pytorch-a3c](https://github.com/ikostrikov/pytorch-a3c)
 * [ghliu/pytorch-ddpg](https://github.com/ghliu/pytorch-ddpg)
+* [MorvanZhou/Reinforcement-learning-with-tensorflow](https://github.com/MorvanZhou/Reinforcement-learning-with-tensorflow)

agent/async_agent.py

@@ -29,6 +29,7 @@ def evaluate(config, task, learning_network):
     test_rewards = []
     test_points = []
     worker = config.worker(config, learning_network, None)
+    # config.logger = Logger('./evaluation_log', gym.logger)
     while True:
         steps = config.total_steps.value
         if steps % config.test_interval == 0:

async_worker/continuous_actor_critic.py

@@ -11,32 +11,46 @@ import torch.nn as nn
 class ContinuousAdvantageActorCritic:
     def __init__(self, config, learning_network, target_network):
         self.config = config
-        self.optimizer = config.optimizer_fn(learning_network.parameters())
+        # self.optimizer = config.optimizer_fn(learning_network.parameters())
+        self.optimizer = config.optimizer_fn(learning_network.actor_params)
+        self.critic_optimizer = config.critic_optimizer_fn(learning_network.critic_params)
         self.worker_network = config.network_fn()
         self.worker_network.load_state_dict(learning_network.state_dict())
         self.task = config.task_fn()
         self.policy = config.policy_fn()
         self.learning_network = learning_network
+        self.counter = 0
 
     def episode(self, deterministic=False):
         config = self.config
         state = self.task.reset()
+        state = config.state_shift_fn(state)
         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]))
+            mean, std, value = self.worker_network.predict(np.stack([state]))
             action = self.policy.sample(mean.data.numpy().flatten(),
-                                        var.data.numpy().flatten(),
-                                        deterministic)
+                                        std.data.numpy().flatten(),
+                                        False)
+            action = self.config.action_shift_fn(action)
             next_state, reward, terminal, _ = self.task.step(action)
+            next_state = config.state_shift_fn(next_state)
+
+            # if deterministic:
+            #     self.config.logger.scalar_summary('reward', reward, self.counter)
+            #     self.config.logger.histo_summary('std', std.data.numpy(), self.counter)
+            #     self.config.logger.histo_summary('mean', mean.data.numpy(), self.counter)
+            #     self.config.logger.histo_summary('action', action, self.counter)
+            #     self.config.logger.scalar_summary('steps', steps, self.counter)
+            #     self.config.logger.histo_summary('states', state, self.counter)
+            #     self.counter += 1
 
             steps += 1
             total_reward += reward
-            if not deterministic:
-                reward = np.clip(reward, -1, 1)
+            reward = config.reward_shift_fn(reward)
 
             if deterministic:
                 if terminal:
@@ -44,19 +58,20 @@ class ContinuousAdvantageActorCritic:
                 state = next_state
                 continue
 
-            pending.append([mean, var, value, action, reward])
+            pending.append([mean, std, value, action, reward])
             with config.steps_lock:
                 config.total_steps.value += 1
 
             if terminal or len(pending) >= config.update_interval:
-                loss = 0
+                critic_loss = 0
+                actor_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]
+                    mean, std, value, action, reward = pending[i]
                     if i == len(pending) - 1:
                         delta = reward + config.discount * R - value.data
                     else:
@@ -64,28 +79,30 @@ class ContinuousAdvantageActorCritic:
                     GAE = config.discount * config.gae_tau * GAE + delta
 
                     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
+                    log_prob = -(action - mean).pow(2) / (2 * std.pow(2)) -\
+                               std.log() - 0.5 * (2 * pi).log().expand_as(std)
+                    actor_loss += -torch.sum(log_prob) * Variable(GAE)
+                    entropy = 0.5 + std.log() + 0.5 * (2 * pi).log().expand_as(std)
+                    actor_loss += -config.entropy_weight * entropy.sum()
 
                     R = reward + config.discount * R
-                    loss += 0.5 * (Variable(R) - value).pow(2)
+                    critic_loss += 0.5 * (Variable(R) - value).pow(2)
 
                 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.critic_optimizer.zero_grad()
+                actor_loss.backward()
+                critic_loss.backward()
+                nn.utils.clip_grad_norm(self.worker_network.actor_params, config.gradient_clip)
+                nn.utils.clip_grad_norm(self.worker_network.critic_params, config.gradient_clip)
                 for param, worker_param in zip(
                         self.learning_network.parameters(), self.worker_network.parameters()):
                     if param.grad is not None:
                         break
                     param._grad = worker_param.grad
                 self.optimizer.step()
+                self.critic_optimizer.step()
                 self.worker_network.load_state_dict(self.learning_network.state_dict())
                 self.worker_network.reset(terminal)
 

component/policy.py

@@ -48,10 +48,10 @@ class SamplePolicy:
         pass
 
 class GaussianPolicy:
-    def sample(self, mean, var, deterministic=False):
+    def sample(self, mean, std, deterministic=False):
         if deterministic:
             return mean
-        return mean + np.sqrt(var) * np.random.randn(*mean.shape)
+        return mean + std * np.random.randn(*mean.shape)
 
     def update_epsilon(self):
         pass

component/task.py

@@ -78,7 +78,7 @@ class PixelAtari(BasicTask):
 
 class Pendulum(BasicTask):
     name = 'Pendulum-v0'
-    success_threshold = 200
+    success_threshold = -10
 
     def __init__(self):
         BasicTask.__init__(self)
@@ -88,14 +88,35 @@ class Pendulum(BasicTask):
         self.state_dim = self.env.observation_space.shape[0]
 
     def step(self, action):
-        # 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
 
+class MountainCarContinuous(BasicTask):
+    name = 'MountainCarContinuous-v0'
+    success_threshold = 90
+
+    def __init__(self):
+        BasicTask.__init__(self)
+        self.env = gym.make(self.name)
+        self.env._max_episode_steps = sys.maxsize
+        self.action_dim = self.env.action_space.shape[0]
+        self.state_dim = self.env.observation_space.shape[0]
+
+    def normalize_state(self, state):
+        state = (state - self.env.unwrapped.low_state) / \
+                (self.env.unwrapped.high_state - self.env.unwrapped.low_state)
+        state = state * 2 - 1
+        return state
+
+    def step(self, action):
+        action = np.clip(action, -1, 1)
+        next_state, reward, done, info = self.env.step(action)
+        return next_state, reward, done, info
+
 class BipedalWalker(BasicTask):
     name = 'BipedalWalker-v2'
-    success_threshold = 2000
+    success_threshold = 300
 
     def __init__(self):
         BasicTask.__init__(self)

main.py

@@ -67,19 +67,47 @@ def a3c_cart_pole():
 def a3c_pendulum():
     config = Config()
     config.task_fn = lambda: Pendulum()
+    config.reward_shift_fn = lambda reward: reward / 10
+    # config.task_fn = lambda: MountainCarContinuous()
     task = config.task_fn()
     config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.0001)
+    config.critic_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])
+        task.state_dim, task.action_dim, 2, F.tanh)
     config.policy_fn = lambda: GaussianPolicy()
     config.worker = ContinuousAdvantageActorCritic
     config.discount = 0.99
     config.max_episode_length = 200
-    config.num_workers = 16
+    config.num_workers = 8
+    config.update_interval = 5
+    config.test_interval = 1
+    config.test_repetitions = 5
+    config.entropy_weight = 0.0001
+    config.gradient_clip = 40
+    config.logger = Logger('./log', gym.logger)
+    agent = AsyncAgent(config)
+    agent.run()
+
+def a3c_walker():
+    config = Config()
+    config.task_fn = lambda: BipedalWalker()
+    shifter = Shifter()
+    config.state_shift_fn = lambda state: shifter(state)
+    task = config.task_fn()
+    config.optimizer_fn = lambda params: torch.optim.Adam(params, 0.0001)
+    config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
+    config.network_fn = lambda: ContinuousActorCriticNet(
+        task.state_dim, task.action_dim, 1, F.tanh)
+    config.policy_fn = lambda: GaussianPolicy()
+    config.worker = ContinuousAdvantageActorCritic
+    config.discount = 0.99
+    config.max_episode_length = 999
+    config.num_workers = 8
     config.update_interval = 20
     config.test_interval = 1
-    config.test_repetitions = 1
-    config.entropy_weight = 0.0001
+    config.test_repetitions = 5
+    config.entropy_weight = 0.01
+    config.gradient_clip = 30
     config.logger = Logger('./log', gym.logger)
     agent = AsyncAgent(config)
     agent.run()
@@ -209,11 +237,12 @@ if __name__ == '__main__':
     # dqn_cart_pole()
     # async_cart_pole()
     # a3c_cart_pole()
-    # a3c_pendulum()
+    a3c_pendulum()
+    # a3c_walker()
 
     # dqn_pixel_atari('PongNoFrameskip-v3')
     # async_pixel_atari('PongNoFrameskip-v3')
-    a3c_pixel_atari('PongNoFrameskip-v3')
+    # a3c_pixel_atari('PongNoFrameskip-v3')
 
     # dqn_pixel_atari('BreakoutNoFrameskip-v3')
     # async_pixel_atari('BreakoutNoFrameskip-v3')

network/continuous_action_network.py

@@ -7,34 +7,41 @@
 from network import *
 
 class ContinuousActorCriticNet(nn.Module, BasicNet):
-    def __init__(self, state_dim, hidden_dim, action_dim):
+    def __init__(self, state_dim, action_dim, action_scale, action_gate):
         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)
+        actor_hidden = 200
+        critic_hidden = 100
+        self.fc_actor = nn.Linear(state_dim, actor_hidden)
+        self.fc_mean = nn.Linear(actor_hidden, action_dim)
+        self.fc_std = nn.Linear(actor_hidden, action_dim)
+        self.action_scale = action_scale
+        self.action_gate = action_gate
+        self.actor_params = list(self.fc_actor.parameters()) + \
+                            list(self.fc_mean.parameters()) + \
+                            list(self.fc_std.parameters())
+
+        self.fc_critic = nn.Linear(state_dim, critic_hidden)
+        self.fc_value = nn.Linear(critic_hidden, 1)
+        self.critic_params = list(self.fc_critic.parameters()) + \
+                             list(self.fc_value.parameters())
+
         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
+        x = self.to_torch_variable(x)
+        value = self.critic(x)
+
+        x = F.relu(self.fc_actor(x))
+        mean = self.action_scale * self.action_gate(self.fc_mean(x))
+        std = F.softplus(self.fc_std(x) + 1e-5)
+
+        return mean, std, value
 
     def critic(self, x):
-        phi = self.forward(x)
-        return self.fc_critic(phi)
+        x = self.to_torch_variable(x)
+        x = F.relu(self.fc_critic(x))
+        x = self.fc_value(x)
+        return x
 
 class DDPGActorNet(nn.Module, BasicNet):
     def __init__(self,

utils/__init__.py

@@ -1,4 +1,5 @@
 from config import *
+from shifter import *
 try:
     from tf_logger import Logger
 except:

utils/config.py

@@ -8,6 +8,7 @@ class Config:
     def __init__(self):
         self.task_fn = None
         self.optimizer_fn = None
+        self.critic_optimizer_fn = None
         self.network_fn = None
         self.policy_fn = None
         self.replay_fn = None
@@ -27,3 +28,6 @@ class Config:
         self.gradient_clip = 40
         self.entropy_weight = 0.01
         self.gae_tau = 1.0
+        self.reward_shift_fn = lambda r: r
+        self.state_shift_fn = lambda s: s
+        self.action_shift_fn = lambda a: a

utils/shifter.py

@@ -0,0 +1,28 @@
+# Adapted from https://github.com/kvfrans/parallel-trpo/blob/master/utils.py
+class Shifter:
+    def __init__(self, filter_mean=True):
+        self.m = 0
+        self.v = 0
+        self.n = 0.
+        self.filter_mean = filter_mean
+
+    def state_dict(self):
+        return {'m': self.m,
+                'v': self.v,
+                'n': self.n}
+
+    def load_state_dict(self, saved):
+        self.m = saved['m']
+        self.v = saved['v']
+        self.n = saved['n']
+
+    def __call__(self, o):
+        self.m = self.m * (self.n / (self.n + 1)) + o * 1 / (1 + self.n)
+        self.v = self.v * (self.n / (self.n + 1)) + (o - self.m) ** 2 * 1 / (1 + self.n)
+        self.std = (self.v + 1e-6) ** .5  # std
+        self.n += 1
+        if self.filter_mean:
+            o_ = (o - self.m) / self.std
+        else:
+            o_ = o / self.std
+        return o_