Support async PPO, which doesn't work

agent/PPO_agent.py

@@ -15,125 +15,153 @@ import os
 import time
 
 class PPOWorker:
-    def __init__(self, config, shared_network, shared_state_shifter):
+    def __init__(self, config, shared_network, extra):
         self.config = config
-        # self.shared_network = shared_network
-        # self.local_netwrok = config.network_fn()
         self.task = config.task_fn()
         self.policy = config.policy_fn()
 
-        self.actor_net = config.actor_network_fn()
-        self.critic_net = config.critic_network_fn()
+        self.shared_network = shared_network
+        self.actor_opt = config.actor_optimizer_fn(shared_network.actor.parameters())
+        self.critic_opt = config.critic_optimizer_fn(shared_network.critic.parameters())
 
-        self.actor_opt = config.actor_optimizer_fn(self.actor_net.parameters())
-        self.critic_opt = config.critic_optimizer_fn(self.critic_net.parameters())
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(shared_network.state_dict())
 
-        # self.shared_state_shifter()
+        self.shared_state_normalizer = extra[0]
+        self.state_normalizer = StaticNormalizer(self.task.state_dim)
+        self.shared_reward_normalizer = extra[1]
+        self.reward_normalizer = StaticNormalizer(1)
 
-    def rollout(self, deterministic=False):
+    def episode(self, deterministic=False):
         config = self.config
+        self.state_normalizer.offline_stats.load(self.shared_state_normalizer.offline_stats)
+        self.reward_normalizer.offline_stats.load(self.shared_reward_normalizer.offline_stats)
+
         replay = config.replay_fn()
         state = self.task.reset()
+        state = self.state_normalizer(state)
+
         episode_length = 0
-        episode_reward = 0
-        reward_history = []
+        batched_rewards = 0
+        batched_steps = 0
+        batched_episode = 0
+
+        actor_net = self.worker_network.actor
+        critic_net = self.worker_network.critic
 
         actor_net_old = config.actor_network_fn()
-        actor_net_old.load_state_dict(self.actor_net.state_dict())
-        ep_count = 0
+        actor_net_old.load_state_dict(actor_net.state_dict())
 
-        while True:
-            # self.local_netwrok.load_state_dict(self.shared_network.state_dict())
+        self.worker_network.load_state_dict(self.shared_network.state_dict())
 
-            batched_episode = 0
-            while not replay.full():
-                states = []
-                actions = []
-                rewards = []
-                values = []
-                returns = []
-                advantages = []
+        while not replay.full():
+            states = []
+            actions = []
+            rewards = []
+            values = []
+            returns = []
+            advantages = []
 
-                for i in range(config.rollout_length):
-                    mean, std, log_std = self.actor_net.predict(np.stack([state]))
-                    value = self.critic_net.predict(np.stack([state]))
-                    action = self.policy.sample(mean.data.numpy().flatten(), std.data.numpy().flatten(), deterministic)
-                    action = self.config.action_shift_fn(action)
-                    states.append(state)
-                    actions.append(action)
-                    values.append(value)
-                    state, reward, done, _ = self.task.step(action)
-                    episode_reward += reward
-                    episode_length += 1
-                    done = (done or (config.max_episode_length and episode_length > config.max_episode_length))
-                    reward = self.config.reward_shift_fn(reward)
-                    rewards.append(reward)
+            for i in range(config.rollout_length):
+                mean, std, log_std = actor_net.predict(np.stack([state]))
+                value = critic_net.predict(np.stack([state]))
+                action = self.policy.sample(mean.data.numpy().flatten(), std.data.numpy().flatten(), deterministic)
+                action = self.config.action_shift_fn(action)
+                states.append(state)
+                actions.append(action)
+                values.append(value)
+                state, reward, done, _ = self.task.step(action)
+                state = self.state_normalizer(state)
+                done = (done or (config.max_episode_length and episode_length > config.max_episode_length))
 
-                    if done:
-                        episode_length = 0
-                        batched_episode += 1
-                        reward_history.append(episode_reward)
-                        # print episode_reward, np.mean(reward_history[-100:])
-                        # episode_reward = 0
-                        state = self.task.reset()
-                        break
+                batched_rewards += reward
+                batched_steps += 1
+                episode_length += 1
 
-                R = torch.zeros((1, 1))
-                if not done:
-                    R = self.critic_net.predict(np.stack([state])).data
 
