Source code for sb3_contrib.trpo.trpo

import copy
import warnings
from functools import partial
from typing import Any, ClassVar, Dict, List, Optional, Tuple, Type, TypeVar, Union

import numpy as np
import torch as th
from gymnasium import spaces
from stable_baselines3.common.buffers import RolloutBuffer
from stable_baselines3.common.distributions import kl_divergence
from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
from stable_baselines3.common.policies import ActorCriticPolicy, BasePolicy
from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, RolloutBufferSamples, Schedule
from stable_baselines3.common.utils import explained_variance
from torch import nn
from torch.nn import functional as F

from sb3_contrib.common.utils import conjugate_gradient_solver, flat_grad
from sb3_contrib.trpo.policies import CnnPolicy, MlpPolicy, MultiInputPolicy

SelfTRPO = TypeVar("SelfTRPO", bound="TRPO")

[docs]class TRPO(OnPolicyAlgorithm): """ Trust Region Policy Optimization (TRPO) Paper: Code: This implementation borrows code from OpenAI Spinning Up ( and Stable Baselines (TRPO from Introduction to TRPO: :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 for the value function, 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 :param batch_size: Minibatch size for the value function :param gamma: Discount factor :param cg_max_steps: maximum number of steps in the Conjugate Gradient algorithm for computing the Hessian vector product :param cg_damping: damping in the Hessian vector product computation :param line_search_shrinking_factor: step-size reduction factor for the line-search (i.e., ``theta_new = theta + alpha^i * step``) :param line_search_max_iter: maximum number of iteration for the backtracking line-search :param n_critic_updates: number of critic updates per policy update :param gae_lambda: Factor for trade-off of bias vs variance for Generalized Advantage Estimator :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 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 normalize_advantage: Whether to normalize or not the advantage :param target_kl: Target Kullback-Leibler divergence between updates. Should be small for stability. Values like 0.01, 0.05. :param sub_sampling_factor: Sub-sample the batch to make computation faster see p40-42 of John Schulman thesis :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: 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: ClassVar[Dict[str, Type[BasePolicy]]] = { "MlpPolicy": MlpPolicy, "CnnPolicy": CnnPolicy, "MultiInputPolicy": MultiInputPolicy, } def __init__( self, policy: Union[str, Type[ActorCriticPolicy]], env: Union[GymEnv, str], learning_rate: Union[float, Schedule] = 1e-3, n_steps: int = 2048, batch_size: int = 128, gamma: float = 0.99, cg_max_steps: int = 15, cg_damping: float = 0.1, line_search_shrinking_factor: float = 0.8, line_search_max_iter: int = 10, n_critic_updates: int = 10, gae_lambda: float = 0.95, use_sde: bool = False, sde_sample_freq: int = -1, rollout_buffer_class: Optional[Type[RolloutBuffer]] = None, rollout_buffer_kwargs: Optional[Dict[str, Any]] = None, normalize_advantage: bool = True, target_kl: float = 0.01, sub_sampling_factor: int = 1, 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, ): super().__init__( policy, env, learning_rate=learning_rate, n_steps=n_steps, gamma=gamma, gae_lambda=gae_lambda, ent_coef=0.0, # entropy bonus is not used by TRPO vf_coef=0.0, # value function is optimized separately max_grad_norm=0.0, use_sde=use_sde, sde_sample_freq=sde_sample_freq, 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, ), ) self.normalize_advantage = normalize_advantage # Sanity check, otherwise it will lead to noisy gradient and NaN # because of the advantage normalization 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 if normalize_advantage: assert buffer_size > 1, ( "`n_steps * n_envs` must be greater than 1. " f"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 # Conjugate gradients parameters self.cg_max_steps = cg_max_steps self.cg_damping = cg_damping # Backtracking line search parameters self.line_search_shrinking_factor = line_search_shrinking_factor self.line_search_max_iter = line_search_max_iter self.target_kl = target_kl self.n_critic_updates = n_critic_updates self.sub_sampling_factor = sub_sampling_factor if _init_setup_model: self._setup_model() def _compute_actor_grad( self, kl_div: th.Tensor, policy_objective: th.Tensor ) -> Tuple[List[nn.Parameter], th.Tensor, th.Tensor, List[Tuple[int, ...]]]: """ Compute actor gradients for kl div and surrogate objectives. :param kl_div: The KL divergence objective :param policy_objective: The surrogate objective ("classic" policy gradient) :return: List of actor params, gradients and gradients shape. """ # This is necessary because not all the parameters in the policy have gradients w.r.t. the KL divergence # The policy objective is also called surrogate objective policy_objective_gradients_list = [] # Contains the gradients of the KL divergence grad_kl_list = [] # Contains the shape of the gradients of the KL divergence w.r.t each parameter # This way the flattened gradient can be reshaped back into the original shapes and applied to # the parameters grad_shape: List[Tuple[int, ...]] = [] # Contains the parameters which have non-zeros KL divergence gradients # The list is used during the line-search to apply the step to each parameters actor_params: List[nn.Parameter] = [] for name, param in self.policy.named_parameters(): # Skip parameters related to value function based on name # this work for built-in policies only (not custom ones) if "value" in name: continue # For each parameter we compute the gradient of the KL divergence w.r.t to that parameter kl_param_grad, *_ = th.autograd.grad( kl_div, param, create_graph=True, retain_graph=True, allow_unused=True, only_inputs=True, ) # If the gradient is not zero (not None), we store the parameter in the actor_params list # and add the gradient and its shape to grad_kl and grad_shape respectively if kl_param_grad is not None: # If the parameter impacts the KL divergence (i.e. the policy) # we compute the gradient of the policy objective w.r.t to the parameter # this avoids computing the gradient if it's not going to be used in the conjugate gradient step policy_objective_grad, *_ = th.autograd.grad(policy_objective, param, retain_graph=True, only_inputs=True) grad_shape.append(kl_param_grad.shape) grad_kl_list.append(kl_param_grad.reshape(-1)) policy_objective_gradients_list.append(policy_objective_grad.reshape(-1)) actor_params.append(param) # Gradients are concatenated before the conjugate gradient step policy_objective_gradients = grad_kl = return actor_params, policy_objective_gradients, grad_kl, grad_shape
[docs] 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) policy_objective_values = [] kl_divergences = [] line_search_results = [] value_losses = [] # This will only loop once (get all data in one go) for rollout_data in self.rollout_buffer.get(batch_size=None): # Optional: sub-sample data for faster computation if self.sub_sampling_factor > 1: rollout_data = RolloutBufferSamples( rollout_data.observations[:: self.sub_sampling_factor], rollout_data.actions[:: self.sub_sampling_factor], None, # type: ignore[arg-type] # old values, not used here rollout_data.old_log_prob[:: self.sub_sampling_factor], rollout_data.advantages[:: self.sub_sampling_factor], None, # type: ignore[arg-type] # returns, not used here ) 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: # batch_size is only used for the value function self.policy.reset_noise(actions.shape[0]) with th.no_grad(): # Note: is copy enough, no need for deepcopy? # If using gSDE and deepcopy, we need to use `old_distribution.distribution` # directly to avoid PyTorch errors. old_distribution = copy.copy(self.policy.get_distribution(rollout_data.observations)) distribution = self.policy.get_distribution(rollout_data.observations) log_prob = distribution.log_prob(actions) advantages = rollout_data.advantages if self.normalize_advantage: advantages = (advantages - advantages.mean()) / (rollout_data.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) # surrogate policy objective policy_objective = (advantages * ratio).mean() # KL divergence kl_div = kl_divergence(distribution, old_distribution).mean() # Surrogate & KL gradient self.policy.optimizer.zero_grad() actor_params, policy_objective_gradients, grad_kl, grad_shape = self._compute_actor_grad(kl_div, policy_objective) # Hessian-vector dot product function used in the conjugate gradient step hessian_vector_product_fn = partial(self.hessian_vector_product, actor_params, grad_kl) # Computing search direction search_direction = conjugate_gradient_solver( hessian_vector_product_fn, policy_objective_gradients, max_iter=self.cg_max_steps, ) # Maximal step length line_search_max_step_size = 2 * self.target_kl line_search_max_step_size /= th.matmul( search_direction, hessian_vector_product_fn(search_direction, retain_graph=False) ) line_search_max_step_size = th.sqrt(line_search_max_step_size) # type: ignore[assignment, arg-type] line_search_backtrack_coeff = 1.0 original_actor_params = [param.detach().clone() for param in actor_params] is_line_search_success = False with th.no_grad(): # Line-search (backtracking) for _ in range(self.line_search_max_iter): start_idx = 0 # Applying the scaled step direction for param, original_param, shape in zip(actor_params, original_actor_params, grad_shape): n_params = param.numel() = ( + line_search_backtrack_coeff * line_search_max_step_size * search_direction[start_idx : (start_idx + n_params)].view(shape) ) start_idx += n_params # Recomputing the policy log-probabilities distribution = self.policy.get_distribution(rollout_data.observations) log_prob = distribution.log_prob(actions) # New policy objective ratio = th.exp(log_prob - rollout_data.old_log_prob) new_policy_objective = (advantages * ratio).mean() # New KL-divergence kl_div = kl_divergence(distribution, old_distribution).mean() # Constraint criteria: # we need to improve the surrogate policy objective # while being close enough (in term of kl div) to the old policy if (kl_div < self.target_kl) and (new_policy_objective > policy_objective): is_line_search_success = True break # Reducing step size if line-search wasn't successful line_search_backtrack_coeff *= self.line_search_shrinking_factor line_search_results.append(is_line_search_success) if not is_line_search_success: # If the line-search wasn't successful we revert to the original parameters for param, original_param in zip(actor_params, original_actor_params): = policy_objective_values.append(policy_objective.item()) kl_divergences.append(0.0) else: policy_objective_values.append(new_policy_objective.item()) kl_divergences.append(kl_div.item()) # Critic update for _ in range(self.n_critic_updates): for rollout_data in self.rollout_buffer.get(self.batch_size): values_pred = self.policy.predict_values(rollout_data.observations) value_loss = F.mse_loss(rollout_data.returns, values_pred.flatten()) value_losses.append(value_loss.item()) self.policy.optimizer.zero_grad() value_loss.backward() # Removing gradients of parameters shared with the actor # otherwise it defeats the purposes of the KL constraint for param in actor_params: param.grad = None self.policy.optimizer.step() self._n_updates += 1 explained_var = explained_variance(self.rollout_buffer.values.flatten(), self.rollout_buffer.returns.flatten()) # Logs self.logger.record("train/policy_objective", np.mean(policy_objective_values)) self.logger.record("train/value_loss", np.mean(value_losses)) self.logger.record("train/kl_divergence_loss", np.mean(kl_divergences)) self.logger.record("train/explained_variance", explained_var) self.logger.record("train/is_line_search_success", np.mean(line_search_results)) 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")
[docs] def hessian_vector_product( self, params: List[nn.Parameter], grad_kl: th.Tensor, vector: th.Tensor, retain_graph: bool = True ) -> th.Tensor: """ Computes the matrix-vector product with the Fisher information matrix. :param params: list of parameters used to compute the Hessian :param grad_kl: flattened gradient of the KL divergence between the old and new policy :param vector: vector to compute the dot product the hessian-vector dot product with :param retain_graph: if True, the graph will be kept after computing the Hessian :return: Hessian-vector dot product (with damping) """ jacobian_vector_product = (grad_kl * vector).sum() return flat_grad(jacobian_vector_product, params, retain_graph=retain_graph) + self.cg_damping * vector
[docs] def learn( self: SelfTRPO, total_timesteps: int, callback: MaybeCallback = None, log_interval: int = 1, tb_log_name: str = "TRPO", reset_num_timesteps: bool = True, progress_bar: bool = False, ) -> SelfTRPO: 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, )