import warnings from typing import Any, ClassVar, Dict, Optional, Type, TypeVar, Union import numpy as np import torch as th from gymnasium import spaces from torch.nn import functional as F from stable_baselines3.common.buffers import RolloutBuffer from ..common.on_policy_algorithm import BetterOnPolicyAlgorithm from ..common.policies import ActorCriticPolicy, BasePolicy from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule from stable_baselines3.common.utils import explained_variance, get_schedule_fn from metastable_projections import BaseProjectionLayer, FrobeniusProjectionLayer, WassersteinProjectionLayer, KLProjectionLayer import metastable_projections SelfTRPL = TypeVar("SelfTRPL", bound="TRPL") def castProjection(proj): if type(proj)==str: return getattr(metastable_projections, proj + 'ProjectionLayer') return proj class TRPL(BetterOnPolicyAlgorithm): """ TODO: Bla :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: Should not be used for normal TRPL usage. Is only here to bridge the gap to PPO. Clipping parameter, it can be a function of the current progress remaining (from 1 to 0). Setting it to None will result in no clipping (default) :param clip_range_vf: Should not be used for normal TRPL usage. Is only here to bridge the gap to PPO. 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: Normally a good idea; but TRPL actually often works better without normalization of the advantage. Whether to normalize or not the advantage. (Default: False) :param ent_coef: Entropy coefficient for the loss calculation :param vf_coef: Value function coefficient for the loss calculation :param max_grad_norm: Should not be used for normal TRPL usage. Is only here to bridge the gap to PPO.. The maximum value for the gradient clipping. Setting it to None will result in no gradient clipping (default) :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_pca: Wether to use Prior Conditioned Annealing. :param rollout_buffer_class: Rollout buffer class to use. If ``None``, it will be automatically selected. :param rollout_buffer_kwargs: Keyword arguments to pass to the rollout buffer on creation :param target_kl: Not part of reference implementation of TRPL, but we still ported it over from sb3's PPO. Default to None: No limit. 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) :param stats_window_size: Window size for the rollout logging, specifying the number of episodes to average the reported success rate, mean episode length, and mean reward over :param tensorboard_log: the log location for tensorboard (if None, no logging) :param policy_kwargs: additional arguments to be passed to the policy on creation :param verbose: Verbosity level: 0 for no output, 1 for info messages (such as device or wrappers used), 2 for debug messages :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: ClassVar[Dict[str, Type[BasePolicy]]] = { "MlpPolicy": ActorCriticPolicy } 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, None] = None, clip_range_vf: Union[None, float, Schedule] = None, normalize_advantage: bool = False, ent_coef: float = 0.0, vf_coef: float = 0.5, max_grad_norm: Union[float, None] = None, use_sde: bool = False, sde_sample_freq: int = -1, use_pca: bool = False, pca_is: bool = False, projection_class: Union[BaseProjectionLayer, str] = BaseProjectionLayer, projection_kwargs: Optional[Dict[str, Any]] = None, rollout_buffer_class: Optional[Type[RolloutBuffer]] = None, rollout_buffer_kwargs: Optional[Dict[str, Any]] = None, target_kl: Optional[float] = None, stats_window_size: int = 100, tensorboard_log: Optional[str] = None, 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, ): self.projection_class = castProjection(projection_class) if projection_kwargs is None: projection_kwargs = {} self.projection_kwargs = projection_kwargs self.projection = self.projection_class(**self.projection_kwargs) if policy_kwargs is None: policy_kwargs = {} policy_kwargs['policy_projection'] = self.projection if 'dist_kwargs' not in policy_kwargs: policy_kwargs['dist_kwargs'] = {} if use_pca: policy_kwargs['dist_kwargs']['msqrt_induces_full'] = self.projection_class == WassersteinProjectionLayer 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, use_pca=use_pca, pca_is=pca_is, rollout_buffer_class=rollout_buffer_class, rollout_buffer_kwargs=rollout_buffer_kwargs, stats_window_size=stats_window_size, tensorboard_log=tensorboard_log, policy_kwargs=policy_kwargs, verbose=verbose, device=device, seed=seed, _init_setup_model=False, supported_action_spaces=( spaces.Box, spaces.Discrete, spaces.MultiDiscrete, spaces.MultiBinary, ), ) print('[i] Using metastable version of TRPL') # 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 or ( not normalize_advantage ), 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 if clip_range == False: clip_range = None self.clip_range = clip_range if clip_range_vf == False: clip_range_vf = None self.clip_range_vf = clip_range_vf self.normalize_advantage = normalize_advantage self.target_kl = target_kl if _init_setup_model: self._setup_model() def _setup_model(self) -> None: super()._setup_model() # Initialize schedules for policy/value clipping if self.clip_range is not None: 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 if self.clip_range is not None: 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) trust_region_losses = [] entropy_losses = [] pg_losses = [] value_losses = [] policy_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 or self.use_pca: self.policy.reset_noise(self.batch_size) values, log_prob, entropy, trust_region_loss = self.policy.evaluate_actions(rollout_data, actions) values = values.flatten() # Normalize advantage advantages = rollout_data.advantages # Normalization does not make sense if mini batchsize == 1, see GH issue #325 if self.normalize_advantage and len(advantages) > 1: advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) # ratio between old and new policy, should be one at the first iteration # With pca_is=True, old_log_prob will be of the conditioned old dist (doing two Importance Sampling in one) ratio = th.exp(log_prob - rollout_data.old_log_prob) # clipped surrogate loss if self.clip_range is None: surrogate_loss = -(advantages * ratio).mean() else: surrogate_loss_1 = advantages * ratio surrogate_loss_2 = advantages * th.clamp(ratio, 1 - clip_range, 1 + clip_range) surrogate_loss = -th.min(surrogate_loss_1, surrogate_loss_2).mean() # Logging pg_losses.append(surrogate_loss.item()) if self.clip_range is not None: clip_fraction = th.mean((th.abs(ratio - 1) > clip_range).float()).item() else: clip_fraction = 0 clip_fractions.append(clip_fraction) if self.clip_range_vf is None: # No clipping values_pred = values else: # Clip the difference 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) policy_loss = trust_region_loss + surrogate_loss loss = policy_loss + self.ent_coef * entropy_loss + self.vf_coef * value_loss entropy_losses.append(entropy_loss.item()) trust_region_losses.append(trust_region_loss.item()) policy_losses.append(policy_loss.item()) # 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 if self.max_grad_norm is not None: th.nn.utils.clip_grad_norm_(self.policy.parameters(), self.max_grad_norm) self.policy.optimizer.step() self._n_updates += 1 if not continue_training: break explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten()) # Logs self.logger.record("train/entropy_loss", np.mean(entropy_losses)) self.logger.record("train/policy_gradient_loss", np.mean(pg_losses)) self.logger.record("train/trust_region_loss", np.mean(trust_region_losses)) self.logger.record("train/policy_loss", np.mean(policy_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") if self.clip_range is not None: 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: SelfTRPL, total_timesteps: int, callback: MaybeCallback = None, log_interval: int = 1, tb_log_name: str = "TRPL", reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfTRPL: return super().learn( total_timesteps=total_timesteps, callback=callback, log_interval=log_interval, tb_log_name=tb_log_name, reset_num_timesteps=reset_num_timesteps, progress_bar=progress_bar, )