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

Reformat

Browse source
This commit is contained in:
Antonin Raffin 2019-12-02 14:27:38 +01:00
parent 3cdd5f20af
commit 1ac1a7cad5
6 changed files with 53 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
torchy_baselines/a2c/a2c.py
View file

@ -4,7 +4,6 @@ import torch.nn.functional as F
from torchy_baselines.common.utils import explained_variance
from torchy_baselines.ppo.ppo import PPO
from torchy_baselines.ppo.policies import PPOPolicy
from torchy_baselines.common import logger
@ -132,5 +131,5 @@ class A2C(PPO):
eval_env=None, eval_freq=-1, n_eval_episodes=5, tb_log_name="A2C", reset_num_timesteps=True):
return super(A2C, self).learn(total_timesteps=total_timesteps, callback=callback, log_interval=log_interval,
eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)
eval_env=eval_env, eval_freq=eval_freq, n_eval_episodes=n_eval_episodes,
tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps)

41
torchy_baselines/common/distributions.py
View file

@ -1,10 +1,9 @@
import numpy as np
import torch as th
import torch.nn as nn
from torch.distributions import Normal, Categorical
import torch.nn.functional as F
from gym import spaces
class Distribution(object):
def __init__(self):
super(Distribution, self).__init__()
@ -13,25 +12,25 @@ class Distribution(object):
"""
returns the log likelihood
:param x: (str) the labels of each index
:return: ([float]) The log likelihood of the distribution
:param x: (object) the taken action
:return: (th.Tensor) The log likelihood of the distribution
"""
raise NotImplementedError
def kl_div(self, other):
"""
Calculates the Kullback-Leibler divergence from the given probabilty distribution
:param other: ([float]) the distribution to compare with
:return: (float) the KL divergence of the two distributions
"""
raise NotImplementedError
# def kl_div(self, other):
# """
# Calculates the Kullback-Leibler divergence from the given probabilty distribution
#
# :param other: ([float]) the distribution to compare with
# :return: (float) the KL divergence of the two distributions
# """
# raise NotImplementedError
def entropy(self):
"""
Returns shannon's entropy of the probability
:return: (float) the entropy
:return: (th.Tensor) the entropy
"""
raise NotImplementedError
@ -39,7 +38,7 @@ class Distribution(object):
"""
returns a sample from the probabilty distribution
:return: (Tensorflow Tensor) the stochastic action
:return: (th.Tensor) the stochastic action
"""
raise NotImplementedError
@ -51,6 +50,7 @@ class DiagGaussianDistribution(Distribution):
:param action_dim: (int) Number of continuous actions
"""
def __init__(self, action_dim):
super(DiagGaussianDistribution, self).__init__()
self.distribution = None
@ -137,6 +137,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
:param action_dim: (int) Number of continuous actions
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, action_dim, epsilon=1e-6):
super(SquashedDiagGaussianDistribution, self).__init__(action_dim)
# Avoid NaN (prevents division by zero or log of zero)
@ -144,7 +145,8 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
self.gaussian_action = None
def proba_distribution(self, mean_actions, log_std, deterministic=False):
action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std, deterministic)
action, _ = super(SquashedDiagGaussianDistribution, self).proba_distribution(mean_actions, log_std,
deterministic)
return action, self
def mode(self):
@ -183,6 +185,7 @@ class CategoricalDistribution(Distribution):
:param action_dim: (int) Number of discrete actions
"""
def __init__(self, action_dim):
super(CategoricalDistribution, self).__init__()
self.distribution = None
@ -246,6 +249,7 @@ class StateDependentNoiseDistribution(Distribution):
`latent_sde` in the code.
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, action_dim, full_std=True, use_expln=False,
squash_output=False, learn_features=False, epsilon=1e-6):
super(StateDependentNoiseDistribution, self).__init__()
@ -256,12 +260,12 @@ class StateDependentNoiseDistribution(Distribution):
self.log_std = None
self.weights_dist = None
self.exploration_mat = None
self.exploration_matrices = None
self.use_expln = use_expln
self.full_std = full_std
self.epsilon = epsilon
self.learn_features = learn_features
if squash_output:
print("== Using TanhBijector ===")
self.bijector = TanhBijector(epsilon)
else:
self.bijector = None
@ -296,6 +300,7 @@ class StateDependentNoiseDistribution(Distribution):
using a centered gaussian distribution.
:param log_std: (th.Tensor)
:param batch_size: (int)
"""
std = self.get_std(log_std)
self.weights_dist = Normal(th.zeros_like(std), std)
@ -333,6 +338,7 @@ class StateDependentNoiseDistribution(Distribution):
:param mean_actions: (th.Tensor)
:param log_std: (th.Tensor)
:param latent_sde: (th.Tensor)
:param deterministic: (bool)
:return: (th.Tensor)
"""
@ -408,6 +414,7 @@ class TanhBijector(object):
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, epsilon=1e-6):
super(TanhBijector, self).__init__()
self.epsilon = epsilon
@ -439,7 +446,7 @@ class TanhBijector(object):
def log_prob_correction(self, x):
# Squash correction (from original SAC implementation)
return th.log(1 - th.tanh(x) ** 2 + self.epsilon)
return th.log(1.0 - th.tanh(x) ** 2 + self.epsilon)
def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):

18
torchy_baselines/ppo/policies.py
View file