-                values.append(Variable(R))
-                A = Variable(torch.zeros((1, 1)))
-                for i in reversed(range(len(rewards))):
-                    R = Variable(torch.FloatTensor([[rewards[i]]]))
-                    ret = R + self.config.discount * values[i + 1]
-                    A = ret - values[i] + self.config.discount * self.config.gae_tau * A
-                    advantages.append(A.detach())
-                    returns.append(ret.detach())
-                advantages = list(reversed(advantages))
-                returns = list(reversed(returns))
-                replay.feed([states, actions, returns, advantages])
 
-            episode_reward /= batched_episode
-            print ep_count, episode_reward
-            ep_count += 1
+                reward = np.asscalar(self.reward_normalizer(np.array([reward])))
+                rewards.append(reward)
 
-            for _ in np.arange(self.config.optimize_epochs):
-                # local_network.load_state_dict(self.shared_network.state_dict())
+                if done:
+                    episode_length = 0
+                    batched_episode += 1
+                    state = self.task.reset()
+                    state = self.state_normalizer(state)
+                    break
 
-                states, actions, returns, advantages = replay.sample()
-                states = self.actor_net.to_torch_variable(np.stack(states))
-                actions = self.actor_net.to_torch_variable(np.stack(actions))
-                returns = torch.cat(returns, 0)
-                advantages = torch.cat(advantages, 0)
-                advantages = (advantages - advantages.mean().expand_as(advantages)) / advantages.std().expand_as(advantages)
+            R = torch.zeros((1, 1))
+            if not done:
+                R = critic_net.predict(np.stack([state])).data
 
-                mean_old, std_old, log_std_old = actor_net_old.predict(states)
-                probs_old = self.actor_net.log_density(actions, mean_old, log_std_old, std_old)
-                mean, std, log_std = self.actor_net.predict(states)
-                probs = self.actor_net.log_density(actions, mean, log_std, std)
-                ratio = (probs - probs_old).exp()
-                obj = ratio * advantages
-                obj_clipped = ratio.clamp(1.0 - self.config.ppo_ratio_clip, 1.0 + self.config.ppo_ratio_clip) * advantages
-                policy_loss = -torch.min(obj, obj_clipped).mean(0)
-                if config.entropy_weight:
-                    policy_loss += -config.entropy_weight * self.actor_net.entropy(std)
+            values.append(Variable(R))
+            A = Variable(torch.zeros((1, 1)))
+            for i in reversed(range(len(rewards))):
+                R = Variable(torch.FloatTensor([[rewards[i]]]))
+                ret = R + self.config.discount * values[i + 1]
+                A = ret - values[i] + self.config.discount * self.config.gae_tau * A
+                advantages.append(A.detach())
+                returns.append(ret.detach())
+            advantages = list(reversed(advantages))
+            returns = list(reversed(returns))
+            replay.feed([states, actions, returns, advantages])
 
-                v = self.critic_net.predict(states)
-                value_loss = 0.5 * (returns - v).pow(2).mean()
+        batched_rewards /= batched_episode
+        batched_steps /= batched_episode
 
-                self.critic_opt.zero_grad()
-                value_loss.backward()
-                nn.utils.clip_grad_norm(self.critic_net.parameters(), config.gradient_clip)
-                self.critic_opt.step()
+        if deterministic:
+            return batched_steps, batched_rewards
 
-                actor_net_old.load_state_dict(self.actor_net.state_dict())
-                self.actor_opt.zero_grad()
-                policy_loss.backward()
-                nn.utils.clip_grad_norm(self.actor_net.parameters(), config.gradient_clip)
-                self.actor_opt.step()
+        with config.steps_lock:
+            config.total_steps.value += replay.memory_size
 
