diff --git a/torchy_baselines/common/replay_buffer.py b/torchy_baselines/common/buffers.py similarity index 73% rename from torchy_baselines/common/replay_buffer.py rename to torchy_baselines/common/buffers.py index ea23c84..7ee555a 100644 --- a/torchy_baselines/common/replay_buffer.py +++ b/torchy_baselines/common/buffers.py @@ -3,14 +3,11 @@ import torch as th from torchy_baselines.common.utils import discount_cumsum +class BaseBuffer(object): + """docstring for BaseBuffer.""" -class ReplayBuffer(object): - """ - Taken from https://github.com/apourchot/CEM-RL - """ def __init__(self, buffer_size, state_dim, action_dim, device='cpu'): - super(ReplayBuffer, self).__init__() - # params + super(BaseBuffer, self).__init__() self.buffer_size = buffer_size self.state_dim = state_dim self.action_dim = action_dim @@ -18,12 +15,6 @@ class ReplayBuffer(object): self.full = False self.device = device - self.states = th.zeros(self.buffer_size, self.state_dim) - self.actions = th.zeros(self.buffer_size, self.action_dim) - self.next_states = th.zeros(self.buffer_size, self.state_dim) - self.rewards = th.zeros(self.buffer_size, 1) - self.dones = th.zeros(self.buffer_size, 1) - def size(self): if self.full: return self.buffer_size @@ -32,6 +23,34 @@ class ReplayBuffer(object): def get_pos(self): return self.pos + def add(self, *args, **kwargs): + raise NotImplementedError() + + 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): + raise NotImplementedError() + +class ReplayBuffer(BaseBuffer): + """ + Taken from https://github.com/apourchot/CEM-RL + """ + def __init__(self, buffer_size, state_dim, action_dim, device='cpu'): + super(ReplayBuffer, self).__init__(buffer_size, state_dim, action_dim, device) + self.states = th.zeros(self.buffer_size, self.state_dim) + self.actions = th.zeros(self.buffer_size, self.action_dim) + self.next_states = th.zeros(self.buffer_size, self.state_dim) + self.rewards = th.zeros(self.buffer_size, 1) + self.dones = th.zeros(self.buffer_size, 1) + def add(self, state, next_state, action, reward, done): self.states[self.pos] = th.FloatTensor(state) @@ -45,17 +64,6 @@ 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), @@ -64,7 +72,7 @@ class ReplayBuffer(object): self.rewards[batch_inds].to(self.device)) -class RolloutBuffer(ReplayBuffer): +class RolloutBuffer(BaseBuffer): 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) @@ -72,7 +80,11 @@ class RolloutBuffer(ReplayBuffer): self.lambda_ = lambda_ self.gamma = gamma # TODO: add n_envs + self.states = th.zeros(self.buffer_size, self.state_dim) + self.actions = th.zeros(self.buffer_size, self.action_dim) + self.rewards = th.zeros(self.buffer_size, 1) self.returns = th.zeros(self.buffer_size, 1) + self.dones = 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) @@ -82,8 +94,7 @@ class RolloutBuffer(ReplayBuffer): """ From https://github.com/openai/spinningup/blob/master/spinup/algos/ppo/ppo.py """ - if self.full: - self.pos = self.buffer_size + # No use of dones? 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) @@ -97,22 +108,33 @@ class RolloutBuffer(ReplayBuffer): self.path_start_idx = self.pos - def add(self, state, next_state, action, reward, done, value, log_prob): + def add(self, 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) + self.states[self.pos] = th.FloatTensor(state) + self.actions[self.pos] = th.FloatTensor(action) + self.rewards[self.pos] = th.FloatTensor([reward]) + self.dones[self.pos] = th.FloatTensor([done]) + self.pos += 1 + if self.pos == self.buffer_size: + self.full = True def reset(self): self.path_start_idx = 0 super(RolloutBuffer, self).reset() + def get(self, batch_size): + assert self.full + indices = th.randperm(self.buffer_size) + minibatch_indices = [] + start_idx = 0 + while start_idx < self.buffer_size: + yield self._get_samples(indices[start_idx:start_idx + batch_size]) + start_idx += batch_size + 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/ppo/policies.py b/torchy_baselines/ppo/policies.py index 66a4f21..f8fd89e 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -25,39 +25,33 @@ class PPOPolicy(BasePolicy): self.shared_net = None self._build(learning_rate) + @staticmethod + def init_weights(module): + if type(module) == nn.Linear: + nn.init.orthogonal_(module.weight, gain=1) + module.bias.data.fill_(0.0) + def _build(self, learning_rate): 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) self.log_std = nn.Parameter(th.zeros(self.action_dim)) + # Init weights: + for module in [self.shared_net, self.actor_net, self.value_net]: + module.apply(self.init_weights) + self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate) - def forward(self, state): + def forward(self, state, deterministic=False): state = th.FloatTensor(state).to(self.device) latent = self.shared_net(state) - # TODO: initialize pi_mean weights properly - # TODO: change when multiple envs - mean_actions = self.actor_net(latent) - action_std = th.ones(mean_actions.size()) * self.log_std.exp() - action_distribution = Normal(mean_actions, action_std) - # Sample from the gaussian - # rsample: reparametrization trick - action = action_distribution.rsample() - # TODO: handle shape properly - # sum(axis=1) - log_prob = action_distribution.log_prob(action) - if len(log_prob.shape) > 1: - log_prob = log_prob.sum(axis=1) - else: - log_prob = log_prob.sum() - # entropy = action_distribution.entropy() value = self.value_net(latent) + action, action_distribution = self._get_action_dist_from_latent(latent, deterministic=deterministic) + log_prob = self._get_log_prob(action_distribution, action) return action, value, log_prob - def actor_forward(self, state, deterministic=False): - latent = self.shared_net(state) - # TODO: initialize pi_mean weights properly + def _get_action_dist_from_latent(self, latent, deterministic=False): mean_actions = self.actor_net(latent) action_std = th.ones(mean_actions.size()) * self.log_std.exp() action_distribution = Normal(mean_actions, action_std) @@ -66,25 +60,30 @@ class PPOPolicy(BasePolicy): action = mean_actions else: action = action_distribution.rsample() - return action + return action, action_distribution - def get_policy_stats(self, state, action): - state = th.FloatTensor(state).to(self.device) - latent = self.shared_net(state) - # TODO: initialize pi_mean weights properly - # TODO: change when multiple envs - mean_actions = self.actor_net(latent) - action_std = th.ones(mean_actions.size()) * self.log_std.exp() - action_distribution = Normal(mean_actions, action_std) + @staticmethod + def _get_log_prob(action_distribution, action): log_prob = action_distribution.log_prob(action) - entropy = action_distribution.entropy() if len(log_prob.shape) > 1: log_prob = log_prob.sum(axis=1) else: log_prob = log_prob.sum() - # entropy = action_distribution.entropy() + return log_prob + + def actor_forward(self, state, deterministic=False): + state = th.FloatTensor(state).to(self.device) + latent = self.shared_net(state) + action, _ = self._get_action_dist_from_latent(latent, deterministic=deterministic) + return action.detach().cpu().numpy() + + def get_policy_stats(self, state, action): + state = th.FloatTensor(state).to(self.device) + latent = self.shared_net(state) + _, action_distribution = self._get_action_dist_from_latent(latent) + log_prob = self._get_log_prob(action_distribution, action) value = self.value_net(latent) - return value, log_prob, entropy + return value, log_prob, action_distribution.entropy() def value_forward(self): pass diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index fabd4f1..5bf7d34 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -7,7 +7,7 @@ import numpy as np from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.ppo.policies import PPOPolicy -from torchy_baselines.common.replay_buffer import RolloutBuffer +from torchy_baselines.common.buffers import RolloutBuffer class PPO(BaseRLModel): @@ -58,7 +58,7 @@ class PPO(BaseRLModel): observation = np.array(observation) with th.no_grad(): observation = th.FloatTensor(observation.reshape(1, -1)).to(self.device) - return self.policy.actor_forward(observation, deterministic=False).cpu().data.numpy().flatten() + return self.policy.actor_forward(observation, deterministic=False).flatten() def predict(self, observation, state=None, mask=None, deterministic=True): """ @@ -90,12 +90,10 @@ class PPO(BaseRLModel): action = action[0].cpu().numpy() # Rescale and perform action - new_obs, reward, done, _ = env.step(np.clip(action, -self.max_action, self.max_action)) + 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 + rollout_buffer.add(obs, action, reward, float(done), value, log_prob) if done: value = 0.0 @@ -110,34 +108,35 @@ class PPO(BaseRLModel): # TODO: replace with iterator? for it in range(n_iterations): # Sample replay buffer - replay_data = self.rollout_buffer.sample(batch_size) - state, action, _, _, _, _, old_log_prob, advantage, return_batch = replay_data + for replay_data in self.rollout_buffer.get(batch_size): + # Unpack + state, action, old_log_prob, advantage, return_batch = replay_data - # _, values, log_prob = self.policy.forward(state) - values, log_prob, entropy = self.policy.get_policy_stats(state, action) + # _, values, log_prob = self.policy.forward(state) + values, log_prob, entropy = self.policy.get_policy_stats(state, action) - # Normalize advantage - # advs = returns - values - advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) + # Normalize advantage + # advs = returns - values + advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8) - 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((return_batch - value)**2) - value_loss = F.mse_loss(return_batch, values) - entropy_loss = th.mean(entropy) - loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss - # loss = policy_loss - # TODO: check kl div - # self.approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - self.old_neglog_pac_ph)) - # 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() + 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((return_batch - value)**2) + value_loss = F.mse_loss(return_batch, values) + entropy_loss = th.mean(entropy) + loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss + # loss = policy_loss + # TODO: check kl div + # self.approxkl = .5 * tf.reduce_mean(tf.square(neglogpac - self.old_neglog_pac_ph)) + # 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="PPO", reset_num_timesteps=True): diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 8ef6b8e..9e85c5f 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -5,7 +5,7 @@ import torch.nn.functional as F import numpy as np from torchy_baselines.common.base_class import BaseRLModel -from torchy_baselines.common.replay_buffer import ReplayBuffer +from torchy_baselines.common.buffers import ReplayBuffer from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.td3.policies import TD3Policy