diff --git a/projections_orig/__init__.py b/projections_orig/__init__.py
deleted file mode 100644
index a578185..0000000
--- a/projections_orig/__init__.py
+++ /dev/null
@@ -1,15 +0,0 @@
-# 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 .
diff --git a/projections_orig/base_projection_layer.py b/projections_orig/base_projection_layer.py
deleted file mode 100644
index 3a881af..0000000
--- a/projections_orig/base_projection_layer.py
+++ /dev/null
@@ -1,374 +0,0 @@
-# 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 .
-
-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()}
diff --git a/projections_orig/frob_projection_layer.py b/projections_orig/frob_projection_layer.py
deleted file mode 100644
index 8d338ce..0000000
--- a/projections_orig/frob_projection_layer.py
+++ /dev/null
@@ -1,97 +0,0 @@
-# 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 .
-
-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
diff --git a/projections_orig/kl_projection_layer.py b/projections_orig/kl_projection_layer.py
deleted file mode 100644
index ca5acd5..0000000
--- a/projections_orig/kl_projection_layer.py
+++ /dev/null
@@ -1,101 +0,0 @@
-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
diff --git a/projections_orig/papi_projection.py b/projections_orig/papi_projection.py
deleted file mode 100644
index b52db75..0000000
--- a/projections_orig/papi_projection.py
+++ /dev/null
@@ -1,233 +0,0 @@
-# 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 .
-
-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())
diff --git a/projections_orig/projection_factory.py b/projections_orig/projection_factory.py
deleted file mode 100644
index 9c38275..0000000
--- a/projections_orig/projection_factory.py
+++ /dev/null
@@ -1,54 +0,0 @@
-# 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 .
-
-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'.")
diff --git a/projections_orig/w2_projection_layer.py b/projections_orig/w2_projection_layer.py
deleted file mode 100644
index bce87a3..0000000
--- a/projections_orig/w2_projection_layer.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# 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 .
-
-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)
\ No newline at end of file
diff --git a/trl_pg/__init__.py b/trl_pg/__init__.py
deleted file mode 100644
index df30cfc..0000000
--- a/trl_pg/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from sb3.trl_pg.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
-from sb3.trl_pg.trl_pg import TRL_PG
diff --git a/trl_pg/policies.py b/trl_pg/policies.py
deleted file mode 100644
index 8b784eb..0000000
--- a/trl_pg/policies.py
+++ /dev/null
@@ -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
diff --git a/trl_pg/trl_pg.py b/trl_pg/trl_pg.py
deleted file mode 100644
index 09c9c58..0000000
--- a/trl_pg/trl_pg.py
+++ /dev/null
@@ -1,338 +0,0 @@
-import warnings
-from typing import Any, Dict, Optional, Type, Union
-
-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
-
-
-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
- see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213)
- 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 _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 = None,
-
- _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
-
- 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)
-
- 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):
- 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)
-
- values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
- 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(policy_loss_1, policy_loss_2).mean()
-
- surrogate_losses.append(surrogate_loss.item())
-
- # Logging
- pg_losses.append(policy_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()#TODO: params
-
- trust_region_losses.append(trust_region_loss.item())
-
- policy_loss = surrogate_loss + self.ent_coef * entropy_loss + trust_region_loss
-
- 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 = "PPO",
- eval_log_path: Optional[str] = None,
- reset_num_timesteps: bool = True,
- ) -> "PPO":
-
- 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,
- )
diff --git a/trl_sac/__init__.py b/trl_sac/__init__.py
deleted file mode 100644
index c0e01b7..0000000
--- a/trl_sac/__init__.py
+++ /dev/null
@@ -1,2 +0,0 @@
-from sb3_trl.trl_sac.policies import CnnPolicy, MlpPolicy, MultiInputPolicy
-from sb3_trl.trl_sac.trl_sac import TRL_SAC
diff --git a/trl_sac/policies.py b/trl_sac/policies.py
deleted file mode 100644
index 6fcbea1..0000000
--- a/trl_sac/policies.py
+++ /dev/null
@@ -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,
- )
diff --git a/trl_sac/trl_sac.py b/trl_sac/trl_sac.py
deleted file mode 100644
index 2e884b9..0000000
--- a/trl_sac/trl_sac.py
+++ /dev/null
@@ -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