-            replay.clear()
+        self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats)
+        self.state_normalizer.online_stats.zero()
+
+        self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats)
+        self.reward_normalizer.online_stats.zero()
+
+        for _ in np.arange(self.config.optimize_epochs):
+            self.worker_network.load_state_dict(self.shared_network.state_dict())
+
+            states, actions, returns, advantages = replay.sample()
+            states = actor_net.to_torch_variable(np.stack(states))
+            actions = actor_net.to_torch_variable(np.stack(actions))
+            returns = torch.cat(returns, 0)
+            advantages = torch.cat(advantages, 0)
+            advantages = (advantages - advantages.mean().expand_as(advantages)) / advantages.std().expand_as(advantages)
+
+            mean_old, std_old, log_std_old = actor_net_old.predict(states)
+            probs_old = actor_net.log_density(actions, mean_old, log_std_old, std_old)
+            mean, std, log_std = actor_net.predict(states)
+            probs = actor_net.log_density(actions, mean, log_std, std)
+            ratio = (probs - probs_old).exp()
+            obj = ratio * advantages
+            obj_clipped = ratio.clamp(1.0 - self.config.ppo_ratio_clip, 1.0 + self.config.ppo_ratio_clip) * advantages
+            policy_loss = -torch.min(obj, obj_clipped).mean(0)
+            if config.entropy_weight:
+                policy_loss += -config.entropy_weight * actor_net.entropy(std)
+
+            v = critic_net.predict(states)
+            value_loss = 0.5 * (returns - v).pow(2).mean()
+            actor_net_old.load_state_dict(self.actor_net.state_dict())
+
+            self.worker_network.zero_grad()
+            self.actor_opt.zero_grad()
+            self.critic_opt.zero_grad()
+            policy_loss.backward()
+            value_loss.backward()
+            nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+
+            for param, worker_param in zip(
+                    self.shared_network.parameters(), self.worker_network.parameters()):
+                if param.grad is not None:
+                    break
+                param._grad = worker_param.grad
+            self.actor_opt.step()
+            self.critic_opt.step()
+
+        return batched_steps, batched_rewards
 
 class PPOAgent:
     def __init__(self, config):

agent/async_agent.py

@@ -14,8 +14,8 @@ import pickle
 import os
 import time
 
-def train(id, config, learning_network, target_network):
-    worker = config.worker(config, learning_network, target_network)
+def train(id, config, learning_network, extra):
+    worker = config.worker(config, learning_network, extra)
     episode = 0
     rewards = []
     while not config.stop_signal.value:
@@ -64,14 +64,14 @@ class AsyncAgent:
         task = config.task_fn()
         learning_network = config.network_fn()
         learning_network.share_memory()
-        target_network = config.network_fn()
-        target_network.share_memory()
-        target_network.load_state_dict(learning_network.state_dict())
 
         os.environ['OMP_NUM_THREADS'] = '1'
         if config.worker == NStepQLearning or config.worker == OneStepQLearning or config.worker == OneStepSarsa:
+            target_network = config.network_fn()
+            target_network.share_memory()
+            target_network.load_state_dict(learning_network.state_dict())
             extra = target_network
-        elif config.worker == ContinuousAdvantageActorCritic:
+        elif config.worker == ContinuousAdvantageActorCritic or config.worker == ProximalPolicyOptimization:
             state_normalizer = StaticNormalizer(task.state_dim)
             reward_normalizer = StaticNormalizer(1)
             extra = [state_normalizer, reward_normalizer]

async_worker/__init__.py

@@ -2,4 +2,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 *
+from one_step_q import *
+from ppo import *

async_worker/continuous_actor_critic.py

@@ -45,7 +45,6 @@ class ContinuousAdvantageActorCritic:
             next_state, reward, terminal, _ = self.task.step(action)
             terminal = (terminal or (config.max_episode_length and steps >= config.max_episode_length))
             next_state = self.state_normalizer(next_state)
-            # next_state = config.state_shift_fn(next_state)
 
             # if deterministic:
             #     self.config.logger.scalar_summary('reward', reward, self.counter)

async_worker/ppo.py

