Continuous A3C

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)