From 6bb7e183d22a0b7407556c55fdb7cb51c32f4c34 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Wed, 18 Sep 2019 15:35:17 +0200 Subject: [PATCH] Running PPO (not working yet) --- tests/test_run.py | 28 ++++ tests/test_td3.py | 21 --- torchy_baselines/__init__.py | 3 +- torchy_baselines/cem_rl/cem_rl.py | 2 +- torchy_baselines/common/policies.py | 1 + torchy_baselines/common/replay_buffer.py | 62 +++++++++ torchy_baselines/common/utils.py | 34 +++++ torchy_baselines/ppo/__init__.py | 1 + torchy_baselines/ppo/policies.py | 56 +------- torchy_baselines/ppo/ppo.py | 164 +++++++++++------------ 10 files changed, 215 insertions(+), 157 deletions(-) create mode 100644 tests/test_run.py delete mode 100644 tests/test_td3.py diff --git a/tests/test_run.py b/tests/test_run.py new file mode 100644 index 0000000..342ced9 --- /dev/null +++ b/tests/test_run.py @@ -0,0 +1,28 @@ +import os + +import gym + +from torchy_baselines import TD3, CEMRL, PPO + +# def test_pendulum(): +# model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1) +# model.learn(total_timesteps=500, eval_freq=100) +# model.save("test_save") +# model.load("test_save") +# os.remove("test_save.pth") +# +# +# def test_cemrl(): +# model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1, +# start_timesteps=100, verbose=1) +# model.learn(total_timesteps=1000, eval_freq=500) +# model.save("test_save") +# model.load("test_save") +# os.remove("test_save.pth") + +def test_ppo(): + model = PPO('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), verbose=1) + model.learn(total_timesteps=1000, eval_freq=500) + # model.save("test_save") + # model.load("test_save") + # os.remove("test_save.pth") diff --git a/tests/test_td3.py b/tests/test_td3.py deleted file mode 100644 index 089e79f..0000000 --- a/tests/test_td3.py +++ /dev/null @@ -1,21 +0,0 @@ -import os - -import gym - -from torchy_baselines import TD3, CEMRL - -def test_pendulum(): - model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1) - model.learn(total_timesteps=500, eval_freq=100) - model.save("test_save") - model.load("test_save") - os.remove("test_save.pth") - - -def test_cemrl(): - model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1, - start_timesteps=100, verbose=1) - model.learn(total_timesteps=1000, eval_freq=500) - model.save("test_save") - model.load("test_save") - os.remove("test_save.pth") diff --git a/torchy_baselines/__init__.py b/torchy_baselines/__init__.py index fad6e63..e95418f 100644 --- a/torchy_baselines/__init__.py +++ b/torchy_baselines/__init__.py @@ -1,4 +1,5 @@ -from torchy_baselines.td3 import TD3 from torchy_baselines.cem_rl import CEMRL +from torchy_baselines.ppo import PPO +from torchy_baselines.td3 import TD3 __version__ = "0.0.2" diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index f3d4336..0715880 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -4,7 +4,7 @@ import torch as th import torch.nn.functional as F import numpy as np -from torchy_baselines import TD3 +from torchy_baselines.td3.td3 import TD3 from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.cem_rl.cem import CEM diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 33f143b..c1ab00f 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -1,3 +1,4 @@ +import torch as th import torch.nn as nn diff --git a/torchy_baselines/common/replay_buffer.py b/torchy_baselines/common/replay_buffer.py index b5a953c..ea23c84 100644 --- a/torchy_baselines/common/replay_buffer.py +++ b/torchy_baselines/common/replay_buffer.py @@ -1,6 +1,8 @@ import numpy as np import torch as th +from torchy_baselines.common.utils import discount_cumsum + class ReplayBuffer(object): """ @@ -43,14 +45,74 @@ class ReplayBuffer(object): self.full = True self.pos = 0 + def reset(self): + self.pos = 0 + self.full = False + def sample(self, batch_size): upper_bound = self.buffer_size if self.full else self.pos batch_inds = th.LongTensor( np.random.randint(0, upper_bound, size=batch_size)) + return self._get_samples(batch_inds) + def _get_samples(self, batch_inds): return (self.states[batch_inds].to(self.device), self.actions[batch_inds].to(self.device), self.next_states[batch_inds].to(self.device), self.dones[batch_inds].to(self.device), self.rewards[batch_inds].to(self.device)) + + +class RolloutBuffer(ReplayBuffer): + def __init__(self, buffer_size, state_dim, action_dim, device='cpu', + lambda_=1, gamma=0.99): + super(RolloutBuffer, self).__init__(buffer_size, state_dim, action_dim, device) + + self.lambda_ = lambda_ + self.gamma = gamma + # TODO: add n_envs + self.returns = th.zeros(self.buffer_size, 1) + self.values = th.zeros(self.buffer_size, 1) + self.log_probs = th.zeros(self.buffer_size, 1) + self.advantages = th.zeros(self.buffer_size, 1) + self.path_start_idx = 0 + + def finish_path(self, last_value=0): + """ + From https://github.com/openai/spinningup/blob/master/spinup/algos/ppo/ppo.py + """ + if self.full: + self.pos = self.buffer_size + path_slice = slice(self.path_start_idx, self.pos) + rewards = np.append(self.rewards[path_slice].detach().cpu().numpy(), last_value) + values = np.append(self.values[path_slice].detach().cpu().numpy(), last_value) + + # the next two lines implement GAE-Lambda advantage calculation + deltas = rewards[:-1] + self.gamma * values[1:] - values[:-1] + + self.advantages[path_slice, 0] = th.FloatTensor(discount_cumsum(deltas, self.gamma * self.lambda_).copy()) + # the next line computes rewards-to-go, to be targets for the value function + self.returns[path_slice, 0] = th.FloatTensor(discount_cumsum(rewards, self.gamma)[:-1].copy()) + + self.path_start_idx = self.pos + + def add(self, state, next_state, action, reward, done, value, log_prob): + self.values[self.pos] = th.FloatTensor([value]) + self.log_probs[self.pos] = th.FloatTensor([log_prob]) + super(RolloutBuffer, self).add(state, next_state, action, reward, done) + + def reset(self): + self.path_start_idx = 0 + super(RolloutBuffer, self).reset() + + def _get_samples(self, batch_inds): + return (self.states[batch_inds].to(self.device), + self.actions[batch_inds].to(self.device), + self.next_states[batch_inds].to(self.device), + self.dones[batch_inds].to(self.device), + self.rewards[batch_inds].to(self.device), + self.values[batch_inds].to(self.device), + self.log_probs[batch_inds].to(self.device), + self.advantages[batch_inds].to(self.device), + self.returns[batch_inds].to(self.device)) diff --git a/torchy_baselines/common/utils.py b/torchy_baselines/common/utils.py index a1030cb..e526c03 100644 --- a/torchy_baselines/common/utils.py +++ b/torchy_baselines/common/utils.py @@ -1,5 +1,6 @@ import random +import scipy.signal import torch as th import numpy as np @@ -18,3 +19,36 @@ def set_random_seed(seed, using_cuda=False): # Make CuDNN Determinist th.backends.cudnn.deterministic = True th.cuda.manual_seed(seed) + + +# From stable_baselines.common.math_util +# def discount(vector, gamma): +# """ +# computes discounted sums along 0th dimension of vector x. +# y[t] = x[t] + gamma*x[t+1] + gamma^2*x[t+2] + ... + gamma^k x[t+k], +# where k = len(x) - t - 1 +# +# :param vector: (np.ndarray) the input vector +# :param gamma: (float) the discount value +# :return: (np.ndarray) the output vector +# """ +# assert vector.ndim >= 1 +# return scipy.signal.lfilter([1], [1, -gamma], vector[::-1], axis=0)[::-1] + + +def discount_cumsum(x, discount): + """ + magic from rllab for computing discounted cumulative sums of vectors. + + input: + vector x, + [x0, + x1, + x2] + + output: + [x0 + discount * x1 + discount^2 * x2, + x1 + discount * x2, + x2] + """ + return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1] diff --git a/torchy_baselines/ppo/__init__.py b/torchy_baselines/ppo/__init__.py index e69de29..2ce3051 100644 --- a/torchy_baselines/ppo/__init__.py +++ b/torchy_baselines/ppo/__init__.py @@ -0,0 +1 @@ +from torchy_baselines.ppo.ppo import PPO diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 53548c4..1aea078 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -2,52 +2,7 @@ import torch as th import torch.nn as nn from torch.distributions import Normal -from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork - - -class Actor(BaseNetwork): - def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU): - super(Actor, self).__init__() - - if net_arch is None: - net_arch = [64, 64] - - # TODO: orthogonal initialization? - actor_net = create_mlp(state_dim, action_dim, net_arch, activation_fn, squash_out=True) - self.actor_net = nn.Sequential(*actor_net) - - def forward(self, x): - return self.actor_net(x) - - -class Critic(BaseNetwork): - def __init__(self, state_dim, action_dim, - net_arch=None, activation_fn=nn.ReLU): - super(Critic, self).__init__() - - if net_arch is None: - net_arch = [400, 300] - - # TODO: solve pytorch parameter registration - # for _ in range(n_critics): - # q_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) - # self.q_net = nn.Sequential(*q_net) - # self.q_networks.append(self.q_net) - - q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) - self.q1_net = nn.Sequential(*q1_net) - - q2_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn) - self.q2_net = nn.Sequential(*q2_net) - - self.q_networks = [self.q1_net, self.q2_net] - - def forward(self, obs, action): - qvalue_input = th.cat([obs, action], dim=1) - return [q_net(qvalue_input) for q_net in self.q_networks] - - def q1_forward(self, obs, action): - return self.q_networks[0](th.cat([obs, action], dim=1)) +from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp class PPOPolicy(BasePolicy): @@ -71,7 +26,7 @@ class PPOPolicy(BasePolicy): self._build(learning_rate) def _build(self, learning_rate): - shared_net = create_mlp(self.state_dim, output_dim=-1, self.net_arch, self.activation_fn) + shared_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn) self.shared_net = nn.Sequential(*shared_net).to(self.device) self.actor_net = nn.Linear(self.net_arch[-1], self.action_dim) self.value_net = nn.Linear(self.net_arch[-1], 1) @@ -79,18 +34,19 @@ class PPOPolicy(BasePolicy): self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate) def forward(self, state): + state = th.FloatTensor(state).to(self.device) latent = self.shared_net(state) # TODO: initialize pi_mean weights properly mean_actions = self.actor_net(latent) action_distribution = Normal(mean_actions, self.log_std) # Sample from the gaussian action = action_distribution.rsample() - log_prob = action_distribution.log_prob() + log_prob = action_distribution.log_prob(action) # entropy = action_distribution.entropy() value = self.value_net(latent) return action, value, log_prob - def actor_forward(self): + def actor_forward(self, state): latent = self.shared_net(state) # TODO: initialize pi_mean weights properly mean_actions = self.actor_net(latent) @@ -98,7 +54,7 @@ class PPOPolicy(BasePolicy): # Sample from the gaussian action = action_distribution.rsample() return action - + def value_forward(self): pass diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index db12a03..45c8ab2 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -5,10 +5,9 @@ import torch.nn.functional as F import numpy as np from torchy_baselines.common.base_class import BaseRLModel -from torchy_baselines.common.utils import set_random_seed from torchy_baselines.common.evaluation import evaluate_policy -from torchy_baselines.ppo.policies import ActorCriticPolicy -from torchy_baselines.common.replay_buffer import ReplayBuffer +from torchy_baselines.ppo.policies import PPOPolicy +from torchy_baselines.common.replay_buffer import RolloutBuffer class PPO(BaseRLModel): @@ -16,15 +15,18 @@ class PPO(BaseRLModel): Implementation of Proximal Policy Optimization (PPO) (clip version) Paper: https://arxiv.org/abs/1707.06347 Code: https://github.com/openai/spinningup/ + and https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail + and stable_baselines """ def __init__(self, policy, env, policy_kwargs=None, verbose=0, learning_rate=1e-3, seed=0, device='auto', - n_optim=5, batch_size=100, n_steps=256, - gamma=0.99, lambda_=0.95, - _init_setup_model=True): + n_optim=5, batch_size=64, n_steps=256, + gamma=0.99, lambda_=0.95, clip_range=0.2, + ent_coef=0.01, vf_coef=0.5, + _init_setup_model=True): - super(PPO, self).__init__(policy, env, ActorCriticPolicy, policy_kwargs, verbose, device) + super(PPO, self).__init__(policy, env, PPOPolicy, policy_kwargs, verbose, device) self.max_action = np.abs(self.action_space.high) self.learning_rate = learning_rate @@ -34,7 +36,10 @@ class PPO(BaseRLModel): self.n_steps = n_steps self.gamma = gamma self.lambda_ = lambda_ - self.buffer_rollouts = None + self.clip_range = clip_range + self.ent_coef = ent_coef + self.vf_coef = vf_coef + self.rollout_buffer = None if _init_setup_model: self._setup_model() @@ -43,10 +48,11 @@ class PPO(BaseRLModel): state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0] self.seed(self._seed) + self.rollout_buffer = RolloutBuffer(self.n_steps, state_dim, action_dim, self.device, + gamma=self.gamma, lambda_=self.lambda_) self.policy = self.policy(self.observation_space, self.action_space, self.learning_rate, device=self.device, **self.policy_kwargs) - def select_action(self, observation): # Normally not needed observation = np.array(observation) @@ -66,52 +72,80 @@ class PPO(BaseRLModel): """ return np.clip(self.select_action(observation), -self.max_action, self.max_action) + def collect_rollouts(self, env, rollout_buffer, n_rollout_steps=256, callback=None, + obs=None): - def train_actor(self, n_iterations=1, batch_size=100, tau_actor=0.005, tau_critic=0.005, replay_data=None): + n_steps = 0 + done = obs is None + rollout_buffer.reset() + while n_steps < n_rollout_steps: + # Reset environment + if done: + obs = env.reset() + + # No grad ok? + with th.no_grad(): + action, value, log_prob = self.policy.forward(obs) + action = action[0].detach().cpu().numpy() + + # Rescale and perform action + new_obs, reward, done, _ = env.step(np.clip(action, -self.max_action, self.max_action)) + + n_steps += 1 + rollout_buffer.add(obs, new_obs, action, reward, float(done), value, log_prob) + + obs = new_obs + + if done: + value = 0.0 + obs = None + + rollout_buffer.finish_path(last_value=value) + + return obs + + def train(self, n_iterations, batch_size=64): + + # TODO: replace with iterator? for it in range(n_iterations): # Sample replay buffer - if replay_data is None: - state, action, next_state, done, reward = self.replay_buffer.sample(batch_size) - else: - state, action, next_state, done, reward = replay_data + replay_data = self.rollout_buffer.sample(batch_size) + state, action, next_state, done, reward, _, old_log_prob, advantage, return_batch = replay_data - # Compute actor loss - actor_loss = -self.critic.q1_forward(state, self.actor(state)).mean() + _, value, log_prob = self.policy.forward(state) - # Optimize the actor - self.actor.optimizer.zero_grad() - actor_loss.backward() - self.actor.optimizer.step() + # Normalize advantage + # advs = returns - values + advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) - # Update the frozen target models - if tau_critic > 0: - for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()): - target_param.data.copy_(tau_critic * param.data + (1 - tau_critic) * target_param.data) - - for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): - target_param.data.copy_(tau_actor * param.data + (1 - tau_actor) * target_param.data) - - def train(self, n_iterations, batch_size=100, discount=0.99, - tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2): - - for it in range(n_iterations): - - # Sample replay buffer - replay_data = self.replay_buffer.sample(batch_size) - self.train_critic(replay_data=replay_data) - - # Delayed policy updates - if it % policy_freq == 0: - self.train_actor(replay_data=replay_data) + ratio = th.exp(log_prob - old_log_prob) + policy_loss_1 = -advantage * ratio + policy_loss_2 = -advantage * th.clamp(ratio, 1 - self.clip_range, 1 + self.clip_range) + policy_loss = -th.min(policy_loss_1, policy_loss_2).mean() + # value_loss = th.mean((returns - value)**2) + value_loss = F.mse_loss(return_batch, value) + # Approximate entropy + # TODO: replace by distribution entropy + entropy_loss = th.mean(-log_prob) + loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss + # TODO: check kl div + # approx_kl_div = th.mean(old_log_prob - log_prob) + # Optimization step + self.policy.optimizer.zero_grad() + loss.backward() + # TODO: clip grad norm? + # nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) + self.policy.optimizer.step() def learn(self, total_timesteps, callback=None, log_interval=100, - eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): + eval_freq=-1, n_eval_episodes=5, tb_log_name="PPO", reset_num_timesteps=True): timesteps_since_eval = 0 episode_num = 0 evaluations = [] start_time = time.time() + obs = None while self.num_timesteps < total_timesteps: @@ -120,21 +154,13 @@ class PPO(BaseRLModel): if callback(locals(), globals()) is False: break - episode_reward, episode_timesteps = self.collect_rollouts(self.env, n_episodes=1, - action_noise_std=self.action_noise_std, - deterministic=False, callback=None, - start_timesteps=self.start_timesteps, - num_timesteps=self.num_timesteps, - replay_buffer=self.buffer_rollouts) + obs = self.collect_rollouts(self.env, self.rollout_buffer, n_rollout_steps=self.n_steps, + obs=obs) episode_num += 1 - self.num_timesteps += episode_timesteps - timesteps_since_eval += episode_timesteps + self.num_timesteps += self.n_steps + timesteps_since_eval += self.n_steps - if self.num_timesteps > 0: - if self.verbose > 1: - print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - self.num_timesteps, episode_num, episode_timesteps, episode_reward)) - self.train(episode_timesteps, batch_size=self.batch_size, policy_freq=self.policy_freq) + self.train(self.n_optim, batch_size=self.batch_size) # Evaluate episode if 0 < eval_freq <= timesteps_since_eval: @@ -158,33 +184,3 @@ class PPO(BaseRLModel): if env is not None: pass self.policy.load_state_dict(th.load(path)) - - -class PPOBuffer(ReplayBuffer): - """docstring for PPOBuffer.""" - - def __init__(self, buffer_size, state_dim, action_dim, device='cpu', - lambda=0.95): - super(PPOBuffer, self).__init__(buffer_size, state_dim, action_dim, device) - - self.returns = th.zeros(self.buffer_size, 1) - self.values = th.zeros(self.buffer_size, 1) - self.log_probs = th.zeros(self.buffer_size, 1) - self.advantages = th.zeros(self.buffer_size, 1) - - def compute_gae(self): - """ - From https://github.com/openai/spinningup/blob/master/spinup/algos/ppo/ppo.py - """ - path_slice = slice(self.path_start_idx, self.pos) - rews = np.append(self.rewards[path_slice], last_val) - vals = np.append(self.val_buf[path_slice], last_val) - - # the next two lines implement GAE-Lambda advantage calculation - deltas = rews[:-1] + self.gamma * vals[1:] - vals[:-1] - self.advantages[path_slice] = core.discount_cumsum(deltas, self.gamma * self.lam) - - # the next line computes rewards-to-go, to be targets for the value function - self.ret_buf[path_slice] = core.discount_cumsum(rews, self.gamma)[:-1] - - self.path_start_idx = self.pos