@@ -0,0 +1,163 @@
+#######################################################################
+# 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.multiprocessing as mp
+from network import *
+from utils import *
+from component import *
+from async_worker import *
+import pickle
+import os
+import time
+
+class ProximalPolicyOptimization:
+    def __init__(self, config, shared_network, extra):
+        self.config = config
+        self.task = config.task_fn()
+        self.policy = config.policy_fn()
+
+        self.shared_network = shared_network
+        self.actor_opt = config.actor_optimizer_fn(shared_network.actor.parameters())
+        self.critic_opt = config.critic_optimizer_fn(shared_network.critic.parameters())
+
+        self.worker_network = config.network_fn()
+        self.worker_network.load_state_dict(shared_network.state_dict())
+
+        self.shared_state_normalizer = extra[0]
+        self.state_normalizer = StaticNormalizer(self.task.state_dim)
+        self.shared_reward_normalizer = extra[1]
+        self.reward_normalizer = StaticNormalizer(1)
+
+    def episode(self, deterministic=False):
+        config = self.config
+        self.state_normalizer.offline_stats.load(self.shared_state_normalizer.offline_stats)
+        self.reward_normalizer.offline_stats.load(self.shared_reward_normalizer.offline_stats)
+
+        replay = config.replay_fn()
+        state = self.task.reset()
+        state = self.state_normalizer(state)
+
+        episode_length = 0
+        batched_rewards = 0
+        batched_steps = 0
+        batched_episode = 0
+
+        actor_net = self.worker_network.actor
+        critic_net = self.worker_network.critic
+
+        actor_net_old = config.actor_network_fn()
+        actor_net_old.load_state_dict(actor_net.state_dict())
+
+        self.worker_network.load_state_dict(self.shared_network.state_dict())
+
+        while not replay.full():
+            states = []
+            actions = []
+            rewards = []
+            values = []
+            returns = []
+            advantages = []
+
+            for i in range(config.rollout_length):
+                mean, std, log_std = actor_net.predict(np.stack([state]))
+                value = critic_net.predict(np.stack([state]))
+                action = self.policy.sample(mean.data.numpy().flatten(), std.data.numpy().flatten(), deterministic)
+                action = self.config.action_shift_fn(action)
+                states.append(state)
+                actions.append(action)
+                values.append(value)
+                state, reward, done, _ = self.task.step(action)
+                state = self.state_normalizer(state)
+                done = (done or (config.max_episode_length and episode_length > config.max_episode_length))
+
+                batched_rewards += reward
+                batched_steps += 1
+                episode_length += 1
+
+
+
+                reward = np.asscalar(self.reward_normalizer(np.array([reward])))
+                rewards.append(reward)
+
+                if done:
+                    episode_length = 0
+                    batched_episode += 1
+                    state = self.task.reset()
+                    state = self.state_normalizer(state)
+                    break
+
+            R = torch.zeros((1, 1))
+            if not done:
+                R = critic_net.predict(np.stack([state])).data
+
+            values.append(Variable(R))
+            A = Variable(torch.zeros((1, 1)))
+            for i in reversed(range(len(rewards))):
+                R = Variable(torch.FloatTensor([[rewards[i]]]))
+                ret = R + self.config.discount * values[i + 1]
+                A = ret - values[i] + self.config.discount * self.config.gae_tau * A
+                advantages.append(A.detach())
+                returns.append(ret.detach())
+            advantages = list(reversed(advantages))
+            returns = list(reversed(returns))
+            replay.feed([states, actions, returns, advantages])
+
+        batched_rewards /= batched_episode
+        batched_steps /= batched_episode
+
+        if deterministic:
+            return batched_steps, batched_rewards
+
+        with config.steps_lock:
+            config.total_steps.value += replay.memory_size
+
+        self.shared_state_normalizer.offline_stats.merge(self.state_normalizer.online_stats)
+        self.state_normalizer.online_stats.zero()
+
+        self.shared_reward_normalizer.offline_stats.merge(self.reward_normalizer.online_stats)
+        self.reward_normalizer.online_stats.zero()
+
+        for _ in np.arange(self.config.optimize_epochs):
+            # self.worker_network.load_state_dict(self.shared_network.state_dict())
+
+            states, actions, returns, advantages = replay.sample()
+            states = actor_net.to_torch_variable(np.stack(states))
+            actions = actor_net.to_torch_variable(np.stack(actions))
+            returns = torch.cat(returns, 0)
+            advantages = torch.cat(advantages, 0)
+            advantages = (advantages - advantages.mean().expand_as(advantages)) / advantages.std().expand_as(advantages)
+
+            mean_old, std_old, log_std_old = actor_net_old.predict(states)
+            probs_old = actor_net.log_density(actions, mean_old, log_std_old, std_old)
+            mean, std, log_std = actor_net.predict(states)
+            probs = actor_net.log_density(actions, mean, log_std, std)
+            ratio = (probs - probs_old).exp()
+            obj = ratio * advantages
+            obj_clipped = ratio.clamp(1.0 - self.config.ppo_ratio_clip, 1.0 + self.config.ppo_ratio_clip) * advantages
+            policy_loss = -torch.min(obj, obj_clipped).mean(0)
+            if config.entropy_weight:
+                policy_loss += -config.entropy_weight * actor_net.entropy(std)
+
+            v = critic_net.predict(states)
+            value_loss = 0.5 * (returns - v).pow(2).mean()
+            actor_net_old.load_state_dict(actor_net.state_dict())
+
+            self.worker_network.zero_grad()
+            self.actor_opt.zero_grad()
+            self.critic_opt.zero_grad()
+            policy_loss.backward()
+            value_loss.backward()
+            nn.utils.clip_grad_norm(self.worker_network.parameters(), config.gradient_clip)
+            for param, worker_param in zip(
+                    self.shared_network.parameters(), self.worker_network.parameters()):
+                if param.grad is not None:
+                    break
+                param._grad = worker_param.grad
+            self.actor_opt.step()
+            self.critic_opt.step()
+
+        return batched_steps, batched_rewards