@ -4,7 +4,8 @@ import torch as th
import torch.nn as nn
import numpy as np
from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, create_mlp, create_sde_feature_extractor
from torchy_baselines.common.policies import BasePolicy, register_policy, MlpExtractor, \
create_sde_feature_extractor
from torchy_baselines.common.distributions import make_proba_distribution,\
DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution
@ -14,7 +15,7 @@ class PPOPolicy(BasePolicy):
Policy class (with both actor and critic) for A2C and derivates (PPO).
:param observation_space: (gym.spaces.Space) Observation space
:param action_dim: (gym.spaces.Space) Action space
:param action_space: (gym.spaces.Space) Action space
:param learning_rate: (callable) Learning rate schedule (could be constant)
:param net_arch: ([int or dict]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.
@ -87,21 +88,24 @@ class PPOPolicy(BasePolicy):
self.mlp_extractor = MlpExtractor(self.features_dim, net_arch=self.net_arch,
activation_fn=self.activation_fn, device=self.device)
latent_dim_pi = self.mlp_extractor.latent_dim_pi
# Separate feature extractor for SDE
if self.sde_net_arch is not None:
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim, self.sde_net_arch,
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(self.features_dim,
self.sde_net_arch,
self.activation_fn)
if isinstance(self.action_dist, DiagGaussianDistribution):
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
log_std_init=self.log_std_init)
elif isinstance(self.action_dist, StateDependentNoiseDistribution):
latent_sde_dim = self.mlp_extractor.latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi,
latent_sde_dim = latent_dim_pi if self.sde_net_arch is None else latent_sde_dim
self.action_net, self.log_std = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi,
latent_sde_dim=latent_sde_dim,
log_std_init=self.log_std_init)
elif isinstance(self.action_dist, CategoricalDistribution):
self.action_net = self.action_dist.proba_distribution_net(latent_dim=self.mlp_extractor.latent_dim_pi)
self.action_net = self.action_dist.proba_distribution_net(latent_dim=latent_dim_pi)
self.value_net = nn.Linear(self.mlp_extractor.latent_dim_vf, 1)
# Init weights: use orthogonal initialization

3
torchy_baselines/ppo/ppo.py
View file

@ -43,7 +43,8 @@ class PPO(BaseRLModel):
:param n_epochs: (int) Number of epoch when optimizing the surrogate loss
:param gamma: (float) Discount factor
:param gae_lambda: (float) Factor for trade-off of bias vs variance for Generalized Advantage Estimator
:param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress (from 1 to 0).
:param clip_range: (float or callable) Clipping parameter, it can be a function of the current progress
(from 1 to 0).
:param clip_range_vf: (float or callable) Clipping parameter for the value function,
it can be a function of the current progress (from 1 to 0).
This is a parameter specific to the OpenAI implementation. If None is passed (default),

14
torchy_baselines/sac/policies.py
View file

@ -1,7 +1,8 @@
import torch as th
import torch.nn as nn
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
create_sde_feature_extractor
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
# CAP the standard deviation of the actor
@ -38,7 +39,8 @@ class Actor(BaseNetwork):
latent_sde_dim = net_arch[-1]
# Separate feature extractor for SDE
if sde_net_arch is not None:
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
activation_fn)
# TODO: check for the learn_features
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
@ -95,10 +97,12 @@ class Actor(BaseNetwork):
mean_actions, log_std, latent_sde = self.get_action_dist_params(obs)
if self.use_sde:
# Note the action is squashed
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde, deterministic=deterministic)
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, latent_sde,
deterministic=deterministic)
else:
# Note the action is squashed
action, _ = self.action_dist.proba_distribution(mean_actions, log_std, deterministic=deterministic)
action, _ = self.action_dist.proba_distribution(mean_actions, log_std,
deterministic=deterministic)
return action
def action_log_prob(self, obs):
@ -145,7 +149,7 @@ class SACPolicy(BasePolicy):
Policy class (with both actor and critic) for SAC.
:param observation_space: (gym.spaces.Space) Observation space
:param action_dim: (gym.spaces.Space) Action space
:param action_space: (gym.spaces.Space) Action space
:param learning_rate: (callable) Learning rate schedule (could be constant)
:param net_arch: ([int or dict]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.

9
torchy_baselines/td3/policies.py
View file

@ -1,7 +1,8 @@
import torch as th
import torch.nn as nn
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, create_sde_feature_extractor
from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
create_sde_feature_extractor
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
@ -43,7 +44,8 @@ class Actor(BaseNetwork):
# Separate feature extractor for SDE
if sde_net_arch is not None:
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch, activation_fn)
self.sde_feature_extractor, latent_sde_dim = create_sde_feature_extractor(obs_dim, sde_net_arch,
activation_fn)
# Create state dependent noise matrix (SDE)
self.action_dist = StateDependentNoiseDistribution(action_dim, full_std=full_std, use_expln=False,
@ -92,7 +94,6 @@ class Actor(BaseNetwork):
:param obs: (th.Tensor)
:param action: (th.Tensor)
:param deterministic: (bool)
:return: (th.Tensor, th.Tensor) log likelihood of taking those actions
and entropy of the action distribution.
"""
@ -162,7 +163,7 @@ class TD3Policy(BasePolicy):
Policy class (with both actor and critic) for TD3.
:param observation_space: (gym.spaces.Space) Observation space
:param action_dim: (gym.spaces.Space) Action space
:param action_space: (gym.spaces.Space) Action space
:param learning_rate: (callable) Learning rate schedule (could be constant)
:param net_arch: ([int or dict]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.