saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-11 17:48:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Renaming

Browse source
This commit is contained in:
Antonin Raffin 2019-09-24 14:53:03 +02:00
parent d22caac616
commit f4fe1362f0
7 changed files with 88 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
torchy_baselines/common/buffers.py
View file

@ -3,10 +3,10 @@ import torch as th
class BaseBuffer(object):
def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1):
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(BaseBuffer, self).__init__()
self.buffer_size = buffer_size
self.state_dim = state_dim
self.obs_dim = obs_dim
self.action_dim = action_dim
self.pos = 0
self.full = False
@ -58,20 +58,20 @@ class ReplayBuffer(BaseBuffer):
Taken from https://github.com/apourchot/CEM-RL
"""
def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1):
super(ReplayBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs)
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu', n_envs=1):
super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
assert n_envs == 1
self.states = th.zeros(self.buffer_size, self.n_envs, self.state_dim)
self.observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim)
self.actions = th.zeros(self.buffer_size, self.n_envs, self.action_dim)
self.next_states = th.zeros(self.buffer_size, self.n_envs, self.state_dim)
self.next_observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim)
self.rewards = th.zeros(self.buffer_size, self.n_envs)
self.dones = th.zeros(self.buffer_size, self.n_envs)
def add(self, state, next_state, action, reward, done):
def add(self, obs, next_obs, action, reward, done):
# Copy to avoid modification by reference
self.states[self.pos] = th.FloatTensor(np.array(state).copy())
self.next_states[self.pos] = th.FloatTensor(np.array(next_state).copy())
self.observations[self.pos] = th.FloatTensor(np.array(obs).copy())
self.next_observations[self.pos] = th.FloatTensor(np.array(next_obs).copy())
self.actions[self.pos] = th.FloatTensor(np.array(action).copy())
self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy())
self.dones[self.pos] = th.FloatTensor(np.array(done).copy())
@ -82,27 +82,27 @@ class ReplayBuffer(BaseBuffer):
self.pos = 0
def _get_samples(self, batch_inds):
return (self.states[batch_inds, 0, :].to(self.device),
return (self.observations[batch_inds, 0, :].to(self.device),
self.actions[batch_inds, 0, :].to(self.device),
self.next_states[batch_inds, 0, :].to(self.device),
self.next_observations[batch_inds, 0, :].to(self.device),
self.dones[batch_inds].to(self.device),
self.rewards[batch_inds].to(self.device))
class RolloutBuffer(BaseBuffer):
def __init__(self, buffer_size, state_dim, action_dim, device='cpu',
lambda_=1, gamma=0.99, n_envs=1):
super(RolloutBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs)
def __init__(self, buffer_size, obs_dim, action_dim, device='cpu',
gae_lambda=1, gamma=0.99, n_envs=1):
super(RolloutBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
# TODO: try the buffer on the gpu?
self.lambda_ = lambda_
self.gae_lambda = gae_lambda
self.gamma = gamma
self.states, self.actions, self.rewards, self.advantages = None, None, None, None
self.observations, self.actions, self.rewards, self.advantages = None, None, None, None
self.returns, self.dones, self.values, self.log_probs = None, None, None, None
self.generator_ready = False
self.reset()
def reset(self):
self.states = th.zeros(self.buffer_size, self.n_envs, self.state_dim)
self.observations = th.zeros(self.buffer_size, self.n_envs, self.obs_dim)
self.actions = th.zeros(self.buffer_size, self.n_envs, self.action_dim)
self.rewards = th.zeros(self.buffer_size, self.n_envs)
self.returns = th.zeros(self.buffer_size, self.n_envs)
@ -126,12 +126,12 @@ class RolloutBuffer(BaseBuffer):
next_non_terminal = 1.0 - self.dones[step + 1]
next_value = self.values[step + 1]
delta = self.rewards[step] + self.gamma * next_value * next_non_terminal - self.values[step]
last_gae_lam = delta + self.gamma * self.lambda_ * next_non_terminal * last_gae_lam
last_gae_lam = delta + self.gamma * self.gae_lambda * next_non_terminal * last_gae_lam
self.advantages[step] = last_gae_lam
self.returns = self.advantages + self.values
def add(self, state, action, reward, done, value, log_prob):
self.states[self.pos] = th.FloatTensor(np.array(state).copy())
def add(self, obs, action, reward, done, value, log_prob):
self.observations[self.pos] = th.FloatTensor(np.array(obs).copy())
self.actions[self.pos] = th.FloatTensor(np.array(action).copy())
self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy())
self.dones[self.pos] = th.FloatTensor(np.array(done).copy())
@ -146,7 +146,7 @@ class RolloutBuffer(BaseBuffer):
indices = th.randperm(self.buffer_size * self.n_envs)
# Prepare the data
if not self.generator_ready:
for tensor in ['states', 'actions', 'values',
for tensor in ['observations', 'actions', 'values',
'log_probs', 'advantages', 'returns']:
self.__dict__[tensor] = self.swap_and_flatten(self.__dict__[tensor])
self.generator_ready = True
@ -157,7 +157,7 @@ class RolloutBuffer(BaseBuffer):
start_idx += batch_size
def _get_samples(self, batch_inds):
return (self.states[batch_inds].to(self.device),
return (self.observations[batch_inds].to(self.device),
self.actions[batch_inds].to(self.device),
self.values[batch_inds].flatten().to(self.device),
self.log_probs[batch_inds].flatten().to(self.device),

33
torchy_baselines/ppo/policies.py
View file

@ -2,18 +2,18 @@ from functools import partial
import torch as th
import torch.nn as nn
from torch.distributions import Normal
import numpy as np
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp
from torchy_baselines.common.distributions import DiagGaussianDistribution, SquashedDiagGaussianDistribution
class PPOPolicy(BasePolicy):
def __init__(self, observation_space, action_space,
learning_rate=1e-3, net_arch=None, device='cpu',
activation_fn=nn.Tanh, adam_epsilon=1e-5):
super(PPOPolicy, self).__init__(observation_space, action_space, device)
self.state_dim = self.observation_space.shape[0]
self.obs_dim = self.observation_space.shape[0]
self.action_dim = self.action_space.shape[0]
if net_arch is None:
net_arch = [64, 64]
@ -21,7 +21,7 @@ class PPOPolicy(BasePolicy):
self.activation_fn = activation_fn
self.adam_epsilon = adam_epsilon
self.net_args = {
'input_dim': self.state_dim,
'input_dim': self.obs_dim,
'output_dim': -1,
'net_arch': self.net_arch,
'activation_fn': self.activation_fn
@ -41,12 +41,12 @@ class PPOPolicy(BasePolicy):
def _build(self, learning_rate):
# TODO: support shared network
# shared_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
# shared_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
# self.shared_net = nn.Sequential(*shared_net).to(self.device)
pi_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
pi_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
self.pi_net = nn.Sequential(*pi_net).to(self.device)
vf_net = create_mlp(self.state_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
vf_net = create_mlp(self.obs_dim, output_dim=-1, net_arch=self.net_arch, activation_fn=self.activation_fn)
self.vf_net = nn.Sequential(*vf_net).to(self.device)
# self.action_net = nn.Linear(self.net_arch[-1], self.action_dim)
@ -67,32 +67,33 @@ class PPOPolicy(BasePolicy):
# TODO: support linear decay of the learning rate
self.optimizer = th.optim.Adam(self.parameters(), lr=learning_rate, eps=self.adam_epsilon)
def forward(self, state, deterministic=False):
state = th.FloatTensor(state).to(self.device)
latent_pi, latent_vf = self._get_latent(state)
def forward(self, obs, deterministic=False):
if not isinstance(obs, th.Tensor):
obs = th.FloatTensor(obs).to(self.device)
latent_pi, latent_vf = self._get_latent(obs)
value = self.value_net(latent_vf)
action, action_distribution = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
log_prob = action_distribution.log_prob(action)
return action, value, log_prob
def _get_latent(self, state):
def _get_latent(self, obs):
if self.shared_net is not None:
latent = self.shared_net(state)
latent = self.shared_net(obs)
return latent, latent
else:
return self.pi_net(state), self.vf_net(state)
return self.pi_net(obs), self.vf_net(obs)
def _get_action_dist_from_latent(self, latent, deterministic=False):
mean_actions = self.action_net(latent)
return self.action_dist.proba_distribution(mean_actions, self.log_std, deterministic=deterministic)
def actor_forward(self, state, deterministic=False):
latent_pi, _ = self._get_latent(state)
def actor_forward(self, obs, deterministic=False):
latent_pi, _ = self._get_latent(obs)
action, _ = self._get_action_dist_from_latent(latent_pi, deterministic=deterministic)
return action.detach().cpu().numpy()
def get_policy_stats(self, state, action):
latent_pi, latent_vf = self._get_latent(state)
def get_policy_stats(self, obs, action):
latent_pi, latent_vf = self._get_latent(obs)
_, action_distribution = self._get_action_dist_from_latent(latent_pi)
log_prob = action_distribution.log_prob(action)
value = self.value_net(latent_vf)

12
torchy_baselines/ppo/ppo.py
View file

@ -31,7 +31,7 @@ class PPO(BaseRLModel):
def __init__(self, policy, env, policy_kwargs=None, verbose=0,
learning_rate=3e-4, seed=0, device='auto',
n_optim=5, batch_size=64, n_steps=256,
gamma=0.99, lambda_=0.95, clip_range=0.2,
gamma=0.99, gae_lambda=0.95, clip_range=0.2,
ent_coef=0.01, vf_coef=0.5, max_grad_norm=0.5,
target_kl=None, clip_range_vf=None, create_eval_env=False,
tensorboard_log=None,
@ -46,7 +46,7 @@ class PPO(BaseRLModel):
self.n_optim = n_optim
self.n_steps = n_steps
self.gamma = gamma
self.lambda_ = lambda_
self.gae_lambda = gae_lambda
self.clip_range = clip_range
self.ent_coef = ent_coef
self.vf_coef = vf_coef
@ -67,7 +67,7 @@ class PPO(BaseRLModel):
self.seed(self._seed)
self.rollout_buffer = RolloutBuffer(self.n_steps, state_dim, action_dim, self.device,
gamma=self.gamma, lambda_=self.lambda_, n_envs=self.n_envs)
gamma=self.gamma, gae_lambda=self.gae_lambda, n_envs=self.n_envs)
self.policy = self.policy(self.observation_space, self.action_space,
self.learning_rate, device=self.device, **self.policy_kwargs)
self.policy = self.policy.to(self.device)
@ -125,8 +125,8 @@ class PPO(BaseRLModel):
# Sample replay buffer
for replay_data in self.rollout_buffer.get(batch_size):
# Unpack
state, action, old_values, old_log_prob, advantage, return_batch = replay_data
values, log_prob, entropy = self.policy.get_policy_stats(state, action)
obs, action, old_values, old_log_prob, advantage, return_batch = replay_data
values, log_prob, entropy = self.policy.get_policy_stats(obs, action)
values = values.flatten()
# Normalize advantage
advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
@ -145,7 +145,7 @@ class PPO(BaseRLModel):
# Clip the different between old and new value
# NOTE: this depends on the reward scaling
values_pred = old_values + th.clamp(values - old_values, -self.clip_range_vf, self.clip_range_vf)
# Value loss using the TD(lambda_) target
# Value loss using the TD(gae_lambda) target
value_loss = F.mse_loss(return_batch, values_pred)

26
torchy_baselines/sac/policies.py
View file

@ -10,51 +10,51 @@ LOG_STD_MIN = -20
class Actor(BaseNetwork):
def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU):
def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU):
super(Actor, self).__init__()
if net_arch is None:
net_arch = [256, 256]
# TODO: orthogonal initialization?
actor_net = create_mlp(state_dim, -1, net_arch, activation_fn)
actor_net = create_mlp(obs_dim, -1, net_arch, activation_fn)
self.actor_net = nn.Sequential(*actor_net)
self.action_dist = SquashedDiagGaussianDistribution(action_dim)
self.mu = nn.Linear(net_arch[-1], action_dim)
self.log_std = nn.Linear(net_arch[-1], action_dim)
def get_action_dist_params(self, state):
latent = self.actor_net(state)
def get_action_dist_params(self, obs):
latent = self.actor_net(obs)
mean_actions, log_std = self.mu(latent), self.log_std(latent)
# Original Implementation to cap the standard deviation
log_std = th.clamp(log_std, LOG_STD_MIN, LOG_STD_MAX)
return mean_actions, log_std
def forward(self, state, deterministic=False):
mean_actions, log_std = self.get_action_dist_params(state)
def forward(self, obs, deterministic=False):
mean_actions, log_std = self.get_action_dist_params(obs)
# Note the action is squashed
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
return action
def action_log_prob(self, state):
mean_actions, log_std = self.get_action_dist_params(state)
def action_log_prob(self, obs):
mean_actions, log_std = self.get_action_dist_params(obs)
action, log_prob = self.action_dist.log_prob_from_params(mean_actions, log_std)
return action, log_prob
class Critic(BaseNetwork):
def __init__(self, state_dim, action_dim,
def __init__(self, obs_dim, action_dim,
net_arch=None, activation_fn=nn.ReLU):
super(Critic, self).__init__()
if net_arch is None:
net_arch = [256, 256]
q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn)
q1_net = create_mlp(obs_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)
q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
self.q2_net = nn.Sequential(*q2_net)
self.q_networks = [self.q1_net, self.q2_net]
@ -72,12 +72,12 @@ class SACPolicy(BasePolicy):
learning_rate=1e-3, net_arch=None, device='cpu',
activation_fn=nn.ReLU):
super(SACPolicy, self).__init__(observation_space, action_space, device)
self.state_dim = self.observation_space.shape[0]
self.obs_dim = self.observation_space.shape[0]
self.action_dim = self.action_space.shape[0]
self.net_arch = net_arch
self.activation_fn = activation_fn
self.net_args = {
'state_dim': self.state_dim,
'obs_dim': self.obs_dim,
'action_dim': self.action_dim,
'net_arch': self.net_arch,
'activation_fn': self.activation_fn

16
torchy_baselines/sac/sac.py
View file

@ -52,7 +52,7 @@ class SAC(BaseRLModel):
self._setup_model()
def _setup_model(self):
state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
self.seed(self._seed)
# Target entropy is used when learning the entropy coefficient
@ -88,7 +88,7 @@ class SAC(BaseRLModel):
# is passed
self.ent_coef = float(self.ent_coef)
self.replay_buffer = ReplayBuffer(self.buffer_size, state_dim, action_dim, self.device)
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
self.policy = self.policy(self.observation_space, self.action_space,
self.learning_rate, device=self.device, **self.policy_kwargs)
self.policy = self.policy.to(self.device)
@ -125,10 +125,10 @@ class SAC(BaseRLModel):
# Sample replay buffer
replay_data = self.replay_buffer.sample(batch_size)
state, action_batch, next_state, done, reward = replay_data
obs, action_batch, next_obs, done, reward = replay_data
# Action by the current actor for the sampled state
action_pi, log_prob = self.actor.action_log_prob(state)
action_pi, log_prob = self.actor.action_log_prob(obs)
log_prob = log_prob.reshape(-1, 1)
ent_coef_loss = None
@ -143,10 +143,10 @@ class SAC(BaseRLModel):
self.ent_coef_optimizer.step()
# Select action according to policy
next_action, next_log_prob = self.actor.action_log_prob(next_state)
next_action, next_log_prob = self.actor.action_log_prob(next_obs)
# Compute the target Q value
target_q1, target_q2 = self.critic_target(next_state, next_action)
target_q1, target_q2 = self.critic_target(next_obs, next_action)
target_q = th.min(target_q1, target_q2)
target_q = reward + ((1 - done) * self.gamma * target_q).detach()
@ -155,7 +155,7 @@ class SAC(BaseRLModel):
# Get current Q estimates
# using action from the replay buffer
current_q1, current_q2 = self.critic(state, action_batch)
current_q1, current_q2 = self.critic(obs, action_batch)
# Compute critic loss
critic_loss = 0.5 * (F.mse_loss(current_q1, q_backup) + F.mse_loss(current_q2, q_backup))
@ -167,7 +167,7 @@ class SAC(BaseRLModel):
# Compute actor loss
# Alternative: actor_loss = th.mean(log_prob - min_qf_pi)
actor_loss = (self.ent_coef * log_prob - self.critic.q1_forward(state, action_pi)).mean()
actor_loss = (self.ent_coef * log_prob - self.critic.q1_forward(obs, action_pi)).mean()
# Optimize the actor
self.actor.optimizer.zero_grad()

21
torchy_baselines/td3/policies.py
View file

@ -5,32 +5,32 @@ from torchy_baselines.common.policies import BasePolicy, register_policy, create
class Actor(BaseNetwork):
def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU):
def __init__(self, obs_dim, action_dim, net_arch=None, activation_fn=nn.ReLU):
super(Actor, self).__init__()
if net_arch is None:
net_arch = [400, 300]
# TODO: orthogonal initialization?
actor_net = create_mlp(state_dim, action_dim, net_arch, activation_fn, squash_out=True)
actor_net = create_mlp(obs_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)
def forward(self, obs):
return self.actor_net(obs)
class Critic(BaseNetwork):
def __init__(self, state_dim, action_dim,
def __init__(self, obs_dim, action_dim,
net_arch=None, activation_fn=nn.ReLU):
super(Critic, self).__init__()
if net_arch is None:
net_arch = [400, 300]
q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn)
q1_net = create_mlp(obs_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)
q2_net = create_mlp(obs_dim + action_dim, 1, net_arch, activation_fn)
self.q2_net = nn.Sequential(*q2_net)
self.q_networks = [self.q1_net, self.q2_net]
@ -48,12 +48,12 @@ class TD3Policy(BasePolicy):
learning_rate=1e-3, net_arch=None, device='cpu',
activation_fn=nn.ReLU):
super(TD3Policy, self).__init__(observation_space, action_space, device)
self.state_dim = self.observation_space.shape[0]
self.obs_dim = self.observation_space.shape[0]
self.action_dim = self.action_space.shape[0]
self.net_arch = net_arch
self.activation_fn = activation_fn
self.net_args = {
'state_dim': self.state_dim,
'obs_dim': self.obs_dim,
'action_dim': self.action_dim,
'net_arch': self.net_arch,
'activation_fn': self.activation_fn
@ -79,6 +79,9 @@ class TD3Policy(BasePolicy):
def make_critic(self):
return Critic(**self.net_args).to(self.device)
def forward(self, obs):
return self.actor(obs)
MlpPolicy = TD3Policy

21
torchy_baselines/td3/td3.py
View file

@ -16,7 +16,6 @@ class TD3(BaseRLModel):
Paper: https://arxiv.org/abs/1802.09477
Code: https://github.com/sfujim/TD3
"""
def __init__(self, policy, env, policy_kwargs=None, verbose=0,
buffer_size=int(1e6), learning_rate=1e-3, seed=0, device='auto',
action_noise_std=0.1, start_timesteps=100, policy_freq=2,
@ -39,9 +38,9 @@ class TD3(BaseRLModel):
self._setup_model()
def _setup_model(self):
state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
obs_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
self.seed(self._seed)
self.replay_buffer = ReplayBuffer(self.buffer_size, state_dim, action_dim, self.device)
self.replay_buffer = ReplayBuffer(self.buffer_size, obs_dim, action_dim, self.device)
self.policy = self.policy(self.observation_space, self.action_space,
self.learning_rate, device=self.device, **self.policy_kwargs)
self.policy = self.policy.to(self.device)
@ -78,22 +77,22 @@ class TD3(BaseRLModel):
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)
obs, action, next_obs, done, reward = self.replay_buffer.sample(batch_size)
else:
state, action, next_state, done, reward = replay_data
obs, action, next_obs, done, reward = replay_data
# Select action according to policy and add clipped noise
noise = action.clone().data.normal_(0, policy_noise)
noise = noise.clamp(-noise_clip, noise_clip)
next_action = (self.actor_target(next_state) + noise).clamp(-1, 1)
next_action = (self.actor_target(next_obs) + noise).clamp(-1, 1)
# Compute the target Q value
target_q1, target_q2 = self.critic_target(next_state, next_action)
target_q1, target_q2 = self.critic_target(next_obs, next_action)
target_q = th.min(target_q1, target_q2)
target_q = reward + ((1 - done) * discount * target_q).detach()
# Get current Q estimates
current_q1, current_q2 = self.critic(state, action)
current_q1, current_q2 = self.critic(obs, action)
# Compute critic loss
critic_loss = F.mse_loss(current_q1, target_q) + F.mse_loss(current_q2, target_q)
@ -115,12 +114,12 @@ class TD3(BaseRLModel):
for it in range(n_iterations):
# Sample replay buffer
if replay_data is None:
state, _, next_state, done, reward = self.replay_buffer.sample(batch_size)
obs, _, next_obs, done, reward = self.replay_buffer.sample(batch_size)
else:
state, _, next_state, done, reward = replay_data
obs, _, next_obs, done, reward = replay_data
# Compute actor loss
actor_loss = -self.critic.q1_forward(state, self.actor(state)).mean()
actor_loss = -self.critic.q1_forward(obs, self.actor(obs)).mean()
# Optimize the actor
self.actor.optimizer.zero_grad()