Rebranding to Metastable Baselines

icon.svg

@@ -0,0 +1,97 @@
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<!-- Created with Inkscape (http://www.inkscape.org/) -->
+
+<svg
+   width="116.92769mm"
+   height="58.176697mm"
+   viewBox="0 0 116.92769 58.176697"
+   version="1.1"
+   id="svg5"
+   sodipodi:docname="metastable_col.svg"
+   inkscape:version="1.2 (dc2aedaf03, 2022-05-15)"
+   xmlns:inkscape="http://www.inkscape.org/namespaces/inkscape"
+   xmlns:sodipodi="http://sodipodi.sourceforge.net/DTD/sodipodi-0.dtd"
+   xmlns:xlink="http://www.w3.org/1999/xlink"
+   xmlns="http://www.w3.org/2000/svg"
+   xmlns:svg="http://www.w3.org/2000/svg">
+  <sodipodi:namedview
+     id="namedview7"
+     pagecolor="#505050"
+     bordercolor="#eeeeee"
+     borderopacity="1"
+     inkscape:showpageshadow="0"
+     inkscape:pageopacity="0"
+     inkscape:pagecheckerboard="0"
+     inkscape:deskcolor="#505050"
+     inkscape:document-units="mm"
+     showgrid="false"
+     inkscape:zoom="2.719879"
+     inkscape:cx="213.61244"
+     inkscape:cy="109.93136"
+     inkscape:window-width="1920"
+     inkscape:window-height="1050"
+     inkscape:window-x="0"
+     inkscape:window-y="0"
+     inkscape:window-maximized="1"
+     inkscape:current-layer="layer1" />
+  <defs
+     id="defs2">
+    <linearGradient
+       inkscape:collect="always"
+       id="linearGradient1075">
+      <stop
+         style="stop-color:#ff0000;stop-opacity:1;"
+         offset="0.40881506"
+         id="stop1071" />
+      <stop
+         style="stop-color:#00ff58;stop-opacity:1;"
+         offset="0.85932338"
+         id="stop1073" />
+    </linearGradient>
+    <inkscape:path-effect
+       effect="bspline"
+       id="path-effect299"
+       is_visible="true"
+       lpeversion="1"
+       weight="33.333333"
+       steps="2"
+       helper_size="0"
+       apply_no_weight="true"
+       apply_with_weight="true"
+       only_selected="false" />
+    <linearGradient
+       inkscape:collect="always"
+       xlink:href="#linearGradient1075"
+       id="linearGradient1077"
+       x1="103.01504"
+       y1="95.063644"
+       x2="104.56599"
+       y2="152.75583"
+       gradientUnits="userSpaceOnUse" />
+  </defs>
+  <g
+     inkscape:label="Layer 1"
+     inkscape:groupmode="layer"
+     id="layer1"
+     transform="translate(-42.437337,-95.11482)">
+    <path
+       style="fill:none;stroke:url(#linearGradient1077);stroke-width:0.865;stroke-linecap:butt;stroke-linejoin:miter;stroke-opacity:1;stroke-dasharray:none;fill-opacity:1"
+       d="m 42.550287,95.183688 c 10.933866,17.931672 21.867327,35.862692 31.926852,38.361672 10.059525,2.49898 19.244372,-10.43374 26.122241,-7.46713 6.87786,2.96661 11.44781,21.83192 20.8447,26.09518 9.39689,4.26327 23.6199,-6.07589 37.84315,-16.41521"
+       id="path297"
+       inkscape:path-effect="#path-effect299"
+       inkscape:original-d="m 42.550287,95.183688 c 10.933939,17.931632 21.867402,35.862642 32.801027,53.794092 9.185654,-12.93322 18.370498,-25.86595 27.555626,-38.79971 4.57049,18.86672 9.14044,37.73203 13.7104,56.59821 14.22385,-10.3393 28.44686,-20.67846 42.66989,-31.01808"
+       sodipodi:nodetypes="ccccc" />
+    <circle
+       style="fill:#000000;fill-opacity:0;stroke:#060301;stroke-width:1.323;stroke-dasharray:2.646, 1.323;stroke-dashoffset:0;stroke-opacity:1"
+       id="path511"
+       cx="126.68443"
+       cy="141.9209"
+       r="8.8512135" />
+    <circle
+       style="fill:#000000;fill-opacity:0;stroke:#060301;stroke-width:1.323;stroke-dasharray:none;stroke-dashoffset:0;stroke-opacity:1"
+       id="path511-3"
+       cx="77.499413"
+       cy="122.66233"
+       r="8.8512135" />
+  </g>
+</svg>

metastable_baselines/distributions/__init__.py

@@ -0,0 +1 @@
+#TODO: License or such

metastable_baselines/distributions/distributions.py

@@ -0,0 +1,197 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch as th
+from torch import nn
+from torch.distributions import Normal, MultivariateNormal
+
+from stable_baselines3.common.preprocessing import get_action_dim
+
+from stable_baselines3.common.distributions import Distribution as SB3_Distribution
+from stable_baselines3.common.distributions import DiagGaussianDistribution
+
+
+class ContextualCovDiagonalGaussianDistribution(DiagGaussianDistribution):
+    """
+    Gaussian distribution with diagonal covariance matrix, for continuous actions.
+    Includes contextual parametrization of the covariance matrix.
+
+    :param action_dim:  Dimension of the action space.
+    """
+
+    def __init__(self, action_dim: int):
+        super(ContextualCovDiagonalGaussianDistribution, self).__init__()
+
+    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
+        standard deviation (log std in fact to allow negative values)
+
+        :param latent_dim: Dimension of the last layer of the policy (before the action layer)
+        :param log_std_init: Initial value for the log standard deviation
+        :return:
+        """
+        mean_actions = nn.Linear(latent_dim, self.action_dim)
+        log_std = nn.Linear(latent_dim, self.action_dim)
+        return mean_actions, log_std
+
+
+class ContextualSqrtCovDiagonalGaussianDistribution(DiagGaussianDistribution):
+    """
+    Gaussian distribution induced by its sqrt(cov), for continuous actions.
+
+    :param action_dim:  Dimension of the action space.
+    """
+
+    def __init__(self, action_dim: int):
+        super(DiagGaussianDistribution, self).__init__()
+        self.action_dim = action_dim
+        self.mean_actions = None
+        self.log_std = None
+
+    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
+        standard deviation (log std in fact to allow negative values)
+
+        :param latent_dim: Dimension of the last layer of the policy (before the action layer)
+        :param log_std_init: Initial value for the log standard deviation
+        :return:
+        """
+        mean_actions = nn.Linear(latent_dim, self.action_dim)
+        # TODO: allow action dependent std
+        log_std = nn.Linear(latent_dim, (self.action_dim, self.action_dim))
+        return mean_actions, log_std
+
+    def proba_distribution(self, mean_actions: th.Tensor, log_std: th.Tensor) -> "DiagGaussianDistribution":
+        """
+        Create the distribution given its parameters (mean, std)
+
+        :param mean_actions:
+        :param log_std:
+        :return:
+        """
+        action_std = th.ones_like(mean_actions) * log_std.exp()
+        self.distribution = Normal(mean_actions, action_std)
+        return self
+
+    def log_prob(self, actions: th.Tensor) -> th.Tensor:
+        """
+        Get the log probabilities of actions according to the distribution.
+        Note that you must first call the ``proba_distribution()`` method.
+
+        :param actions:
+        :return:
+        """
+        log_prob = self.distribution.log_prob(actions)
+        return sum_independent_dims(log_prob)
+
+    def entropy(self) -> th.Tensor:
+        return sum_independent_dims(self.distribution.entropy())
+
+    def sample(self) -> th.Tensor:
+        # Reparametrization trick to pass gradients
+        return self.distribution.rsample()
+
+    def mode(self) -> th.Tensor:
+        return self.distribution.mean
+
+    def actions_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor, deterministic: bool = False) -> th.Tensor:
+        # Update the proba distribution
+        self.proba_distribution(mean_actions, log_std)
+        return self.get_actions(deterministic=deterministic)
+
+    def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Compute the log probability of taking an action
+        given the distribution parameters.
+
+        :param mean_actions:
+        :param log_std:
+        :return:
+        """
+        actions = self.actions_from_params(mean_actions, log_std)
+        log_prob = self.log_prob(actions)
+        return actions, log_prob
+
+
+class DiagGaussianDistribution(SB3_Distribution):
+    """
+    Gaussian distribution with full covariance matrix, for continuous actions.
+
+    :param action_dim:  Dimension of the action space.
+    """
+
+    def __init__(self, action_dim: int):
+        super(DiagGaussianDistribution, self).__init__()
+        self.action_dim = action_dim
+        self.mean_actions = None
+        self.log_std = None
+
+    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
+        standard deviation (log std in fact to allow negative values)
+
+        :param latent_dim: Dimension of the last layer of the policy (before the action layer)
+        :param log_std_init: Initial value for the log standard deviation
+        :return:
+        """
+        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, requires_grad=True)
+        return mean_actions, log_std
+
+    def proba_distribution(self, mean_actions: th.Tensor, log_std: th.Tensor) -> "DiagGaussianDistribution":
+        """
+        Create the distribution given its parameters (mean, std)
+
+        :param mean_actions:
+        :param log_std:
+        :return:
+        """
+        action_std = th.ones_like(mean_actions) * log_std.exp()
+        self.distribution = Normal(mean_actions, action_std)
+        return self
+
+    def log_prob(self, actions: th.Tensor) -> th.Tensor:
+        """
+        Get the log probabilities of actions according to the distribution.
+        Note that you must first call the ``proba_distribution()`` method.
+
+        :param actions:
+        :return:
+        """
+        log_prob = self.distribution.log_prob(actions)
+        return sum_independent_dims(log_prob)
+
+    def entropy(self) -> th.Tensor:
+        return sum_independent_dims(self.distribution.entropy())
+
+    def sample(self) -> th.Tensor:
+        # Reparametrization trick to pass gradients
+        return self.distribution.rsample()
+
+    def mode(self) -> th.Tensor:
+        return self.distribution.mean
+
+    def actions_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor, deterministic: bool = False) -> th.Tensor:
+        # Update the proba distribution
+        self.proba_distribution(mean_actions, log_std)
+        return self.get_actions(deterministic=deterministic)
+
+    def log_prob_from_params(self, mean_actions: th.Tensor, log_std: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Compute the log probability of taking an action
+        given the distribution parameters.
+
+        :param mean_actions:
+        :param log_std:
+        :return:
+        """
+        actions = self.actions_from_params(mean_actions, log_std)
+        log_prob = self.log_prob(actions)
+        return actions, log_prob

metastable_baselines/misc/distTools.py

@@ -17,6 +17,18 @@ def get_mean_and_chol(p, expand=False):
         raise Exception('Dist-Type not implemented')
 
 
+def get_mean_and_sqrt(p):
+    raise Exception('Not yet implemented...')
+    if isinstance(p, th.distributions.Normal):
+        return p.mean, p.stddev
+    elif isinstance(p, th.distributions.MultivariateNormal):
+        return p.mean, p.scale_tril
+    elif isinstance(p, SB3_Distribution):
+        return get_mean_and_chol(p.distribution)
+    else:
+        raise Exception('Dist-Type not implemented')
+
+
 def get_cov(p):
     if isinstance(p, th.distributions.Normal):
         return th.diag_embed(p.variance)

metastable_baselines/misc/norm.py

@@ -7,9 +7,9 @@ def mahalanobis_alt(u, v, std):
     return th.triangular_solve(delta, std, upper=False)[0].pow(2).sum([-2, -1])
 
 
-def mahalanobis(u, v, cov):
+def mahalanobis(u, v, chol):
     delta = u - v
-    return _batch_mahalanobis(cov, delta)
+    return _batch_mahalanobis(chol, delta)
 
 
 def frob_sq(diff, is_spd=False):

metastable_baselines/projections_orig/__init__.py

@@ -0,0 +1,15 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.

metastable_baselines/projections_orig/base_projection_layer.py

@@ -0,0 +1,374 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import copy
+import math
+import torch as ch
+from typing import Tuple, Union
+
+from trust_region_projections.models.policy.abstract_gaussian_policy import AbstractGaussianPolicy
+from trust_region_projections.utils.network_utils import get_optimizer
+from trust_region_projections.utils.projection_utils import gaussian_kl, get_entropy_schedule
+from trust_region_projections.utils.torch_utils import generate_minibatches, select_batch, tensorize
+
+
+def entropy_inequality_projection(policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                  beta: Union[float, ch.Tensor]):
+    """
+    Projects std to satisfy an entropy INEQUALITY constraint.
+    Args:
+        policy: policy instance
+        p: current distribution
+        beta: target entropy for EACH std or general bound for all stds
+
+    Returns:
+        projected std that satisfies the entropy bound
+    """
+    mean, std = p
+    k = std.shape[-1]
+    batch_shape = std.shape[:-2]
+
+    ent = policy.entropy(p)
+    mask = ent < beta
+
+    # if nothing has to be projected skip computation
+    if (~mask).all():
+        return p
+
+    alpha = ch.ones(batch_shape, dtype=std.dtype, device=std.device)
+    alpha[mask] = ch.exp((beta[mask] - ent[mask]) / k)
+
+    proj_std = ch.einsum('ijk,i->ijk', std, alpha)
+    return mean, ch.where(mask[..., None, None], proj_std, std)
+
+
+def entropy_equality_projection(policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                beta: Union[float, ch.Tensor]):
+    """
+    Projects std to satisfy an entropy EQUALITY constraint.
+    Args:
+        policy: policy instance
+        p: current distribution
+        beta: target entropy for EACH std or general bound for all stds
+
+    Returns:
+        projected std that satisfies the entropy bound
+    """
+    mean, std = p
+    k = std.shape[-1]
+
+    ent = policy.entropy(p)
+    alpha = ch.exp((beta - ent) / k)
+    proj_std = ch.einsum('ijk,i->ijk', std, alpha)
+    return mean, proj_std
+
+
+def mean_projection(mean: ch.Tensor, old_mean: ch.Tensor, maha: ch.Tensor, eps: ch.Tensor):
+    """
+    Projects the mean based on the Mahalanobis objective and trust region.
+    Args:
+        mean: current mean vectors
+        old_mean: old mean vectors
+        maha: Mahalanobis distance between the two mean vectors
+        eps: trust region bound
+
+    Returns:
+        projected mean that satisfies the trust region
+    """
+    batch_shape = mean.shape[:-1]
+    mask = maha > eps
+
+    ################################################################################################################
+    # mean projection maha
+
+    # if nothing has to be projected skip computation
+    if mask.any():
+        omega = ch.ones(batch_shape, dtype=mean.dtype, device=mean.device)
+        omega[mask] = ch.sqrt(maha[mask] / eps) - 1.
+        omega = ch.max(-omega, omega)[..., None]
+
+        m = (mean + omega * old_mean) / (1 + omega + 1e-16)
+        proj_mean = ch.where(mask[..., None], m, mean)
+    else:
+        proj_mean = mean
+
+    return proj_mean
+
+
+class BaseProjectionLayer(object):
+
+    def __init__(self,
+                 proj_type: str = "",
+                 mean_bound: float = 0.03,
+                 cov_bound: float = 1e-3,
+                 trust_region_coeff: float = 0.0,
+                 scale_prec: bool = True,
+
+                 entropy_schedule: Union[None, str] = None,
+                 action_dim: Union[None, int] = None,
+                 total_train_steps: Union[None, int] = None,
+                 target_entropy: float = 0.0,
+                 temperature: float = 0.5,
+                 entropy_eq: bool = False,
+                 entropy_first: bool = False,
+
+                 do_regression: bool = False,
+                 regression_iters: int = 1000,
+                 regression_lr: int = 3e-4,
+                 optimizer_type_reg: str = "adam",
+
+                 cpu: bool = True,
+                 dtype: ch.dtype = ch.float32,
+                 ):
+
+        """
+        Base projection layer, which can be used to compute metrics for non-projection approaches.
+        Args:
+           proj_type: Which type of projection to use. None specifies no projection and uses the TRPO objective.
+           mean_bound: projection bound for the step size w.r.t. mean
+           cov_bound: projection bound for the step size w.r.t. covariance matrix
+           trust_region_coeff: Coefficient for projection regularization loss term.
+           scale_prec: If true used mahalanobis distance for projections instead of euclidean with Sigma_old^-1.
+           entropy_schedule: Schedule type for entropy projection, one of 'linear', 'exp', None.
+           action_dim: number of action dimensions to scale exp decay correctly.
+           total_train_steps: total number of training steps to compute appropriate decay over time.
+           target_entropy: projection bound for the entropy of the covariance matrix
+           temperature: temperature decay for exponential entropy bound
+           entropy_eq: Use entropy equality constraints.
+           entropy_first: Project entropy before trust region.
+           do_regression: Conduct additional regression steps after the the policy steps to match projection and policy.
+           regression_iters: Number of regression steps.
+           regression_lr: Regression learning rate.
+           optimizer_type_reg: Optimizer for regression.
+           cpu: Compute on CPU only.
+           dtype: Data type to use, either of float32 or float64. The later might be necessary for higher
+                   dimensions in order to learn the full covariance.
+        """
+
+        # projection and bounds
+        self.proj_type = proj_type
+        self.mean_bound = tensorize(mean_bound, cpu=cpu, dtype=dtype)
+        self.cov_bound = tensorize(cov_bound, cpu=cpu, dtype=dtype)
+        self.trust_region_coeff = trust_region_coeff
+        self.scale_prec = scale_prec
+
+        # projection utils
+        assert (action_dim and total_train_steps) if entropy_schedule else True
+        self.entropy_proj = entropy_equality_projection if entropy_eq else entropy_inequality_projection
+        self.entropy_schedule = get_entropy_schedule(entropy_schedule, total_train_steps, dim=action_dim)
+        self.target_entropy = tensorize(target_entropy, cpu=cpu, dtype=dtype)
+        self.entropy_first = entropy_first
+        self.entropy_eq = entropy_eq
+        self.temperature = temperature
+        self._initial_entropy = None
+
+        # regression
+        self.do_regression = do_regression
+        self.regression_iters = regression_iters
+        self.lr_reg = regression_lr
+        self.optimizer_type_reg = optimizer_type_reg
+
+    def __call__(self, policy, p: Tuple[ch.Tensor, ch.Tensor], q, step, *args, **kwargs):
+        # entropy_bound = self.policy.entropy(q) - self.target_entropy
+        entropy_bound = self.entropy_schedule(self.initial_entropy, self.target_entropy, self.temperature,
+                                              step) * p[0].new_ones(p[0].shape[0])
+        return self._projection(policy, p, q, self.mean_bound, self.cov_bound, entropy_bound, **kwargs)
+
+    def _trust_region_projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                 q: Tuple[ch.Tensor, ch.Tensor], eps: ch.Tensor, eps_cov: ch.Tensor, **kwargs):
+        """
+        Hook for implementing the specific trust region projection
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: mean trust region bound
+            eps_cov: covariance trust region bound
+            **kwargs:
+
+        Returns:
+            projected
+        """
+        return p
+
+    # @final
+    def _projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                    q: Tuple[ch.Tensor, ch.Tensor], eps: ch.Tensor, eps_cov: ch.Tensor, beta: ch.Tensor, **kwargs):
+        """
+        Template method with hook _trust_region_projection() to encode specific functionality.
+        (Optional) entropy projection is executed before or after as specified by entropy_first.
+        Do not override this. For Python >= 3.8 you can use the @final decorator to enforce not overwriting.
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: mean trust region bound
+            eps_cov: covariance trust region bound
+            beta: entropy bound
+            **kwargs:
+
+        Returns:
+            projected mean, projected std
+        """
+
+        ####################################################################################################################
+        # entropy projection in the beginning
+        if self.entropy_first:
+            p = self.entropy_proj(policy, p, beta)
+
+        ####################################################################################################################
+        # trust region projection for mean and cov bounds
+        proj_mean, proj_std = self._trust_region_projection(policy, p, q, eps, eps_cov, **kwargs)
+
+        ####################################################################################################################
+        # entropy projection in the end
+        if self.entropy_first:
+            return proj_mean, proj_std
+
+        return self.entropy_proj(policy, (proj_mean, proj_std), beta)
+
+    @property
+    def initial_entropy(self):
+        return self._initial_entropy
+
+    @initial_entropy.setter
+    def initial_entropy(self, entropy):
+        if self.initial_entropy is None:
+            self._initial_entropy = entropy
+
+    def trust_region_value(self, policy, p, q):
+        """
+        Computes the KL divergence between two Gaussian distributions p and q.
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+        Returns:
+            Mean and covariance part of the trust region metric.
+        """
+        return gaussian_kl(policy, p, q)
+
+    def get_trust_region_loss(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                              proj_p: Tuple[ch.Tensor, ch.Tensor]):
+        """
+        Compute the trust region loss to ensure policy output and projection stay close.
+        Args:
+            policy: policy instance
+            proj_p: projected distribution
+            p: predicted distribution from network output
+
+        Returns:
+            trust region loss
+        """
+        p_target = (proj_p[0].detach(), proj_p[1].detach())
+        mean_diff, cov_diff = self.trust_region_value(policy, p, p_target)
+
+        delta_loss = (mean_diff + cov_diff if policy.contextual_std else mean_diff).mean()
+
+        return delta_loss * self.trust_region_coeff
+
+    def compute_metrics(self, policy, p, q) -> dict:
+        """
+        Returns dict with constraint metrics.
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+
+        Returns:
+            dict with constraint metrics
+        """
+        with ch.no_grad():
+            entropy_old = policy.entropy(q)
+            entropy = policy.entropy(p)
+            mean_kl, cov_kl = gaussian_kl(policy, p, q)
+            kl = mean_kl + cov_kl
+
+            mean_diff, cov_diff = self.trust_region_value(policy, p, q)
+
+            combined_constraint = mean_diff + cov_diff
+            entropy_diff = entropy_old - entropy
+
+        return {'kl': kl.detach().mean(),
+                'constraint': combined_constraint.mean(),
+                'mean_constraint': mean_diff.mean(),
+                'cov_constraint': cov_diff.mean(),
+                'entropy': entropy.mean(),
+                'entropy_diff': entropy_diff.mean(),
+                'kl_max': kl.max(),
+                'constraint_max': combined_constraint.max(),
+                'mean_constraint_max': mean_diff.max(),
+                'cov_constraint_max': cov_diff.max(),
+                'entropy_max': entropy.max(),
+                'entropy_diff_max': entropy_diff.max()
+                }
+
+    def trust_region_regression(self, policy: AbstractGaussianPolicy, obs: ch.Tensor, q: Tuple[ch.Tensor, ch.Tensor],
+                                n_minibatches: int, global_steps: int):
+        """
+        Take additional regression steps to match projection output and policy output.
+        The policy parameters are updated in-place.
+        Args:
+            policy: policy instance
+            obs: collected observations from trajectories
+            q: old distributions
+            n_minibatches: split the rollouts into n_minibatches.
+            global_steps: current number of steps, required for projection
+        Returns:
+            dict with mean of regession loss
+        """
+
+        if not self.do_regression:
+            return {}
+
+        policy_unprojected = copy.deepcopy(policy)
+        optim_reg = get_optimizer(self.optimizer_type_reg, policy_unprojected.parameters(), learning_rate=self.lr_reg)
+        optim_reg.reset()
+
+        reg_losses = obs.new_tensor(0.)
+
+        # get current projected values --> targets for regression
+        p_flat = policy(obs)
+        p_target = self(policy, p_flat, q, global_steps)
+
+        for _ in range(self.regression_iters):
+            batch_indices = generate_minibatches(obs.shape[0], n_minibatches)
+
+            # Minibatches SGD
+            for indices in batch_indices:
+                batch = select_batch(indices, obs, p_target[0], p_target[1])
+                b_obs, b_target_mean, b_target_std = batch
+                proj_p = (b_target_mean.detach(), b_target_std.detach())
+
+                p = policy_unprojected(b_obs)
+
+                # invert scaling with coeff here as we do not have to balance with other losses
+                loss = self.get_trust_region_loss(policy, p, proj_p) / self.trust_region_coeff
+
+                optim_reg.zero_grad()
+                loss.backward()
+                optim_reg.step()
+                reg_losses += loss.detach()
+
+        policy.load_state_dict(policy_unprojected.state_dict())
+
+        if not policy.contextual_std:
+            # set policy with projection value.
+            # In non-contextual cases we have only one cov, so the projection is the same.
+            policy.set_std(p_target[1][0])
+
+        steps = self.regression_iters * (math.ceil(obs.shape[0] / n_minibatches))
+        return {"regression_loss": (reg_losses / steps).detach()}

metastable_baselines/projections_orig/frob_projection_layer.py

@@ -0,0 +1,97 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import torch as ch
+from typing import Tuple
+
+from trust_region_projections.models.policy.abstract_gaussian_policy import AbstractGaussianPolicy
+from trust_region_projections.projections.base_projection_layer import BaseProjectionLayer, mean_projection
+from trust_region_projections.utils.projection_utils import gaussian_frobenius
+
+
+class FrobeniusProjectionLayer(BaseProjectionLayer):
+
+    def _trust_region_projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                 q: Tuple[ch.Tensor, ch.Tensor], eps: ch.Tensor, eps_cov: ch.Tensor, **kwargs):
+        """
+        Runs Frobenius projection layer and constructs cholesky of covariance
+
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: (modified) kl bound/ kl bound for mean part
+            eps_cov: (modified) kl bound for cov part
+            beta: (modified) entropy bound
+            **kwargs:
+        Returns: mean, cov cholesky
+        """
+
+        mean, chol = p
+        old_mean, old_chol = q
+        batch_shape = mean.shape[:-1]
+
+        ####################################################################################################################
+        # precompute mean and cov part of frob projection, which are used for the projection.
+        mean_part, cov_part, cov, cov_old = gaussian_frobenius(policy, p, q, self.scale_prec, True)
+
+        ################################################################################################################
+        # mean projection maha/euclidean
+
+        proj_mean = mean_projection(mean, old_mean, mean_part, eps)
+
+        ################################################################################################################
+        # cov projection frobenius
+
+        cov_mask = cov_part > eps_cov
+
+        if cov_mask.any():
+            # alpha = ch.where(fro_norm_sq > eps_cov, ch.sqrt(fro_norm_sq / eps_cov) - 1., ch.tensor(1.))
+            eta = ch.ones(batch_shape, dtype=chol.dtype, device=chol.device)
+            eta[cov_mask] = ch.sqrt(cov_part[cov_mask] / eps_cov) - 1.
+            eta = ch.max(-eta, eta)
+
+            new_cov = (cov + ch.einsum('i,ijk->ijk', eta, cov_old)) / (1. + eta + 1e-16)[..., None, None]
+            proj_chol = ch.where(cov_mask[..., None, None], ch.cholesky(new_cov), chol)
+        else:
+            proj_chol = chol
+
+        return proj_mean, proj_chol
+
+    def trust_region_value(self, policy, p, q):
+        """
+        Computes the Frobenius metric between two Gaussian distributions p and q.
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+        Returns:
+            mean and covariance part of Frobenius metric
+        """
+        return gaussian_frobenius(policy, p, q, self.scale_prec)
+
+    def get_trust_region_loss(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                              proj_p: Tuple[ch.Tensor, ch.Tensor]):
+
+        mean_diff, _ = self.trust_region_value(policy, p, proj_p)
+        if policy.contextual_std:
+            # Compute MSE here, because we found the Frobenius norm tends to generate values that explode for the cov
+            cov_diff = (p[1] - proj_p[1]).pow(2).sum([-1, -2])
+            delta_loss = (mean_diff + cov_diff).mean()
+        else:
+            delta_loss = mean_diff.mean()
+
+        return delta_loss * self.trust_region_coeff

metastable_baselines/projections_orig/kl_projection_layer.py

@@ -0,0 +1,101 @@
+import cpp_projection
+import numpy as np
+import torch as ch
+from typing import Any, Tuple
+
+from trust_region_projections.models.policy.abstract_gaussian_policy import AbstractGaussianPolicy
+from trust_region_projections.projections.base_projection_layer import BaseProjectionLayer, mean_projection
+from trust_region_projections.utils.projection_utils import gaussian_kl
+from trust_region_projections.utils.torch_utils import get_numpy
+
+
+class KLProjectionLayer(BaseProjectionLayer):
+
+    def _trust_region_projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                 q: Tuple[ch.Tensor, ch.Tensor], eps: ch.Tensor, eps_cov: ch.Tensor, **kwargs):
+        """
+        Runs KL projection layer and constructs cholesky of covariance
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: (modified) kl bound/ kl bound for mean part
+            eps_cov: (modified) kl bound for cov part
+            **kwargs:
+
+        Returns:
+            projected mean, projected cov cholesky
+        """
+        mean, std = p
+        old_mean, old_std = q
+
+        if not policy.contextual_std:
+            # only project first one to reduce number of numerical optimizations
+            std = std[:1]
+            old_std = old_std[:1]
+
+        ################################################################################################################
+        # project mean with closed form
+        mean_part, _ = gaussian_kl(policy, p, q)
+        proj_mean = mean_projection(mean, old_mean, mean_part, eps)
+
+        cov = policy.covariance(std)
+        old_cov = policy.covariance(old_std)
+
+        if policy.is_diag:
+            proj_cov = KLProjectionGradFunctionDiagCovOnly.apply(cov.diagonal(dim1=-2, dim2=-1),
+                                                                 old_cov.diagonal(dim1=-2, dim2=-1),
+                                                                 eps_cov)
+            proj_std = proj_cov.sqrt().diag_embed()
+        else:
+            raise NotImplementedError("The KL projection currently does not support full covariance matrices.")
+
+        if not policy.contextual_std:
+            # scale first std back to batchsize
+            proj_std = proj_std.expand(mean.shape[0], -1, -1)
+
+        return proj_mean, proj_std
+
+
+class KLProjectionGradFunctionDiagCovOnly(ch.autograd.Function):
+    projection_op = None
+
+    @staticmethod
+    def get_projection_op(batch_shape, dim, max_eval=100):
+        if not KLProjectionGradFunctionDiagCovOnly.projection_op:
+            KLProjectionGradFunctionDiagCovOnly.projection_op = \
+                cpp_projection.BatchedDiagCovOnlyProjection(batch_shape, dim, max_eval=max_eval)
+        return KLProjectionGradFunctionDiagCovOnly.projection_op
+
+    @staticmethod
+    def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        std, old_std, eps_cov = args
+
+        batch_shape = std.shape[0]
+        dim = std.shape[-1]
+
+        cov_np = get_numpy(std)
+        old_std = get_numpy(old_std)
+        eps = get_numpy(eps_cov) * np.ones(batch_shape)
+
+        # p_op = cpp_projection.BatchedDiagCovOnlyProjection(batch_shape, dim)
+        # ctx.proj = projection_op
+
+        p_op = KLProjectionGradFunctionDiagCovOnly.get_projection_op(batch_shape, dim)
+        ctx.proj = p_op
+
+        proj_std = p_op.forward(eps, old_std, cov_np)
+
+        return std.new(proj_std)
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        projection_op = ctx.proj
+        d_std, = grad_outputs
+
+        d_std_np = get_numpy(d_std)
+        d_std_np = np.atleast_2d(d_std_np)
+        df_stds = projection_op.backward(d_std_np)
+        df_stds = np.atleast_2d(df_stds)
+
+        return d_std.new(df_stds), None, None

metastable_baselines/projections_orig/papi_projection.py

@@ -0,0 +1,233 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import logging
+
+import copy
+import numpy as np
+import torch as ch
+from typing import Tuple, Union
+
+from trust_region_projections.utils.projection_utils import gaussian_kl
+from trust_region_projections.models.policy.abstract_gaussian_policy import AbstractGaussianPolicy
+from trust_region_projections.projections.base_projection_layer import BaseProjectionLayer
+from trust_region_projections.utils.torch_utils import torch_batched_trace
+
+logger = logging.getLogger("papi_projection")
+
+
+class PAPIProjection(BaseProjectionLayer):
+
+    def __init__(self,
+                 proj_type: str = "",
+                 mean_bound: float = 0.015,
+                 cov_bound: float = 0.0,
+
+                 entropy_eq: bool = False,
+                 entropy_first: bool = True,
+
+                 cpu: bool = True,
+                 dtype: ch.dtype = ch.float32,
+                 **kwargs
+                 ):
+
+        """
+        PAPI projection, which can be used after each training epoch to satisfy the trust regions.
+        Args:
+           proj_type: Which type of projection to use. None specifies no projection and uses the TRPO objective.
+           mean_bound: projection bound for the step size,
+                       PAPI only has a joint KL constraint, mean and cov bound are summed for this bound.
+           cov_bound: projection bound for the step size,
+                      PAPI only has a joint KL constraint, mean and cov bound are summed for this bound.
+           entropy_eq: Use entropy equality constraints.
+           entropy_first: Project entropy before trust region.
+           cpu: Compute on CPU only.
+           dtype: Data type to use, either of float32 or float64. The later might be necessary for higher
+                   dimensions in order to learn the full covariance.
+        """
+
+        assert entropy_first
+        super().__init__(proj_type, mean_bound, cov_bound, 0.0, False, None, None, None, 0.0, 0.0, entropy_eq,
+                         entropy_first, cpu, dtype)
+
+        self.last_policies = []
+
+    def __call__(self, policy, p, q, step=0, *args, **kwargs):
+        if kwargs.get("obs"):
+            self._papi_steps(policy, q, **kwargs)
+        else:
+            return p
+
+    def _trust_region_projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                 q: Tuple[ch.Tensor, ch.Tensor], eps: Union[ch.Tensor, float],
+                                 eps_cov: Union[ch.Tensor, float], **kwargs):
+        """
+        runs papi projection layer and constructs sqrt of covariance
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: (modified) kl bound/ kl bound for mean part
+            eps_cov: (modified) kl bound for cov part
+            **kwargs:
+
+        Returns:
+            mean, cov sqrt
+        """
+
+        mean, chol = p
+        old_mean, old_chol = q
+        intermed_mean = kwargs.get('intermed_mean')
+
+        dtype = mean.dtype
+        device = mean.device
+
+        dim = mean.shape[-1]
+
+        ################################################################################################################
+        # Precompute basic matrices
+
+        # Joint bound
+        eps += eps_cov
+
+        I = ch.eye(dim, dtype=dtype, device=device)
+        old_precision = ch.cholesky_solve(I, old_chol)[0]
+        logdet_old = policy.log_determinant(old_chol)
+        cov = policy.covariance(chol)
+
+        ################################################################################################################
+        # compute expected KL
+        maha_part, cov_part = gaussian_kl(policy, p, q)
+        maha_part = maha_part.mean()
+        cov_part = cov_part.mean()
+
+        if intermed_mean is not None:
+            maha_intermediate = 0.5 * policy.maha(intermed_mean, old_mean, old_chol).mean()
+            mm = ch.min(maha_part, maha_intermediate)
+
+        ################################################################################################################
+        # matrix rotation/rescaling projection
+        if maha_part + cov_part > eps + 1e-6:
+            old_cov = policy.covariance(old_chol)
+
+            maha_delta = eps if intermed_mean is None else (eps - mm)
+            eta_rot = maha_delta / ch.max(maha_part + cov_part, ch.tensor(1e-16, dtype=dtype, device=device))
+            new_cov = (1 - eta_rot) * old_cov + eta_rot * cov
+            proj_chol = ch.cholesky(new_cov)
+
+            # recompute covariance part of KL for new chol
+            trace_term = 0.5 * (torch_batched_trace(old_precision @ new_cov) - dim).mean()  # rotation difference
+            entropy_diff = 0.5 * (logdet_old - policy.log_determinant(proj_chol)).mean()
+
+            cov_part = trace_term + entropy_diff
+
+        else:
+            proj_chol = chol
+
+        ################################################################################################################
+        # mean interpolation projection
+        if maha_part + cov_part > eps + 1e-6:
+
+            if intermed_mean is not None:
+                a = 0.5 * policy.maha(mean, intermed_mean, old_chol).mean()
+                b = 0.5 * ((mean - intermed_mean) @ old_precision @ (intermed_mean - old_mean).T).mean()
+                c = maha_intermediate - ch.max(eps - cov_part, ch.tensor(0., dtype=dtype, device=device))
+                eta_mean = (-b + ch.sqrt(ch.max(b * b - a * c, ch.tensor(1e-16, dtype=dtype, device=device)))) / \
+                           ch.max(a, ch.tensor(1e-16, dtype=dtype, device=device))
+            else:
+                eta_mean = ch.sqrt(
+                    ch.max(eps - cov_part, ch.tensor(1e-16, dtype=dtype, device=device)) /
+                    ch.max(maha_part, ch.tensor(1e-16, dtype=dtype, device=device)))
+        else:
+            eta_mean = ch.tensor(1., dtype=dtype, device=device)
+
+        return eta_mean, proj_chol
+
+    def _papi_steps(self, policy: AbstractGaussianPolicy, q: Tuple[ch.Tensor, ch.Tensor], obs: ch.Tensor, lr_schedule,
+                    lr_schedule_vf=None):
+        """
+        Take PAPI steps after PPO finished its steps. Policy parameters are updated in-place.
+        Args:
+            policy: policy instance
+            q: old distribution
+            obs: collected observations from trajectories
+            lr_schedule: lr schedule for policy
+            lr_schedule_vf: lr schedule for vf
+
+        Returns:
+
+        """
+        assert not policy.contextual_std
+
+        # save latest policy in history
+        self.last_policies.append(copy.deepcopy(policy))
+
+        ################################################################################################################
+        # policy backtracking: out of last n policies and current one find one that satisfies the kl constraint
+
+        intermed_policy = None
+        n_backtracks = 0
+
+        for i, pi in enumerate(reversed(self.last_policies)):
+            p_prime = pi(obs)
+            mean_part, cov_part = pi.kl_divergence(p_prime, q)
+            if (mean_part + cov_part).mean() <= self.mean_bound + self.cov_bound:
+                intermed_policy = pi
+                n_backtracks = i
+                break
+
+        ################################################################################################################
+        # LR update
+
+        # reduce learning rate when appropriate policy not within the last 4 epochs
+        if n_backtracks >= 4 or intermed_policy is None:
+            # Linear learning rate annealing
+            lr_schedule.step()
+            if lr_schedule_vf:
+                lr_schedule_vf.step()
+
+        if intermed_policy is None:
+            # pop last policy and make it current one, as the updated one was poor
+            # do not keep last policy in history, otherwise we could stack the same policy multiple times.
+            if len(self.last_policies) >= 1:
+                policy.load_state_dict(self.last_policies.pop().state_dict())
+                logger.warning(f"No suitable policy found in backtracking of {len(self.last_policies)} policies.")
+            return
+
+        ################################################################################################################
+        # PAPI iterations
+
+        # We assume only non contextual covariances here, therefore we only need to project for one
+        q = (q[0], q[1][:1])  # (means, covs[:1])
+
+        # This is A from Alg. 2 [Akrour et al., 2019]
+        intermed_weight = intermed_policy.get_last_layer().detach().clone()
+        # This is A @ phi(s)
+        intermed_mean = p_prime[0].detach().clone()
+
+        entropy = policy.entropy(q)
+        entropy_bound = obs.new_tensor([-np.inf]) if entropy / self.initial_entropy > 0.5 \
+            else entropy - (self.mean_bound + self.cov_bound)
+
+        for _ in range(20):
+            eta, proj_chol = self._projection(intermed_policy, (p_prime[0], p_prime[1][:1]), q,
+                                              self.mean_bound, self.cov_bound, entropy_bound,
+                                              intermed_mean=intermed_mean)
+            intermed_policy.papi_weight_update(eta, intermed_weight)
+            intermed_policy.set_std(proj_chol[0])
+            p_prime = intermed_policy(obs)
+
+        policy.load_state_dict(intermed_policy.state_dict())

metastable_baselines/projections_orig/projection_factory.py

@@ -0,0 +1,54 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from trust_region_projections.projections.base_projection_layer import BaseProjectionLayer
+from trust_region_projections.projections.frob_projection_layer import FrobeniusProjectionLayer
+from trust_region_projections.projections.kl_projection_layer import KLProjectionLayer
+from trust_region_projections.projections.papi_projection import PAPIProjection
+from trust_region_projections.projections.w2_projection_layer import WassersteinProjectionLayer
+
+
+def get_projection_layer(proj_type: str = "", **kwargs) -> BaseProjectionLayer:
+    """
+    Factory to generate the projection layers for all projections.
+    Args:
+        proj_type: One of None/' ', 'ppo', 'papi', 'w2', 'w2_non_com', 'frob', 'kl', or 'entropy'
+        **kwargs: arguments for projection layer
+
+    Returns:
+
+    """
+    if not proj_type or proj_type.isspace() or proj_type.lower() in ["ppo", "sac", "td3", "mpo", "entropy"]:
+        return BaseProjectionLayer(proj_type, **kwargs)
+
+    elif proj_type.lower() == "w2":
+        return WassersteinProjectionLayer(proj_type, **kwargs)
+
+    elif proj_type.lower() == "frob":
+        return FrobeniusProjectionLayer(proj_type, **kwargs)
+
+    elif proj_type.lower() == "kl":
+        return KLProjectionLayer(proj_type, **kwargs)
+
+    elif proj_type.lower() == "papi":
+        # papi has a different approach compared to our projections.
+        # It has to be applied after the training with PPO.
+        return PAPIProjection(proj_type, **kwargs)
+
+    else:
+        raise ValueError(
+            f"Invalid projection type {proj_type}."
+            f" Choose one of None/' ', 'ppo', 'papi', 'w2', 'w2_non_com', 'frob', 'kl', or 'entropy'.")

metastable_baselines/projections_orig/w2_projection_layer.py

@@ -0,0 +1,84 @@
+#   Copyright (c) 2021 Robert Bosch GmbH
+#   Author: Fabian Otto
+#
+#   This program is free software: you can redistribute it and/or modify
+#   it under the terms of the GNU Affero General Public License as published
+#   by the Free Software Foundation, either version 3 of the License, or
+#   (at your option) any later version.
+#
+#   This program is distributed in the hope that it will be useful,
+#   but WITHOUT ANY WARRANTY; without even the implied warranty of
+#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#   GNU Affero General Public License for more details.
+#
+#   You should have received a copy of the GNU Affero General Public License
+#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import torch as ch
+from typing import Tuple
+
+from trust_region_projections.models.policy.abstract_gaussian_policy import AbstractGaussianPolicy
+from trust_region_projections.projections.base_projection_layer import BaseProjectionLayer, mean_projection
+from trust_region_projections.utils.projection_utils import gaussian_wasserstein_commutative
+
+
+class WassersteinProjectionLayer(BaseProjectionLayer):
+
+    def _trust_region_projection(self, policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
+                                 q: Tuple[ch.Tensor, ch.Tensor], eps: ch.Tensor, eps_cov: ch.Tensor, **kwargs):
+        """
+        Runs commutative Wasserstein projection layer and constructs sqrt of covariance
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+            eps: (modified) kl bound/ kl bound for mean part
+            eps_cov: (modified) kl bound for cov part
+            **kwargs:
+
+        Returns:
+            mean, cov sqrt
+        """
+        mean, sqrt = p
+        old_mean, old_sqrt = q
+        batch_shape = mean.shape[:-1]
+
+        ####################################################################################################################
+        # precompute mean and cov part of W2, which are used for the projection.
+        # Both parts differ based on precision scaling.
+        # If activated, the mean part is the maha distance and the cov has a more complex term in the inner parenthesis.
+        mean_part, cov_part = gaussian_wasserstein_commutative(policy, p, q, self.scale_prec)
+
+        ####################################################################################################################
+        # project mean (w/ or w/o precision scaling)
+        proj_mean = mean_projection(mean, old_mean, mean_part, eps)
+
+        ####################################################################################################################
+        # project covariance (w/ or w/o precision scaling)
+
+        cov_mask = cov_part > eps_cov
+
+        if cov_mask.any():
+            # gradient issue with ch.where, it executes both paths and gives NaN gradient.
+            eta = ch.ones(batch_shape, dtype=sqrt.dtype, device=sqrt.device)
+            eta[cov_mask] = ch.sqrt(cov_part[cov_mask] / eps_cov) - 1.
+            eta = ch.max(-eta, eta)
+
+            new_sqrt = (sqrt + ch.einsum('i,ijk->ijk', eta, old_sqrt)) / (1. + eta + 1e-16)[..., None, None]
+            proj_sqrt = ch.where(cov_mask[..., None, None], new_sqrt, sqrt)
+        else:
+            proj_sqrt = sqrt
+
+        return proj_mean, proj_sqrt
+
+    def trust_region_value(self, policy, p, q):
+        """
+        Computes the Wasserstein distance between two Gaussian distributions p and q.
+        Args:
+            policy: policy instance
+            p: current distribution
+            q: old distribution
+        Returns:
+            mean and covariance part of Wasserstein distance
+        """
+        return gaussian_wasserstein_commutative(policy, p, q, scale_prec=self.scale_prec)

replay.py

@@ -10,7 +10,7 @@ from stable_baselines3 import SAC, PPO, A2C
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, MultiInputActorCriticPolicy
 
-from sb3_trl.trl_pg import TRL_PG
+from metastable_baselines.trl_pg import TRL_PG
 import columbus
 
 

sb3_trl/trl_pg/__init__.py

@@ -1,2 +0,0 @@
-from sb3_trl.trl_pg.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
-from sb3_trl.trl_pg.trl_pg import TRL_PG

sb3_trl/trl_pg/policies.py

@@ -1,7 +0,0 @@
-# This file is here just to define MlpPolicy/CnnPolicy
-# that work for TRL_PG
-from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, MultiInputActorCriticPolicy
-
-MlpPolicy = ActorCriticPolicy
-CnnPolicy = ActorCriticCnnPolicy
-MultiInputPolicy = MultiInputActorCriticPolicy

sb3_trl/trl_pg/trl_pg.py

@@ -1,520 +0,0 @@
-import warnings
-from typing import Any, Dict, Optional, Type, Union, NamedTuple
-
-import numpy as np
-import torch as th
-from gym import spaces
-from torch.nn import functional as F
-
-from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
-from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, BasePolicy, MultiInputActorCriticPolicy
-from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
-from stable_baselines3.common.utils import explained_variance, get_schedule_fn
-from stable_baselines3.common.vec_env import VecEnv
-from stable_baselines3.common.buffers import RolloutBuffer
-from stable_baselines3.common.callbacks import BaseCallback
-from stable_baselines3.common.utils import obs_as_tensor
-from stable_baselines3.common.vec_env import VecNormalize
-
-from ..projections.base_projection_layer import BaseProjectionLayer
-from ..projections.frob_projection_layer import FrobeniusProjectionLayer
-from ..projections.w2_projection_layer import WassersteinProjectionLayer
-
-from ..misc.rollout_buffer import GaussianRolloutBuffer, GaussianRolloutBufferSamples
-
-
-class TRL_PG(OnPolicyAlgorithm):
-    """
-    Differential Trust Region Layer (TRL) for Policy Gradient (PG)
-
-    Paper: https://arxiv.org/abs/2101.09207
-    Code: This implementation borrows (/steals most) code from SB3's PPO implementation https://github.com/DLR-RM/stable-baselines3/blob/master/stable_baselines3/ppo/ppo.py
-    The implementation of the TRL-specific parts borrows from https://github.com/boschresearch/trust-region-layers/blob/main/trust_region_projections/algorithms/pg/pg.py
-
-    :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
-    :param env: The environment to learn from (if registered in Gym, can be str)
-    :param learning_rate: The learning rate, it can be a function
-        of the current progress remaining (from 1 to 0)
-    :param n_steps: The number of steps to run for each environment per update
-        (i.e. rollout buffer size is n_steps * n_envs where n_envs is number of environment copies running in parallel)
-        NOTE: n_steps * n_envs must be greater than 1 (because of the advantage normalization)
-        See https://github.com/pytorch/pytorch/issues/29372
-    :param batch_size: Minibatch size
-    :param n_epochs: Number of epoch when optimizing the surrogate loss
-    :param gamma: Discount factor
-    :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator
-    :param clip_range: Clipping parameter, it can be a function of the current progress
-        remaining (from 1 to 0).
-    :param clip_range_vf: Clipping parameter for the value function,
-        it can be a function of the current progress remaining (from 1 to 0).
-        This is a parameter specific to the OpenAI implementation. If None is passed (default),
-        no clipping will be done on the value function.
-        IMPORTANT: this clipping depends on the reward scaling.
-    :param normalize_advantage: Whether to normalize or not the advantage
-    :param ent_coef: Entropy coefficient for the loss calculation
-    :param vf_coef: Value function coefficient for the loss calculation
-    :param max_grad_norm: The maximum value for the gradient clipping
-    :param use_sde: Whether to use generalized State Dependent Exploration (gSDE)
-        instead of action noise exploration (default: False)
-    :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
-        Default: -1 (only sample at the beginning of the rollout)
-    :param target_kl: Limit the KL divergence between updates,
-        because the clipping is not enough to prevent large update
-        # 213 (cf https://github.com/hill-a/stable-baselines/issues/213)
-        see issue
-        By default, there is no limit on the kl div.
-    :param tensorboard_log: the log location for tensorboard (if None, no logging)
-    :param create_eval_env: Whether to create a second environment that will be
-        used for evaluating the agent periodically. (Only available when passing string for the environment)
-    :param policy_kwargs: additional arguments to be passed to the policy on creation
-    :param verbose: the verbosity level: 0 no output, 1 info, 2 debug
-    :param seed: Seed for the pseudo random generators
-    :param device: Device (cpu, cuda, ...) on which the code should be run.
-        Setting it to auto, the code will be run on the GPU if possible.
-    :param projection: What kind of Projection to use
-    :param _init_setup_model: Whether or not to build the network at the creation of the instance
-    """
-
-    policy_aliases: Dict[str, Type[BasePolicy]] = {
-        "MlpPolicy": ActorCriticPolicy,
-        "CnnPolicy": ActorCriticCnnPolicy,
-        "MultiInputPolicy": MultiInputActorCriticPolicy,
-    }
-
-    def __init__(
-        self,
-        policy: Union[str, Type[ActorCriticPolicy]],
-        env: Union[GymEnv, str],
-        learning_rate: Union[float, Schedule] = 3e-4,
-        n_steps: int = 2048,
-        batch_size: int = 64,
-        n_epochs: int = 10,
-        gamma: float = 0.99,
-        gae_lambda: float = 0.95,
-        clip_range: Union[float, Schedule] = 0.2,
-        clip_range_vf: Union[None, float, Schedule] = None,
-        normalize_advantage: bool = True,
-        ent_coef: float = 0.0,
-        vf_coef: float = 0.5,
-        max_grad_norm: float = 0.5,
-        use_sde: bool = False,
-        sde_sample_freq: int = -1,
-        target_kl: Optional[float] = None,
-        tensorboard_log: Optional[str] = None,
-        create_eval_env: bool = False,
-        policy_kwargs: Optional[Dict[str, Any]] = None,
-        verbose: int = 0,
-        seed: Optional[int] = None,
-        device: Union[th.device, str] = "auto",
-
-        # Different from PPO:
-        projection: BaseProjectionLayer = WassersteinProjectionLayer(),
-        #projection: BaseProjectionLayer = FrobeniusProjectionLayer(),
-        #projection: BaseProjectionLayer = BaseProjectionLayer(),
-
-        _init_setup_model: bool = True,
-    ):
-
-        super().__init__(
-            policy,
-            env,
-            learning_rate=learning_rate,
-            n_steps=n_steps,
-            gamma=gamma,
-            gae_lambda=gae_lambda,
-            ent_coef=ent_coef,
-            vf_coef=vf_coef,
-            max_grad_norm=max_grad_norm,
-            use_sde=use_sde,
-            sde_sample_freq=sde_sample_freq,
-            tensorboard_log=tensorboard_log,
-            policy_kwargs=policy_kwargs,
-            verbose=verbose,
-            device=device,
-            create_eval_env=create_eval_env,
-            seed=seed,
-            _init_setup_model=False,
-            supported_action_spaces=(
-                spaces.Box,
-                # spaces.Discrete,
-                # spaces.MultiDiscrete,
-                # spaces.MultiBinary,
-            ),
-        )
-
-        # Sanity check, otherwise it will lead to noisy gradient and NaN
-        # because of the advantage normalization
-        if normalize_advantage:
-            assert (
-                batch_size > 1
-            ), "`batch_size` must be greater than 1. See https://github.com/DLR-RM/stable-baselines3/issues/440"
-
-        if self.env is not None:
-            # Check that `n_steps * n_envs > 1` to avoid NaN
-            # when doing advantage normalization
-            buffer_size = self.env.num_envs * self.n_steps
-            assert (
-                buffer_size > 1
-            ), f"`n_steps * n_envs` must be greater than 1. Currently n_steps={self.n_steps} and n_envs={self.env.num_envs}"
-            # Check that the rollout buffer size is a multiple of the mini-batch size
-            untruncated_batches = buffer_size // batch_size
-            if buffer_size % batch_size > 0:
-                warnings.warn(
-                    f"You have specified a mini-batch size of {batch_size},"
-                    f" but because the `RolloutBuffer` is of size `n_steps * n_envs = {buffer_size}`,"
-                    f" after every {untruncated_batches} untruncated mini-batches,"
-                    f" there will be a truncated mini-batch of size {buffer_size % batch_size}\n"
-                    f"We recommend using a `batch_size` that is a factor of `n_steps * n_envs`.\n"
-                    f"Info: (n_steps={self.n_steps} and n_envs={self.env.num_envs})"
-                )
-        self.batch_size = batch_size
-        self.n_epochs = n_epochs
-        self.clip_range = clip_range
-        self.clip_range_vf = clip_range_vf
-        self.normalize_advantage = normalize_advantage
-        self.target_kl = target_kl
-
-        # Different from PPO:
-        self.projection = projection
-        self._global_steps = 0
-
-        if _init_setup_model:
-            self._setup_model()
-
-    def _setup_model(self) -> None:
-        super()._setup_model()
-
-        # Initialize schedules for policy/value clipping
-        self.clip_range = get_schedule_fn(self.clip_range)
-        if self.clip_range_vf is not None:
-            if isinstance(self.clip_range_vf, (float, int)):
-                assert self.clip_range_vf > 0, "`clip_range_vf` must be positive, " "pass `None` to deactivate vf clipping"
-
-            self.clip_range_vf = get_schedule_fn(self.clip_range_vf)
-
-        # Changed from PPO: We need a bigger RolloutBuffer
-        self.rollout_buffer = GaussianRolloutBuffer(
-            self.n_steps,
-            self.observation_space,
-            self.action_space,
-            device=self.device,
-            gamma=self.gamma,
-            gae_lambda=self.gae_lambda,
-            n_envs=self.n_envs,
-        )
-
-    def train(self) -> None:
-        """
-        Update policy using the currently gathered rollout buffer.
-        """
-        # Switch to train mode (this affects batch norm / dropout)
-        self.policy.set_training_mode(True)
-        # Update optimizer learning rate
-        self._update_learning_rate(self.policy.optimizer)
-        # Compute current clip range
-        clip_range = self.clip_range(self._current_progress_remaining)
-        # Optional: clip range for the value function
-        if self.clip_range_vf is not None:
-            clip_range_vf = self.clip_range_vf(
-                self._current_progress_remaining)
-
-        surrogate_losses = []
-        entropy_losses = []
-        trust_region_losses = []
-        pg_losses, value_losses = [], []
-        clip_fractions = []
-
-        continue_training = True
-
-        # train for n_epochs epochs
-        for epoch in range(self.n_epochs):
-            approx_kl_divs = []
-            # Do a complete pass on the rollout buffer
-            for rollout_data in self.rollout_buffer.get(self.batch_size):
-                # This is new compared to PPO.
-                # Calculating the TR-Projections we need to know the step number
-                self._global_steps += 1
-
-                actions = rollout_data.actions
-                if isinstance(self.action_space, spaces.Discrete):
-                    # Convert discrete action from float to long
-                    actions = rollout_data.actions.long().flatten()
-
-                # Re-sample the noise matrix because the log_std has changed
-                if self.use_sde:
-                    self.policy.reset_noise(self.batch_size)
-
-                # old code for PPO
-                # values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
-
-                # src in TRL reference code:
-                # Stolen from Fabian's Code (Public Version):
-                # p = self.policy(rollout_data.observations)
-                # proj_p = self.projection(self.policy, p, b_q = (b_old_mean, b_old_std), self._global_step)
-                # new_logpacs = self.policy.log_probability(proj_p, b_actions)
-
-                # src of evaluate_actions:
-                # pol = self.policy
-                # features = pol.extract_features(rollout_data.observations)
-                # latent_pi, latent_vf = pol.mlp_extractor(features)
-                # distribution = pol._get_action_dist_from_latent(latent_pi)
-                # log_prob = distribution.log_prob(actions)
-                # values = pol.value_net(latent_vf)
-                # return values, log_prob, distribution.entropy()
-                # entropy = distribution.entropy()
-
-                # here we go:
-                pol = self.policy
-                features = pol.extract_features(rollout_data.observations)
-                latent_pi, latent_vf = pol.mlp_extractor(features)
-                p = pol._get_action_dist_from_latent(latent_pi)
-                p_dist = p.distribution
-                # q_means = rollout_data.means
-                # if len(rollout_data.stds.shape) == 1:  # only diag
-                #    q_stds = th.diag(rollout_data.stds)
-                # else:
-                #    q_stds = rollout_data.stds
-                # q_dist = th.distributions.MultivariateNormal(
-                #    q_means, q_stds)
-                q_dist = th.distributions.Normal(
-                    rollout_data.means, rollout_data.stds)
-                proj_p = self.projection(p_dist, q_dist, self._global_steps)
-                log_prob = proj_p.log_prob(actions).sum(dim=1)
-                values = self.policy.value_net(latent_vf)
-                entropy = proj_p.entropy()
-
-                values = values.flatten()
-                # Normalize advantage
-                advantages = rollout_data.advantages
-                if self.normalize_advantage:
-                    advantages = (advantages - advantages.mean()
-                                  ) / (advantages.std() + 1e-8)
-
-                # ratio between old and new policy, should be one at the first iteration
-                ratio = th.exp(log_prob - rollout_data.old_log_prob)
-
-                # Difference from PPO: We renamed 'policy_loss' to 'surrogate_loss'
-                # clipped surrogate loss
-                surrogate_loss_1 = advantages * ratio
-                surrogate_loss_2 = advantages * \
-                    th.clamp(ratio, 1 - clip_range, 1 + clip_range)
-                surrogate_loss = - \
-                    th.min(surrogate_loss_1, surrogate_loss_2).mean()
-
-                surrogate_losses.append(surrogate_loss.item())
-
-                clip_fraction = th.mean(
-                    (th.abs(ratio - 1) > clip_range).float()).item()
-                clip_fractions.append(clip_fraction)
-
-                if self.clip_range_vf is None:
-                    # No clipping
-                    values_pred = values
-                else:
-                    # Clip the different between old and new value
-                    # NOTE: this depends on the reward scaling
-                    values_pred = rollout_data.old_values + th.clamp(
-                        values - rollout_data.old_values, -clip_range_vf, clip_range_vf
-                    )
-                # Value loss using the TD(gae_lambda) target
-                value_loss = F.mse_loss(rollout_data.returns, values_pred)
-                value_losses.append(value_loss.item())
-
-                # Entropy loss favor exploration
-                if entropy is None:
-                    # Approximate entropy when no analytical form
-                    entropy_loss = -th.mean(-log_prob)
-                else:
-                    entropy_loss = -th.mean(entropy)
-
-                entropy_losses.append(entropy_loss.item())
-
-                # Difference to PPO: Added trust_region_loss; policy_loss includes entropy_loss + trust_region_loss
-                trust_region_loss = self.projection.get_trust_region_loss(
-                    p, proj_p)
-
-                trust_region_losses.append(trust_region_loss.item())
-
-                policy_loss = surrogate_loss + self.ent_coef * entropy_loss + trust_region_loss
-                pg_losses.append(policy_loss.item())
-
-                loss = policy_loss + self.vf_coef * value_loss
-
-                # Calculate approximate form of reverse KL Divergence for early stopping
-                # see issue #417: https://github.com/DLR-RM/stable-baselines3/issues/417
-                # and discussion in PR #419: https://github.com/DLR-RM/stable-baselines3/pull/419
-                # and Schulman blog: http://joschu.net/blog/kl-approx.html
-                with th.no_grad():
-                    log_ratio = log_prob - rollout_data.old_log_prob
-                    approx_kl_div = th.mean(
-                        (th.exp(log_ratio) - 1) - log_ratio).cpu().numpy()
-                    approx_kl_divs.append(approx_kl_div)
-
-                if self.target_kl is not None and approx_kl_div > 1.5 * self.target_kl:
-                    continue_training = False
-                    if self.verbose >= 1:
-                        print(
-                            f"Early stopping at step {epoch} due to reaching max kl: {approx_kl_div:.2f}")
-                    break
-
-                # Optimization step
-                self.policy.optimizer.zero_grad()
-                loss.backward()
-                # Clip grad norm
-                th.nn.utils.clip_grad_norm_(
-                    self.policy.parameters(), self.max_grad_norm)
-                self.policy.optimizer.step()
-
-            if not continue_training:
-                break
-
-        self._n_updates += self.n_epochs
-        explained_var = explained_variance(
-            self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten())
-
-        # Logs
-        self.logger.record("train/surrogate_loss", np.mean(surrogate_losses))
-        self.logger.record("train/entropy_loss", np.mean(entropy_losses))
-        self.logger.record("train/trust_region_loss",
-                           np.mean(trust_region_losses))
-        self.logger.record("train/policy_gradient_loss", np.mean(pg_losses))
-        self.logger.record("train/value_loss", np.mean(value_losses))
-        self.logger.record("train/approx_kl", np.mean(approx_kl_divs))
-        self.logger.record("train/clip_fraction", np.mean(clip_fractions))
-        self.logger.record("train/loss", loss.item())
-        self.logger.record("train/explained_variance", explained_var)
-        if hasattr(self.policy, "log_std"):
-            self.logger.record(
-                "train/std", th.exp(self.policy.log_std).mean().item())
-
-        self.logger.record("train/n_updates",
-                           self._n_updates, exclude="tensorboard")
-        self.logger.record("train/clip_range", clip_range)
-        if self.clip_range_vf is not None:
-            self.logger.record("train/clip_range_vf", clip_range_vf)
-
-    def learn(
-        self,
-        total_timesteps: int,
-        callback: MaybeCallback = None,
-        log_interval: int = 1,
-        eval_env: Optional[GymEnv] = None,
-        eval_freq: int = -1,
-        n_eval_episodes: int = 5,
-        tb_log_name: str = "TRL_PG",
-        eval_log_path: Optional[str] = None,
-        reset_num_timesteps: bool = True,
-    ) -> "TRL_PG":
-
-        return super().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,
-            eval_log_path=eval_log_path,
-            reset_num_timesteps=reset_num_timesteps,
-        )
-
-    # This is new compared to PPO.
-    # TRL requires us to also save the original mean and std in our rollouts
-    def collect_rollouts(
-        self,
-        env: VecEnv,
-        callback: BaseCallback,
-        rollout_buffer: RolloutBuffer,
-        n_rollout_steps: int,
-    ) -> bool:
-        """
-        Collect experiences using the current policy and fill a ``RolloutBuffer``.
-        The term rollout here refers to the model-free notion and should not
-        be used with the concept of rollout used in model-based RL or planning.
-        :param env: The training environment
-        :param callback: Callback that will be called at each step
-            (and at the beginning and end of the rollout)
-        :param rollout_buffer: Buffer to fill with rollouts
-        :param n_steps: Number of experiences to collect per environment
-        :return: True if function returned with at least `n_rollout_steps`
-            collected, False if callback terminated rollout prematurely.
-        """
-        assert self._last_obs is not None, "No previous observation was provided"
-        # Switch to eval mode (this affects batch norm / dropout)
-        self.policy.set_training_mode(False)
-
-        n_steps = 0
-        rollout_buffer.reset()
-        # Sample new weights for the state dependent exploration
-        if self.use_sde:
-            self.policy.reset_noise(env.num_envs)
-
-        callback.on_rollout_start()
-
-        while n_steps < n_rollout_steps:
-            if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
-                # Sample a new noise matrix
-                self.policy.reset_noise(env.num_envs)
-
-            with th.no_grad():
-                # Convert to pytorch tensor or to TensorDict
-                obs_tensor = obs_as_tensor(self._last_obs, self.device)
-                actions, values, log_probs = self.policy(obs_tensor)
-                dist = self.policy.get_distribution(obs_tensor).distribution
-                mean, std = dist.mean, dist.stddev
-            actions = actions.cpu().numpy()
-
-            # Rescale and perform action
-            clipped_actions = actions
-            # Clip the actions to avoid out of bound error
-            if isinstance(self.action_space, spaces.Box):
-                clipped_actions = np.clip(
-                    actions, self.action_space.low, self.action_space.high)
-
-            new_obs, rewards, dones, infos = env.step(clipped_actions)
-
-            self.num_timesteps += env.num_envs
-
-            # Give access to local variables
-            callback.update_locals(locals())
-            if callback.on_step() is False:
-                return False
-
-            self._update_info_buffer(infos)
-            n_steps += 1
-
-            if isinstance(self.action_space, spaces.Discrete):
-                # Reshape in case of discrete action
-                actions = actions.reshape(-1, 1)
-
-            # Handle timeout by bootstraping with value function
-            # see GitHub issue #633
-            for idx, done in enumerate(dones):
-                if (
-                    done
-                    and infos[idx].get("terminal_observation") is not None
-                    and infos[idx].get("TimeLimit.truncated", False)
-                ):
-                    terminal_obs = self.policy.obs_to_tensor(
-                        infos[idx]["terminal_observation"])[0]
-                    with th.no_grad():
-                        terminal_value = self.policy.predict_values(terminal_obs)[
-                            0]
-                    rewards[idx] += self.gamma * terminal_value
-
-            rollout_buffer.add(self._last_obs, actions, rewards,
-                               self._last_episode_starts, values, log_probs, mean, std)
-            self._last_obs = new_obs
-            self._last_episode_starts = dones
-
-        with th.no_grad():
-            # Compute value for the last timestep
-            values = self.policy.predict_values(
-                obs_as_tensor(new_obs, self.device))
-
-        rollout_buffer.compute_returns_and_advantage(
-            last_values=values, dones=dones)
-
-        callback.on_rollout_end()
-
-        return True

sb3_trl/trl_sac/__init__.py

@@ -1,2 +0,0 @@
-from sb3_trl.trl_sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
-from sb3_trl.trl_sac.trl_sac import TRL_SAC

sb3_trl/trl_sac/policies.py

@@ -1,516 +0,0 @@
-import warnings
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
-
-import gym
-import torch as th
-from torch import nn
-
-from stable_baselines3.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
-from stable_baselines3.common.policies import BasePolicy, ContinuousCritic
-from stable_baselines3.common.preprocessing import get_action_dim
-from stable_baselines3.common.torch_layers import (
-    BaseFeaturesExtractor,
-    CombinedExtractor,
-    FlattenExtractor,
-    NatureCNN,
-    create_mlp,
-    get_actor_critic_arch,
-)
-from stable_baselines3.common.type_aliases import Schedule
-
-# CAP the standard deviation of the actor
-LOG_STD_MAX = 2
-LOG_STD_MIN = -20
-
-
-class Actor(BasePolicy):
-    """
-    Actor network (policy) for SAC.
-
-    :param observation_space: Obervation space
-    :param action_space: Action space
-    :param net_arch: Network architecture
-    :param features_extractor: Network to extract features
-        (a CNN when using images, a nn.Flatten() layer otherwise)
-    :param features_dim: Number of features
-    :param activation_fn: Activation function
-    :param use_sde: Whether to use State Dependent Exploration or not
-    :param log_std_init: Initial value for the log standard deviation
-    :param full_std: Whether to use (n_features x n_actions) parameters
-        for the std instead of only (n_features,) when using gSDE.
-    :param sde_net_arch: Network architecture for extracting features
-        when using gSDE. If None, the latent features from the policy will be used.
-        Pass an empty list to use the states as features.
-    :param use_expln: Use ``expln()`` function instead of ``exp()`` when using gSDE to ensure
-        a positive standard deviation (cf paper). It allows to keep variance
-        above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
-    :param clip_mean: Clip the mean output when using gSDE to avoid numerical instability.
-    :param normalize_images: Whether to normalize images or not,
-         dividing by 255.0 (True by default)
-    """
-
-    def __init__(
-        self,
-        observation_space: gym.spaces.Space,
-        action_space: gym.spaces.Space,
-        net_arch: List[int],
-        features_extractor: nn.Module,
-        features_dim: int,
-        activation_fn: Type[nn.Module] = nn.ReLU,
-        use_sde: bool = False,
-        log_std_init: float = -3,
-        full_std: bool = True,
-        sde_net_arch: Optional[List[int]] = None,
-        use_expln: bool = False,
-        clip_mean: float = 2.0,
-        normalize_images: bool = True,
-    ):
-        super().__init__(
-            observation_space,
-            action_space,
-            features_extractor=features_extractor,
-            normalize_images=normalize_images,
-            squash_output=True,
-        )
-
-        # Save arguments to re-create object at loading
-        self.use_sde = use_sde
-        self.sde_features_extractor = None
-        self.net_arch = net_arch
-        self.features_dim = features_dim
-        self.activation_fn = activation_fn
-        self.log_std_init = log_std_init
-        self.sde_net_arch = sde_net_arch
-        self.use_expln = use_expln
-        self.full_std = full_std
-        self.clip_mean = clip_mean
-
-        if sde_net_arch is not None:
-            warnings.warn("sde_net_arch is deprecated and will be removed in SB3 v2.4.0.", DeprecationWarning)
-
-        action_dim = get_action_dim(self.action_space)
-        latent_pi_net = create_mlp(features_dim, -1, net_arch, activation_fn)
-        self.latent_pi = nn.Sequential(*latent_pi_net)
-        last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim
-
-        if self.use_sde:
-            self.action_dist = StateDependentNoiseDistribution(
-                action_dim, full_std=full_std, use_expln=use_expln, learn_features=True, squash_output=True
-            )
-            self.mu, self.log_std = self.action_dist.proba_distribution_net(
-                latent_dim=last_layer_dim, latent_sde_dim=last_layer_dim, log_std_init=log_std_init
-            )
-            # Avoid numerical issues by limiting the mean of the Gaussian
-            # to be in [-clip_mean, clip_mean]
-            if clip_mean > 0.0:
-                self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-clip_mean, max_val=clip_mean))
-        else:
-            self.action_dist = SquashedDiagGaussianDistribution(action_dim)
-            self.mu = nn.Linear(last_layer_dim, action_dim)
-            self.log_std = nn.Linear(last_layer_dim, action_dim)
-
-    def _get_constructor_parameters(self) -> Dict[str, Any]:
-        data = super()._get_constructor_parameters()
-
-        data.update(
-            dict(
-                net_arch=self.net_arch,
-                features_dim=self.features_dim,
-                activation_fn=self.activation_fn,
-                use_sde=self.use_sde,
-                log_std_init=self.log_std_init,
-                full_std=self.full_std,
-                use_expln=self.use_expln,
-                features_extractor=self.features_extractor,
-                clip_mean=self.clip_mean,
-            )
-        )
-        return data
-
-    def get_std(self) -> th.Tensor:
-        """
-        Retrieve the standard deviation of the action distribution.
-        Only useful when using gSDE.
-        It corresponds to ``th.exp(log_std)`` in the normal case,
-        but is slightly different when using ``expln`` function
-        (cf StateDependentNoiseDistribution doc).
-
-        :return:
-        """
-        msg = "get_std() is only available when using gSDE"
-        assert isinstance(self.action_dist, StateDependentNoiseDistribution), msg
-        return self.action_dist.get_std(self.log_std)
-
-    def reset_noise(self, batch_size: int = 1) -> None:
-        """
-        Sample new weights for the exploration matrix, when using gSDE.
-
-        :param batch_size:
-        """
-        msg = "reset_noise() is only available when using gSDE"
-        assert isinstance(self.action_dist, StateDependentNoiseDistribution), msg
-        self.action_dist.sample_weights(self.log_std, batch_size=batch_size)
-
-    def get_action_dist_params(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor, Dict[str, th.Tensor]]:
-        """
-        Get the parameters for the action distribution.
-
-        :param obs:
-        :return:
-            Mean, standard deviation and optional keyword arguments.
-        """
-        features = self.extract_features(obs)
-        latent_pi = self.latent_pi(features)
-        mean_actions = self.mu(latent_pi)
-
-        if self.use_sde:
-            return mean_actions, self.log_std, dict(latent_sde=latent_pi)
-        # Unstructured exploration (Original implementation)
-        log_std = self.log_std(latent_pi)
-        # 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, obs: th.Tensor, deterministic: bool = False) -> th.Tensor:
-        mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
-        # Note: the action is squashed
-        return self.action_dist.actions_from_params(mean_actions, log_std, deterministic=deterministic, **kwargs)
-
-    def action_log_prob(self, obs: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
-        mean_actions, log_std, kwargs = self.get_action_dist_params(obs)
-        # return action and associated log prob
-        return self.action_dist.log_prob_from_params(mean_actions, log_std, **kwargs)
-
-    def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
-        return self(observation, deterministic)
-
-
-class SACPolicy(BasePolicy):
-    """
-    Policy class (with both actor and critic) for SAC.
-
-    :param observation_space: Observation space
-    :param action_space: Action space
-    :param lr_schedule: Learning rate schedule (could be constant)
-    :param net_arch: The specification of the policy and value networks.
-    :param activation_fn: Activation function
-    :param use_sde: Whether to use State Dependent Exploration or not
-    :param log_std_init: Initial value for the log standard deviation
-    :param sde_net_arch: Network architecture for extracting features
-        when using gSDE. If None, the latent features from the policy will be used.
-        Pass an empty list to use the states as features.
-    :param use_expln: Use ``expln()`` function instead of ``exp()`` when using gSDE to ensure
-        a positive standard deviation (cf paper). It allows to keep variance
-        above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
-    :param clip_mean: Clip the mean output when using gSDE to avoid numerical instability.
-    :param features_extractor_class: Features extractor to use.
-    :param features_extractor_kwargs: Keyword arguments
-        to pass to the features extractor.
-    :param normalize_images: Whether to normalize images or not,
-         dividing by 255.0 (True by default)
-    :param optimizer_class: The optimizer to use,
-        ``th.optim.Adam`` by default
-    :param optimizer_kwargs: Additional keyword arguments,
-        excluding the learning rate, to pass to the optimizer
-    :param n_critics: Number of critic networks to create.
-    :param share_features_extractor: Whether to share or not the features extractor
-        between the actor and the critic (this saves computation time)
-    """
-
-    def __init__(
-        self,
-        observation_space: gym.spaces.Space,
-        action_space: gym.spaces.Space,
-        lr_schedule: Schedule,
-        net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
-        activation_fn: Type[nn.Module] = nn.ReLU,
-        use_sde: bool = False,
-        log_std_init: float = -3,
-        sde_net_arch: Optional[List[int]] = None,
-        use_expln: bool = False,
-        clip_mean: float = 2.0,
-        features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
-        features_extractor_kwargs: Optional[Dict[str, Any]] = None,
-        normalize_images: bool = True,
-        optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
-        n_critics: int = 2,
-        share_features_extractor: bool = True,
-    ):
-        super().__init__(
-            observation_space,
-            action_space,
-            features_extractor_class,
-            features_extractor_kwargs,
-            optimizer_class=optimizer_class,
-            optimizer_kwargs=optimizer_kwargs,
-            squash_output=True,
-        )
-
-        if net_arch is None:
-            if features_extractor_class == NatureCNN:
-                net_arch = []
-            else:
-                net_arch = [256, 256]
-
-        actor_arch, critic_arch = get_actor_critic_arch(net_arch)
-
-        self.net_arch = net_arch
-        self.activation_fn = activation_fn
-        self.net_args = {
-            "observation_space": self.observation_space,
-            "action_space": self.action_space,
-            "net_arch": actor_arch,
-            "activation_fn": self.activation_fn,
-            "normalize_images": normalize_images,
-        }
-        self.actor_kwargs = self.net_args.copy()
-
-        if sde_net_arch is not None:
-            warnings.warn("sde_net_arch is deprecated and will be removed in SB3 v2.4.0.", DeprecationWarning)
-
-        sde_kwargs = {
-            "use_sde": use_sde,
-            "log_std_init": log_std_init,
-            "use_expln": use_expln,
-            "clip_mean": clip_mean,
-        }
-        self.actor_kwargs.update(sde_kwargs)
-        self.critic_kwargs = self.net_args.copy()
-        self.critic_kwargs.update(
-            {
-                "n_critics": n_critics,
-                "net_arch": critic_arch,
-                "share_features_extractor": share_features_extractor,
-            }
-        )
-
-        self.actor, self.actor_target = None, None
-        self.critic, self.critic_target = None, None
-        self.share_features_extractor = share_features_extractor
-
-        self._build(lr_schedule)
-
-    def _build(self, lr_schedule: Schedule) -> None:
-        self.actor = self.make_actor()
-        self.actor.optimizer = self.optimizer_class(self.actor.parameters(), lr=lr_schedule(1), **self.optimizer_kwargs)
-
-        if self.share_features_extractor:
-            self.critic = self.make_critic(features_extractor=self.actor.features_extractor)
-            # Do not optimize the shared features extractor with the critic loss
-            # otherwise, there are gradient computation issues
-            critic_parameters = [param for name, param in self.critic.named_parameters() if "features_extractor" not in name]
-        else:
-            # Create a separate features extractor for the critic
-            # this requires more memory and computation
-            self.critic = self.make_critic(features_extractor=None)
-            critic_parameters = self.critic.parameters()
-
-        # Critic target should not share the features extractor with critic
-        self.critic_target = self.make_critic(features_extractor=None)
-        self.critic_target.load_state_dict(self.critic.state_dict())
-
-        self.critic.optimizer = self.optimizer_class(critic_parameters, lr=lr_schedule(1), **self.optimizer_kwargs)
-
-        # Target networks should always be in eval mode
-        self.critic_target.set_training_mode(False)
-
-    def _get_constructor_parameters(self) -> Dict[str, Any]:
-        data = super()._get_constructor_parameters()
-
-        data.update(
-            dict(
-                net_arch=self.net_arch,
-                activation_fn=self.net_args["activation_fn"],
-                use_sde=self.actor_kwargs["use_sde"],
-                log_std_init=self.actor_kwargs["log_std_init"],
-                use_expln=self.actor_kwargs["use_expln"],
-                clip_mean=self.actor_kwargs["clip_mean"],
-                n_critics=self.critic_kwargs["n_critics"],
-                lr_schedule=self._dummy_schedule,  # dummy lr schedule, not needed for loading policy alone
-                optimizer_class=self.optimizer_class,
-                optimizer_kwargs=self.optimizer_kwargs,
-                features_extractor_class=self.features_extractor_class,
-                features_extractor_kwargs=self.features_extractor_kwargs,
-            )
-        )
-        return data
-
-    def reset_noise(self, batch_size: int = 1) -> None:
-        """
-        Sample new weights for the exploration matrix, when using gSDE.
-
-        :param batch_size:
-        """
-        self.actor.reset_noise(batch_size=batch_size)
-
-    def make_actor(self, features_extractor: Optional[BaseFeaturesExtractor] = None) -> Actor:
-        actor_kwargs = self._update_features_extractor(self.actor_kwargs, features_extractor)
-        return Actor(**actor_kwargs).to(self.device)
-
-    def make_critic(self, features_extractor: Optional[BaseFeaturesExtractor] = None) -> ContinuousCritic:
-        critic_kwargs = self._update_features_extractor(self.critic_kwargs, features_extractor)
-        return ContinuousCritic(**critic_kwargs).to(self.device)
-
-    def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor:
-        return self._predict(obs, deterministic=deterministic)
-
-    def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
-        return self.actor(observation, deterministic)
-
-    def set_training_mode(self, mode: bool) -> None:
-        """
-        Put the policy in either training or evaluation mode.
-
-        This affects certain modules, such as batch normalisation and dropout.
-
-        :param mode: if true, set to training mode, else set to evaluation mode
-        """
-        self.actor.set_training_mode(mode)
-        self.critic.set_training_mode(mode)
-        self.training = mode
-
-
-MlpPolicy = SACPolicy
-
-
-class CnnPolicy(SACPolicy):
-    """
-    Policy class (with both actor and critic) for SAC.
-
-    :param observation_space: Observation space
-    :param action_space: Action space
-    :param lr_schedule: Learning rate schedule (could be constant)
-    :param net_arch: The specification of the policy and value networks.
-    :param activation_fn: Activation function
-    :param use_sde: Whether to use State Dependent Exploration or not
-    :param log_std_init: Initial value for the log standard deviation
-    :param sde_net_arch: Network architecture for extracting features
-        when using gSDE. If None, the latent features from the policy will be used.
-        Pass an empty list to use the states as features.
-    :param use_expln: Use ``expln()`` function instead of ``exp()`` when using gSDE to ensure
-        a positive standard deviation (cf paper). It allows to keep variance
-        above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
-    :param clip_mean: Clip the mean output when using gSDE to avoid numerical instability.
-    :param features_extractor_class: Features extractor to use.
-    :param normalize_images: Whether to normalize images or not,
-         dividing by 255.0 (True by default)
-    :param optimizer_class: The optimizer to use,
-        ``th.optim.Adam`` by default
-    :param optimizer_kwargs: Additional keyword arguments,
-        excluding the learning rate, to pass to the optimizer
-    :param n_critics: Number of critic networks to create.
-    :param share_features_extractor: Whether to share or not the features extractor
-        between the actor and the critic (this saves computation time)
-    """
-
-    def __init__(
-        self,
-        observation_space: gym.spaces.Space,
-        action_space: gym.spaces.Space,
-        lr_schedule: Schedule,
-        net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
-        activation_fn: Type[nn.Module] = nn.ReLU,
-        use_sde: bool = False,
-        log_std_init: float = -3,
-        sde_net_arch: Optional[List[int]] = None,
-        use_expln: bool = False,
-        clip_mean: float = 2.0,
-        features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN,
-        features_extractor_kwargs: Optional[Dict[str, Any]] = None,
-        normalize_images: bool = True,
-        optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
-        n_critics: int = 2,
-        share_features_extractor: bool = True,
-    ):
-        super().__init__(
-            observation_space,
-            action_space,
-            lr_schedule,
-            net_arch,
-            activation_fn,
-            use_sde,
-            log_std_init,
-            sde_net_arch,
-            use_expln,
-            clip_mean,
-            features_extractor_class,
-            features_extractor_kwargs,
-            normalize_images,
-            optimizer_class,
-            optimizer_kwargs,
-            n_critics,
-            share_features_extractor,
-        )
-
-
-class MultiInputPolicy(SACPolicy):
-    """
-    Policy class (with both actor and critic) for SAC.
-
-    :param observation_space: Observation space
-    :param action_space: Action space
-    :param lr_schedule: Learning rate schedule (could be constant)
-    :param net_arch: The specification of the policy and value networks.
-    :param activation_fn: Activation function
-    :param use_sde: Whether to use State Dependent Exploration or not
-    :param log_std_init: Initial value for the log standard deviation
-    :param sde_net_arch: Network architecture for extracting features
-        when using gSDE. If None, the latent features from the policy will be used.
-        Pass an empty list to use the states as features.
-    :param use_expln: Use ``expln()`` function instead of ``exp()`` when using gSDE to ensure
-        a positive standard deviation (cf paper). It allows to keep variance
-        above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
-    :param clip_mean: Clip the mean output when using gSDE to avoid numerical instability.
-    :param features_extractor_class: Features extractor to use.
-    :param normalize_images: Whether to normalize images or not,
-         dividing by 255.0 (True by default)
-    :param optimizer_class: The optimizer to use,
-        ``th.optim.Adam`` by default
-    :param optimizer_kwargs: Additional keyword arguments,
-        excluding the learning rate, to pass to the optimizer
-    :param n_critics: Number of critic networks to create.
-    :param share_features_extractor: Whether to share or not the features extractor
-        between the actor and the critic (this saves computation time)
-    """
-
-    def __init__(
-        self,
-        observation_space: gym.spaces.Space,
-        action_space: gym.spaces.Space,
-        lr_schedule: Schedule,
-        net_arch: Optional[Union[List[int], Dict[str, List[int]]]] = None,
-        activation_fn: Type[nn.Module] = nn.ReLU,
-        use_sde: bool = False,
-        log_std_init: float = -3,
-        sde_net_arch: Optional[List[int]] = None,
-        use_expln: bool = False,
-        clip_mean: float = 2.0,
-        features_extractor_class: Type[BaseFeaturesExtractor] = CombinedExtractor,
-        features_extractor_kwargs: Optional[Dict[str, Any]] = None,
-        normalize_images: bool = True,
-        optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam,
-        optimizer_kwargs: Optional[Dict[str, Any]] = None,
-        n_critics: int = 2,
-        share_features_extractor: bool = True,
-    ):
-        super().__init__(
-            observation_space,
-            action_space,
-            lr_schedule,
-            net_arch,
-            activation_fn,
-            use_sde,
-            log_std_init,
-            sde_net_arch,
-            use_expln,
-            clip_mean,
-            features_extractor_class,
-            features_extractor_kwargs,
-            normalize_images,
-            optimizer_class,
-            optimizer_kwargs,
-            n_critics,
-            share_features_extractor,
-        )

sb3_trl/trl_sac/trl_sac.py

@@ -1,324 +0,0 @@
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
-
-import gym
-import numpy as np
-import torch as th
-from torch.nn import functional as F
-
-from stable_baselines3.common.buffers import ReplayBuffer
-from stable_baselines3.common.noise import ActionNoise
-from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
-from stable_baselines3.common.policies import BasePolicy
-from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule
-from stable_baselines3.common.utils import polyak_update
-from stable_baselines3.sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy, SACPolicy
-
-
-class TRL_SAC(OffPolicyAlgorithm):
-    """
-    Trust Region Layers (TRL) based on SAC (Soft Actor Critic)
-    This implementation is almost a 1:1-copy of the sb3-code for SAC.
-    Only minor changes have been made to implement Differential Trust Region Layers
-
-    Description from original SAC implementation:
-    Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor,
-    This implementation borrows code from original implementation (https://github.com/haarnoja/sac)
-    from OpenAI Spinning Up (https://github.com/openai/spinningup), from the softlearning repo
-    (https://github.com/rail-berkeley/softlearning/)
-    and from Stable Baselines (https://github.com/hill-a/stable-baselines)
-    Paper: https://arxiv.org/abs/1801.01290
-    Introduction to SAC: https://spinningup.openai.com/en/latest/algorithms/sac.html
-
-    Note: we use double q target and not value target as discussed
-    in https://github.com/hill-a/stable-baselines/issues/270
-
-    :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
-    :param env: The environment to learn from (if registered in Gym, can be str)
-    :param learning_rate: learning rate for adam optimizer,
-        the same learning rate will be used for all networks (Q-Values, Actor and Value function)
-        it can be a function of the current progress remaining (from 1 to 0)
-    :param buffer_size: size of the replay buffer
-    :param learning_starts: how many steps of the model to collect transitions for before learning starts
-    :param batch_size: Minibatch size for each gradient update
-    :param tau: the soft update coefficient ("Polyak update", between 0 and 1)
-    :param gamma: the discount factor
-    :param train_freq: Update the model every ``train_freq`` steps. Alternatively pass a tuple of frequency and unit
-        like ``(5, "step")`` or ``(2, "episode")``.
-    :param gradient_steps: How many gradient steps to do after each rollout (see ``train_freq``)
-        Set to ``-1`` means to do as many gradient steps as steps done in the environment
-        during the rollout.
-    :param action_noise: the action noise type (None by default), this can help
-        for hard exploration problem. Cf common.noise for the different action noise type.
-    :param replay_buffer_class: Replay buffer class to use (for instance ``HerReplayBuffer``).
-        If ``None``, it will be automatically selected.
-    :param replay_buffer_kwargs: Keyword arguments to pass to the replay buffer on creation.
-    :param optimize_memory_usage: Enable a memory efficient variant of the replay buffer
-        at a cost of more complexity.
-        See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195
-    :param ent_coef: Entropy regularization coefficient. (Equivalent to
-        inverse of reward scale in the original SAC paper.)  Controlling exploration/exploitation trade-off.
-        Set it to 'auto' to learn it automatically (and 'auto_0.1' for using 0.1 as initial value)
-    :param target_update_interval: update the target network every ``target_network_update_freq``
-        gradient steps.
-    :param target_entropy: target entropy when learning ``ent_coef`` (``ent_coef = 'auto'``)
-    :param use_sde: Whether to use generalized State Dependent Exploration (gSDE)
-        instead of action noise exploration (default: False)
-    :param sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
-        Default: -1 (only sample at the beginning of the rollout)
-    :param use_sde_at_warmup: Whether to use gSDE instead of uniform sampling
-        during the warm up phase (before learning starts)
-    :param create_eval_env: Whether to create a second environment that will be
-        used for evaluating the agent periodically. (Only available when passing string for the environment)
-    :param policy_kwargs: additional arguments to be passed to the policy on creation
-    :param verbose: the verbosity level: 0 no output, 1 info, 2 debug
-    :param seed: Seed for the pseudo random generators
-    :param device: Device (cpu, cuda, ...) on which the code should be run.
-        Setting it to auto, the code will be run on the GPU if possible.
-    :param _init_setup_model: Whether or not to build the network at the creation of the instance
-    """
-
-    policy_aliases: Dict[str, Type[BasePolicy]] = {
-        "MlpPolicy": MlpPolicy,
-        "CnnPolicy": CnnPolicy,
-        "MultiInputPolicy": MultiInputPolicy,
-    }
-
-    def __init__(
-        self,
-        policy: Union[str, Type[SACPolicy]],
-        env: Union[GymEnv, str],
-        learning_rate: Union[float, Schedule] = 3e-4,
-        buffer_size: int = 1_000_000,  # 1e6
-        learning_starts: int = 100,
-        batch_size: int = 256,
-        tau: float = 0.005,
-        gamma: float = 0.99,
-        train_freq: Union[int, Tuple[int, str]] = 1,
-        gradient_steps: int = 1,
-        action_noise: Optional[ActionNoise] = None,
-        replay_buffer_class: Optional[ReplayBuffer] = None,
-        replay_buffer_kwargs: Optional[Dict[str, Any]] = None,
-        optimize_memory_usage: bool = False,
-        ent_coef: Union[str, float] = "auto",
-        target_update_interval: int = 1,
-        target_entropy: Union[str, float] = "auto",
-        use_sde: bool = False,
-        sde_sample_freq: int = -1,
-        use_sde_at_warmup: bool = False,
-        tensorboard_log: Optional[str] = None,
-        create_eval_env: bool = False,
-        policy_kwargs: Optional[Dict[str, Any]] = None,
-        verbose: int = 0,
-        seed: Optional[int] = None,
-        device: Union[th.device, str] = "auto",
-        _init_setup_model: bool = True,
-    ):
-
-        super().__init__(
-            policy,
-            env,
-            learning_rate,
-            buffer_size,
-            learning_starts,
-            batch_size,
-            tau,
-            gamma,
-            train_freq,
-            gradient_steps,
-            action_noise,
-            replay_buffer_class=replay_buffer_class,
-            replay_buffer_kwargs=replay_buffer_kwargs,
-            policy_kwargs=policy_kwargs,
-            tensorboard_log=tensorboard_log,
-            verbose=verbose,
-            device=device,
-            create_eval_env=create_eval_env,
-            seed=seed,
-            use_sde=use_sde,
-            sde_sample_freq=sde_sample_freq,
-            use_sde_at_warmup=use_sde_at_warmup,
-            optimize_memory_usage=optimize_memory_usage,
-            supported_action_spaces=(gym.spaces.Box),
-            support_multi_env=True,
-        )
-
-        self.target_entropy = target_entropy
-        self.log_ent_coef = None  # type: Optional[th.Tensor]
-        # Entropy coefficient / Entropy temperature
-        # Inverse of the reward scale
-        self.ent_coef = ent_coef
-        self.target_update_interval = target_update_interval
-        self.ent_coef_optimizer = None
-
-        if _init_setup_model:
-            self._setup_model()
-
-    def _setup_model(self) -> None:
-        super()._setup_model()
-        self._create_aliases()
-        # Target entropy is used when learning the entropy coefficient
-        if self.target_entropy == "auto":
-            # automatically set target entropy if needed
-            self.target_entropy = -np.prod(self.env.action_space.shape).astype(np.float32)
-        else:
-            # Force conversion
-            # this will also throw an error for unexpected string
-            self.target_entropy = float(self.target_entropy)
-
-        # The entropy coefficient or entropy can be learned automatically
-        # see Automating Entropy Adjustment for Maximum Entropy RL section
-        # of https://arxiv.org/abs/1812.05905
-        if isinstance(self.ent_coef, str) and self.ent_coef.startswith("auto"):
-            # Default initial value of ent_coef when learned
-            init_value = 1.0
-            if "_" in self.ent_coef:
-                init_value = float(self.ent_coef.split("_")[1])
-                assert init_value > 0.0, "The initial value of ent_coef must be greater than 0"
-
-            # Note: we optimize the log of the entropy coeff which is slightly different from the paper
-            # as discussed in https://github.com/rail-berkeley/softlearning/issues/37
-            self.log_ent_coef = th.log(th.ones(1, device=self.device) * init_value).requires_grad_(True)
-            self.ent_coef_optimizer = th.optim.Adam([self.log_ent_coef], lr=self.lr_schedule(1))
-        else:
-            # Force conversion to float
-            # this will throw an error if a malformed string (different from 'auto')
-            # is passed
-            self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(self.device)
-
-    def _create_aliases(self) -> None:
-        self.actor = self.policy.actor
-        self.critic = self.policy.critic
-        self.critic_target = self.policy.critic_target
-
-    def train(self, gradient_steps: int, batch_size: int = 64) -> None:
-        # Switch to train mode (this affects batch norm / dropout)
-        self.policy.set_training_mode(True)
-        # Update optimizers learning rate
-        optimizers = [self.actor.optimizer, self.critic.optimizer]
-        if self.ent_coef_optimizer is not None:
-            optimizers += [self.ent_coef_optimizer]
-
-        # Update learning rate according to lr schedule
-        self._update_learning_rate(optimizers)
-
-        ent_coef_losses, ent_coefs = [], []
-        actor_losses, critic_losses = [], []
-
-        for gradient_step in range(gradient_steps):
-            # Sample replay buffer
-            replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env)
-
-            # We need to sample because `log_std` may have changed between two gradient steps
-            if self.use_sde:
-                self.actor.reset_noise()
-
-            # Action by the current actor for the sampled state
-            actions_pi, log_prob = self.actor.action_log_prob(replay_data.observations)
-            log_prob = log_prob.reshape(-1, 1)
-
-            ent_coef_loss = None
-            if self.ent_coef_optimizer is not None:
-                # Important: detach the variable from the graph
-                # so we don't change it with other losses
-                # see https://github.com/rail-berkeley/softlearning/issues/60
-                ent_coef = th.exp(self.log_ent_coef.detach())
-                ent_coef_loss = -(self.log_ent_coef * (log_prob + self.target_entropy).detach()).mean()
-                ent_coef_losses.append(ent_coef_loss.item())
-            else:
-                ent_coef = self.ent_coef_tensor
-
-            ent_coefs.append(ent_coef.item())
-
-            # Optimize entropy coefficient, also called
-            # entropy temperature or alpha in the paper
-            if ent_coef_loss is not None:
-                self.ent_coef_optimizer.zero_grad()
-                ent_coef_loss.backward()
-                self.ent_coef_optimizer.step()
-
-            with th.no_grad():
-                # Select action according to policy
-                next_actions, next_log_prob = self.actor.action_log_prob(replay_data.next_observations)
-                # Compute the next Q values: min over all critics targets
-                next_q_values = th.cat(self.critic_target(replay_data.next_observations, next_actions), dim=1)
-                next_q_values, _ = th.min(next_q_values, dim=1, keepdim=True)
-                # add entropy term
-                next_q_values = next_q_values - ent_coef * next_log_prob.reshape(-1, 1)
-                # td error + entropy term
-                target_q_values = replay_data.rewards + (1 - replay_data.dones) * self.gamma * next_q_values
-
-            # Get current Q-values estimates for each critic network
-            # using action from the replay buffer
-            current_q_values = self.critic(replay_data.observations, replay_data.actions)
-
-            # Compute critic loss
-            critic_loss = 0.5 * sum(F.mse_loss(current_q, target_q_values) for current_q in current_q_values)
-            critic_losses.append(critic_loss.item())
-
-            # Optimize the critic
-            self.critic.optimizer.zero_grad()
-            critic_loss.backward()
-            self.critic.optimizer.step()
-
-            # Compute actor loss
-            # Alternative: actor_loss = th.mean(log_prob - qf1_pi)
-            # Mean over all critic networks
-            q_values_pi = th.cat(self.critic(replay_data.observations, actions_pi), dim=1)
-            min_qf_pi, _ = th.min(q_values_pi, dim=1, keepdim=True)
-            actor_loss = (ent_coef * log_prob - min_qf_pi).mean()
-            actor_losses.append(actor_loss.item())
-
-            # Optimize the actor
-            self.actor.optimizer.zero_grad()
-            actor_loss.backward()
-            self.actor.optimizer.step()
-
-            # Update target networks
-            if gradient_step % self.target_update_interval == 0:
-                polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau)
-
-        self._n_updates += gradient_steps
-
-        self.logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
-        self.logger.record("train/ent_coef", np.mean(ent_coefs))
-        self.logger.record("train/actor_loss", np.mean(actor_losses))
-        self.logger.record("train/critic_loss", np.mean(critic_losses))
-        if len(ent_coef_losses) > 0:
-            self.logger.record("train/ent_coef_loss", np.mean(ent_coef_losses))
-
-    def learn(
-        self,
-        total_timesteps: int,
-        callback: MaybeCallback = None,
-        log_interval: int = 4,
-        eval_env: Optional[GymEnv] = None,
-        eval_freq: int = -1,
-        n_eval_episodes: int = 5,
-        tb_log_name: str = "SAC",
-        eval_log_path: Optional[str] = None,
-        reset_num_timesteps: bool = True,
-    ) -> OffPolicyAlgorithm:
-
-        return super().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,
-            eval_log_path=eval_log_path,
-            reset_num_timesteps=reset_num_timesteps,
-        )
-
-    def _excluded_save_params(self) -> List[str]:
-        return super()._excluded_save_params() + ["actor", "critic", "critic_target"]
-
-    def _get_torch_save_params(self) -> Tuple[List[str], List[str]]:
-        state_dicts = ["policy", "actor.optimizer", "critic.optimizer"]
-        if self.ent_coef_optimizer is not None:
-            saved_pytorch_variables = ["log_ent_coef"]
-            state_dicts.append("ent_coef_optimizer")
-        else:
-            saved_pytorch_variables = ["ent_coef_tensor"]
-        return state_dicts, saved_pytorch_variables

test.py

@@ -10,16 +10,16 @@ from stable_baselines3 import SAC, PPO, A2C
 from stable_baselines3.common.evaluation import evaluate_policy
 from stable_baselines3.common.policies import ActorCriticCnnPolicy, ActorCriticPolicy, MultiInputActorCriticPolicy
 
-from sb3_trl.trl_pg import TRL_PG
+from metastable_baselines.trl_pg import TRL_PG
 import columbus
 
 #root_path = os.getcwd()
 root_path = '.'
 
 
-def main(env_name='ColumbusCandyland_Aux10-v0', timesteps=200_000, showRes=True, saveModel=True, n_eval_episodes=0):
+def main(env_name='ColumbusCandyland_Aux10-v0', timesteps=10_000_000, showRes=True, saveModel=True, n_eval_episodes=0):
     env = gym.make(env_name)
-    use_sde = False
+    use_sde = True
     ppo = PPO(
         "MlpPolicy",
         env,
@@ -100,7 +100,7 @@ def testModel(model, timesteps, showRes=False, saveModel=False, n_eval_episodes=
 
 
 if __name__ == '__main__':
-    main('LunarLanderContinuous-v2')
+    # main('LunarLanderContinuous-v2')
     # main('ColumbusJustState-v0')
     # main('ColumbusStateWithBarriers-v0')
-    #main('ColumbusEasierObstacles-v0')
+    main('ColumbusEasierObstacles-v0')