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