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

Reformat and type the distributions

Browse source
This commit is contained in:
Antonin Raffin 2020-02-13 13:46:22 +01:00
parent f1a4fa2d3f
commit aa8b4eb22a
4 changed files with 75 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
torchy_baselines/common/base_class.py
View file

@ -225,13 +225,14 @@ class BaseRLModel(ABC):
return self._vec_normalize_env
@staticmethod
def check_env(env, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> bool:
def check_env(env: GymEnv, observation_space: gym.spaces.Space, action_space: gym.spaces.Space) -> bool:
"""
Checks the validity of the environment and returns if it is consistent.
Checked parameters:
- observation_space
- action_space
:param env: (GymEnv)
:param observation_space: (gym.spaces.Space)
:param action_space: (gym.spaces.Space)
:return: (bool) True if environment seems to be coherent
@ -404,7 +405,6 @@ class BaseRLModel(ABC):
# return clipped_actions, state
return clipped_actions
@classmethod
def load(cls, load_path: str, env: Optional[GymEnv] = None, **kwargs):
"""
@ -774,6 +774,7 @@ class OffPolicyRLModel(BaseRLModel):
:param use_sde_at_warmup: (bool) Whether to use SDE instead of uniform sampling
during the warm up phase (before learning starts)
"""
def __init__(self,
policy: Type[BasePolicy],
env: Union[GymEnv, str],
@ -790,8 +791,8 @@ class OffPolicyRLModel(BaseRLModel):
use_sde_at_warmup: bool = False):
super(OffPolicyRLModel, self).__init__(policy, env, policy_base, policy_kwargs, verbose,
device, support_multi_env, create_eval_env, monitor_wrapper,
seed, use_sde, sde_sample_freq)
device, support_multi_env, create_eval_env, monitor_wrapper,
seed, use_sde, sde_sample_freq)
# For SDE only
self.rollout_data = None
self.on_policy_exploration = False

114
torchy_baselines/common/distributions.py
View file

@ -1,5 +1,6 @@
from typing import Optional
from typing import Optional, Tuple, Dict, Any
import gym
import torch as th
import torch.nn as nn
from torch.distributions import Normal, Categorical
@ -45,14 +46,15 @@ class DiagGaussianDistribution(Distribution):
:param action_dim: (int) Number of continuous actions
"""
def __init__(self, action_dim):
def __init__(self, action_dim: int):
super(DiagGaussianDistribution, self).__init__()
self.distribution = None
self.action_dim = action_dim
self.mean_actions = None
self.log_std = None
def proba_distribution_net(self, latent_dim, log_std_init=0.0):
def proba_distribution_net(self, latent_dim: int,
log_std_init: float = 0.0) -> Tuple[nn.Module, nn.Parameter]:
"""
Create the layers and parameter that represent the distribution:
one output will be the mean of the gaussian, the other parameter will be the
@ -64,10 +66,12 @@ class DiagGaussianDistribution(Distribution):
"""
mean_actions = nn.Linear(latent_dim, self.action_dim)
# TODO: allow action dependent std
log_std = nn.Parameter(th.ones(self.action_dim) * log_std_init)
log_std = nn.Parameter(th.ones(self.action_dim) * log_std_init, requires_grad=True)
return mean_actions, log_std
def proba_distribution(self, mean_actions, log_std, deterministic=False):
def proba_distribution(self, mean_actions: th.Tensor,
log_std: th.Tensor,
deterministic: bool = False) -> Tuple[th.Tensor, 'DiagGaussianDistribution']:
"""
Create and sample for the distribution given its parameters (mean, std)
@ -84,29 +88,29 @@ class DiagGaussianDistribution(Distribution):
action = self.sample()
return action, self
def mode(self):
def mode(self) -> th.Tensor:
return self.distribution.mean
def sample(self):
def sample(self) -> th.Tensor:
return self.distribution.rsample()
def entropy(self):
def entropy(self) -> th.Tensor:
return self.distribution.entropy()
def log_prob_from_params(self, mean_actions, log_std):
def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
"""
Compute the log probabilty of taking an action
given the distribution parameters.
:param mean_actions: (th.Tensor)
:param log_std: (th.Tensor)
:return: (th.Tensor, th.Tensor)
:return: (Tuple[th.Tensor, th.Tensor])
"""
action, _ = self.proba_distribution(mean_actions, log_std)
log_prob = self.log_prob(action)
return action, log_prob
def log_prob(self, action):
def log_prob(self, action: th.Tensor) -> th.Tensor:
"""
Get the log probabilty of an action given a distribution.
Note that you must call `proba_distribution()` method
@ -132,7 +136,7 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, action_dim, epsilon=1e-6):
def __init__(self, action_dim: int, epsilon: float = 1e-6):
super(SquashedDiagGaussianDistribution, self).__init__(action_dim)
# Avoid NaN (prevents division by zero or log of zero)
self.epsilon = epsilon
@ -143,26 +147,26 @@ class SquashedDiagGaussianDistribution(DiagGaussianDistribution):
deterministic)
return action, self
def mode(self):
def mode(self) -> th.Tensor:
self.gaussian_action = self.distribution.mean
# Squash the output
return th.tanh(self.gaussian_action)
def entropy(self):
def entropy(self) -> Optional[th.Tensor]:
# No analytical form,
# entropy needs to be estimated using -log_prob.mean()
return None
def sample(self):
def sample(self) -> th.Tensor:
self.gaussian_action = self.distribution.rsample()
return th.tanh(self.gaussian_action)
def log_prob_from_params(self, mean_actions, log_std):
def log_prob_from_params(self, mean_actions, log_std) -> Tuple[th.Tensor, th.Tensor]:
action, _ = self.proba_distribution(mean_actions, log_std)
log_prob = self.log_prob(action, self.gaussian_action)
return action, log_prob
def log_prob(self, action, gaussian_action=None):
def log_prob(self, action: th.Tensor, gaussian_action: Optional[th.Tensor] = None) -> th.Tensor:
# Inverse tanh
# Naive implementation (not stable): 0.5 * torch.log((1 + x) / (1 - x))
# We use numpy to avoid numerical instability
@ -185,12 +189,12 @@ class CategoricalDistribution(Distribution):
:param action_dim: (int) Number of discrete actions
"""
def __init__(self, action_dim):
def __init__(self, action_dim: int):
super(CategoricalDistribution, self).__init__()
self.distribution = None
self.action_dim = action_dim
def proba_distribution_net(self, latent_dim):
def proba_distribution_net(self, latent_dim: int) -> nn.Module:
"""
Create the layer that represents the distribution:
it will be the logits of the Categorical distribution.
@ -202,7 +206,8 @@ class CategoricalDistribution(Distribution):
action_logits = nn.Linear(latent_dim, self.action_dim)
return action_logits
def proba_distribution(self, action_logits, deterministic=False):
def proba_distribution(self, action_logits: th.Tensor,
deterministic: bool = False) -> Tuple[th.Tensor, 'CategoricalDistribution']:
self.distribution = Categorical(logits=action_logits)
if deterministic:
action = self.mode()
@ -210,21 +215,21 @@ class CategoricalDistribution(Distribution):
action = self.sample()
return action, self
def mode(self):
def mode(self) -> th.Tensor:
return th.argmax(self.distribution.probs, dim=1)
def sample(self):
def sample(self) -> th.Tensor:
return self.distribution.sample()
def entropy(self):
def entropy(self) -> th.Tensor:
return self.distribution.entropy()
def log_prob_from_params(self, action_logits):
def log_prob_from_params(self, action_logits: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
action, _ = self.proba_distribution(action_logits)
log_prob = self.log_prob(action)
return action, log_prob
def log_prob(self, action):
def log_prob(self, action: th.Tensor) -> th.Tensor:
log_prob = self.distribution.log_prob(action)
return log_prob
@ -249,8 +254,12 @@ class StateDependentNoiseDistribution(Distribution):
: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):
def __init__(self, action_dim: int,
full_std: bool = True,
use_expln: bool = False,
squash_output: bool = False,
learn_features: bool = False,
epsilon: float = 1e-6):
super(StateDependentNoiseDistribution, self).__init__()
self.distribution = None
self.action_dim = action_dim
@ -269,7 +278,7 @@ class StateDependentNoiseDistribution(Distribution):
else:
self.bijector = None
def get_std(self, log_std):
def get_std(self, log_std: th.Tensor) -> th.Tensor:
"""
Get the standard deviation from the learned parameter
(log of it by default). This ensures that the std is positive.
@ -294,7 +303,7 @@ class StateDependentNoiseDistribution(Distribution):
# Reduce the number of parameters:
return th.ones(self.latent_sde_dim, self.action_dim).to(log_std.device) * std
def sample_weights(self, log_std, batch_size=1):
def sample_weights(self, log_std: th.Tensor, batch_size: int = 1) -> None:
"""
Sample weights for the noise exploration matrix,
using a centered Gaussian distribution.
@ -307,7 +316,8 @@ class StateDependentNoiseDistribution(Distribution):
self.exploration_mat = self.weights_dist.rsample()
self.exploration_matrices = self.weights_dist.rsample((batch_size,))
def proba_distribution_net(self, latent_dim, log_std_init=-2.0, latent_sde_dim=None):
def proba_distribution_net(self, latent_dim: int, log_std_init: float = -2.0,
latent_sde_dim: Optional[th.Tensor] = None) -> Tuple[nn.Module, nn.Parameter]:
"""
Create the layers and parameter that represent the distribution:
one output will be the deterministic action, the other parameter will be the
@ -327,12 +337,15 @@ class StateDependentNoiseDistribution(Distribution):
# Reduce the number of parameters if needed
log_std = th.ones(self.latent_sde_dim, self.action_dim) if self.full_std else th.ones(self.latent_sde_dim, 1)
# Transform it to a parameter so it can be optimized
log_std = nn.Parameter(log_std * log_std_init)
log_std = nn.Parameter(log_std * log_std_init, requires_grad=True)
# Sample an exploration matrix
self.sample_weights(log_std)
return mean_actions_net, log_std
def proba_distribution(self, mean_actions, log_std, latent_sde, deterministic=False):
def proba_distribution(self, mean_actions: th.Tensor,
log_std: th.Tensor,
latent_sde: th.Tensor,
deterministic: bool = False) -> Tuple[th.Tensor, 'StateDependentNoiseDistribution']:
"""
Create and sample for the distribution given its parameters (mean, std)
@ -340,7 +353,7 @@ class StateDependentNoiseDistribution(Distribution):
:param log_std: (th.Tensor)
:param latent_sde: (th.Tensor)
:param deterministic: (bool)
:return: (th.Tensor)
:return: (Tuple[th.Tensor, Distribution])
"""
# Stop gradient if we don't want to influence the features
latent_sde = latent_sde if self.learn_features else latent_sde.detach()
@ -353,13 +366,13 @@ class StateDependentNoiseDistribution(Distribution):
action = self.sample(latent_sde)
return action, self
def mode(self):
def mode(self) -> th.Tensor:
action = self.distribution.mean
if self.bijector is not None:
return self.bijector.forward(action)
return action
def get_noise(self, latent_sde):
def get_noise(self, latent_sde: th.Tensor) -> th.Tensor:
latent_sde = latent_sde if self.learn_features else latent_sde.detach()
# Default case: only one exploration matrix
if len(latent_sde) == 1 or len(latent_sde) != len(self.exploration_matrices):
@ -371,26 +384,28 @@ class StateDependentNoiseDistribution(Distribution):
noise = th.bmm(latent_sde, self.exploration_matrices)
return noise.squeeze(1)
def sample(self, latent_sde):
def sample(self, latent_sde: th.Tensor) -> th.Tensor:
noise = self.get_noise(latent_sde)
action = self.distribution.mean + noise
if self.bijector is not None:
return self.bijector.forward(action)
return action
def entropy(self):
def entropy(self) -> Optional[th.Tensor]:
# No analytical form,
# entropy needs to be estimated using -log_prob.mean()
if self.bijector is not None:
return None
return self.distribution.entropy()
def log_prob_from_params(self, mean_actions, log_std, latent_sde):
def log_prob_from_params(self, mean_actions: th.Tensor,
log_std: th.Tensor,
latent_sde: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
action, _ = self.proba_distribution(mean_actions, log_std, latent_sde)
log_prob = self.log_prob(action)
return action, log_prob
def log_prob(self, action):
def log_prob(self, action: th.Tensor) -> th.Tensor:
if self.bijector is not None:
gaussian_action = self.bijector.inverse(action)
else:
@ -418,16 +433,16 @@ class TanhBijector(object):
:param epsilon: (float) small value to avoid NaN due to numerical imprecision.
"""
def __init__(self, epsilon=1e-6):
def __init__(self, epsilon: float = 1e-6):
super(TanhBijector, self).__init__()
self.epsilon = epsilon
@staticmethod
def forward(x):
def forward(x: th.Tensor) -> th.Tensor:
return th.tanh(x)
@staticmethod
def atanh(x):
def atanh(x: th.Tensor) -> th.Tensor:
"""
Inverse of Tanh
@ -437,7 +452,7 @@ class TanhBijector(object):
return 0.5 * (x.log1p() - (-x).log1p())
@staticmethod
def inverse(y):
def inverse(y: th.Tensor) -> th.Tensor:
"""
Inverse tanh.
@ -448,19 +463,21 @@ class TanhBijector(object):
# Clip the action to avoid NaN
return TanhBijector.atanh(y.clamp(min=-1. + eps, max=1. - eps))
def log_prob_correction(self, x):
def log_prob_correction(self, x: th.Tensor) -> th.Tensor:
# Squash correction (from original SAC implementation)
return th.log(1.0 - th.tanh(x) ** 2 + self.epsilon)
def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
def make_proba_distribution(action_space: gym.spaces.Space,
use_sde: bool = False,
dist_kwargs: Optional[Dict[str, Any]] = None) -> Distribution:
"""
Return an instance of Distribution for the correct type of action space
:param action_space: (Gym Space) the input action space
:param action_space: (gym.spaces.Space) the input action space
:param use_sde: (bool) Force the use of StateDependentNoiseDistribution
instead of DiagGaussianDistribution
:param dist_kwargs: (dict) Keyword arguments to pass to the probabilty distribution
:param dist_kwargs: (Optional[Dict[str, Any]]) Keyword arguments to pass to the probabilty distribution
:return: (Distribution) the approriate Distribution object
"""
if dist_kwargs is None:
@ -478,5 +495,6 @@ def make_proba_distribution(action_space, use_sde=False, dist_kwargs=None):
# elif isinstance(action_space, spaces.MultiBinary):
# return BernoulliDistribution(action_space.n, **dist_kwargs)
else:
raise NotImplementedError(f"Error: probability distribution, not implemented for action space of type {type(action_space)}."
raise NotImplementedError("Error: probability distribution, not implemented for action space"
f"of type {type(action_space)}."
" Must be of type Gym Spaces: Box, Discrete, MultiDiscrete or MultiBinary.")

4
torchy_baselines/common/results_plotter.py
View file

@ -2,12 +2,12 @@ from typing import Tuple, Callable, List, Optional
import numpy as np
import pandas as pd
import matplotlib
# import matplotlib
# matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode
import matplotlib.pyplot as plt
from torchy_baselines.common.monitor import load_results
# matplotlib.use('TkAgg') # Can change to 'Agg' for non-interactive mode
plt.rcParams['svg.fonttype'] = 'none'
X_TIMESTEPS = 'timesteps'

2
torchy_baselines/ppo/policies.py
View file

@ -161,6 +161,8 @@ class PPOPolicy(BasePolicy):
elif isinstance(self.action_dist, StateDependentNoiseDistribution):
return self.action_dist.proba_distribution(mean_actions, self.log_std, latent_sde,
deterministic=deterministic)
else:
raise ValueError('Invalid action distribution')
def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
latent_pi, _, latent_sde = self._get_latent(observation)