component/task.py

@@ -87,16 +87,10 @@ class Pendulum(BasicTask):
         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.observation_space.low) / \
-        #         (self.env.observation_space.high - self.env.observation_space.low)
-        # state = state * 2 - 1
-        return state
-
     def step(self, action):
         action = np.clip(action, -2, 2)
         next_state, reward, done, info = self.env.step(action)
-        return self.normalize_state(next_state), reward, done, info
+        return next_state, reward, done, info
 
 class BipedalWalker(BasicTask):
     name = 'BipedalWalker-v2'

main.py

@@ -303,36 +303,37 @@ def hrmsdqn_fruit():
 
 def ppo_pendulum():
     config = Config()
-    config.task_fn = lambda: Pendulum()
-    # config.task_fn = lambda: BipedalWalker()
+    # config.task_fn = lambda: Pendulum()
+    config.task_fn = lambda: BipedalWalker()
     # config.reward_shift_fn = lambda reward: reward / 10
     task = config.task_fn()
     config.actor_network_fn = lambda: GaussianActorNet(task.state_dim, task.action_dim)
     config.critic_network_fn = lambda: GaussianCriticNet(task.state_dim)
-    config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
+    config.network_fn = lambda: DisjointActorCriticNet(config.actor_network_fn, config.critic_network_fn)
+    config.actor_optimizer_fn = lambda params: torch.optim.Adam(params, 0.0001)
     config.critic_optimizer_fn = lambda params: torch.optim.Adam(params, 0.001)
 
     config.policy_fn = lambda: GaussianPolicy()
     config.replay_fn = lambda: GeneralReplay(memory_size=2048, batch_size=64)
-    config.worker = ContinuousAdvantageActorCritic
+    config.worker = ProximalPolicyOptimization
     config.discount = 0.99
     config.gae_tau = 0.97
     config.max_episode_length = 200
-    config.num_workers = None
-    config.test_interval = None
-    config.test_repetitions = None
+    config.num_workers = 8
+    config.test_interval = 1
+    config.test_repetitions = 1
     config.entropy_weight = 0
     config.gradient_clip = 40
     config.rollout_length = 10000
     config.optimize_epochs = 10
     config.ppo_ratio_clip = 0.2
     config.logger = Logger('./log', gym.logger)
-    agent = PPOAgent(config)
+    agent = AsyncAgent(config)
     agent.run()
 
 if __name__ == '__main__':
-    # gym.logger.setLevel(logging.DEBUG)
-    gym.logger.setLevel(logging.INFO)
+    gym.logger.setLevel(logging.DEBUG)
+    # gym.logger.setLevel(logging.INFO)
 
     # dqn_cart_pole()
     # async_cart_pole()
@@ -341,7 +342,7 @@ if __name__ == '__main__':
     # a3c_walker()
     # ddpg_pendulum()
     # ddpg_walker()
-    # ppo_pendulum()
+    ppo_pendulum()
 
     # dqn_fruit()
     # hrdqn_fruit()

utils/normalizer.py

@@ -1,3 +1,8 @@
+#######################################################################
+# 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
 
 class StaticNormalizer:

utils/shared_grad.py

@@ -0,0 +1,17 @@
+# adapted from https://github.com/alexis-jacq/Pytorch-DPPO/blob/master/model.py
+
+import torch
+
+class SharedGrad():
+    def __init__(self, model):
+        self.grads = {}
+        for name, p in model.named_parameters():
+            self.grads[name+'_grad'] = torch.ones(p.size()).share_memory_()
+
+    def add_gradient(self, model):
+        for name, p in model.named_parameters():
+            self.grads[name+'_grad'] += p.grad.data
+
+    def reset(self):
+        for name,grad in self.grads.items():
+            self.grads[name].fill_(0)