diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 4aa2dbe..33eae78 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/torchy_baselines/a2c/a2c.py @@ -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) diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 8802687..05a5588 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -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): diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 886d958..19e20a5 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -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 diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index 99ed1b4..2ffc6e6 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -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), diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index be32660..cd1b35b 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -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. diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 65f2171..a472b58 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -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.