saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-30 20:18:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Refactor buffer

Browse source
This commit is contained in:
Antonin RAFFIN 2019-09-19 11:43:15 +02:00
parent 149148d4c7
commit 26f0c8d8e5
4 changed files with 116 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

86
torchy_baselines/common/replay_buffer.py → torchy_baselines/common/buffers.py
View file

@ -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))

63
torchy_baselines/ppo/policies.py
View file

@ -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

61
torchy_baselines/ppo/ppo.py
View file

@ -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):

2
torchy_baselines/td3/td3.py
View file

@ -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