From 2614f44f3f997d9a2da78ccbb14893d894431523 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sat, 10 Oct 2020 19:07:54 +0200 Subject: [PATCH] Remove CQL and pretrain --- stable_baselines3/sac/sac.py | 263 +---------------------------------- stable_baselines3/tqc/tqc.py | 173 ----------------------- 2 files changed, 2 insertions(+), 434 deletions(-) diff --git a/stable_baselines3/sac/sac.py b/stable_baselines3/sac/sac.py index 77ff060..98ac9ec 100644 --- a/stable_baselines3/sac/sac.py +++ b/stable_baselines3/sac/sac.py @@ -3,14 +3,12 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union import numpy as np import torch as th from torch.nn import functional as F -from tqdm import tqdm from stable_baselines3.common import logger 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.preprocessing import get_action_dim -from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, ReplayBufferSamples +from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback from stable_baselines3.common.utils import polyak_update from stable_baselines3.sac.policies import SACPolicy @@ -93,11 +91,6 @@ class SAC(OffPolicyAlgorithm): ent_coef: Union[str, float] = "auto", target_update_interval: int = 1, target_entropy: Union[str, float] = "auto", - initial_alpha: float = 5.0, - alpha_threshold: float = 10.0, - alpha_coef: Union[str, float] = "auto", - n_action_samples: int = 10, - use_cql: bool = False, use_sde: bool = False, sde_sample_freq: int = -1, use_sde_at_warmup: bool = False, @@ -145,14 +138,6 @@ class SAC(OffPolicyAlgorithm): self.ent_coef = ent_coef self.target_update_interval = target_update_interval self.ent_coef_optimizer = None - # CQL - self.use_cql = use_cql - self.initial_alpha = initial_alpha - self.alpha_threshold = alpha_threshold - self.n_action_samples = n_action_samples - self.alpha_coef = alpha_coef - self.alpha_coef_tensor = None - self.alpha_optimizer = None if _init_setup_model: self._setup_model() @@ -193,110 +178,24 @@ class SAC(OffPolicyAlgorithm): # is passed self.ent_coef_tensor = th.tensor(float(self.ent_coef)).to(self.device) - # CQL - if self.use_cql: - if self.alpha_coef == "auto": - self.log_alpha = th.log(th.ones(1, device=self.device) * self.initial_alpha).requires_grad_(True) - self.alpha_optimizer = th.optim.Adam([self.log_alpha], lr=self.lr_schedule(1)) - else: - self.alpha_coef_tensor = th.tensor(float(self.alpha_coef)).to(self.device) - self.log_alpha = th.log(self.alpha_coef_tensor) - def _create_aliases(self) -> None: self.actor = self.policy.actor self.critic = self.policy.critic self.critic_target = self.policy.critic_target - def update_alpha(self, replay_data: ReplayBufferSamples): - loss = -self._compute_conservative_loss(replay_data) - self.alpha_optimizer.zero_grad() - loss.backward() - self.alpha_optimizer.step() - - def _compute_conservative_loss(self, replay_data: ReplayBufferSamples): - # from https://github.com/takuseno/d3rlpy - obs_t = replay_data.observations - act_t = replay_data.actions - action_dim = get_action_dim(self.action_space) - n_critics = self.critic.n_critics - assert n_critics == 1 - with th.no_grad(): - policy_actions, n_log_probs = [], [] - for _ in range(self.n_action_samples): - if self.use_sde: - self.actor.reset_noise() - actions_pi, log_prob = self.actor.action_log_prob(replay_data.observations) - policy_actions.append(actions_pi) - n_log_probs.append(log_prob) - # (batch, n, action_dim) - policy_actions = th.cat(policy_actions, dim=1).view( - len(replay_data.observations), self.n_action_samples, action_dim - ) - # assert policy_actions.shape == (len(replay_data.observations), self.n_action_samples, action_dim) - # (batch, n, 1) - n_log_probs = th.cat(n_log_probs).view(len(replay_data.observations), self.n_action_samples, 1) - # assert n_log_probs.shape == (len(replay_data.observations), self.n_action_samples, 1) - - repeated_obs_t = obs_t.expand(self.n_action_samples, *obs_t.shape) - # (n, batch, obs_dim) -> (batch, n, obs_dim) - transposed_obs_t = repeated_obs_t.transpose(0, 1) - # (batch, n, obs_dim) -> (batch * n, obs_dim) - flat_obs_t = transposed_obs_t.reshape(-1, *obs_t.shape[1:]) - # (batch, n, action_dim) -> (batch * n, action_dim) - flat_policy_acts = policy_actions.reshape(-1, action_dim) - - # estimate action-values for policy actions - policy_values = self.critic(flat_obs_t, flat_policy_acts)[0] - policy_values = policy_values.view(n_critics, obs_t.shape[0], self.n_action_samples, 1) - log_probs = n_log_probs.view(1, -1, self.n_action_samples, 1) - - # estimate action-values for actions from uniform distribution - # uniform distribution between [-1.0, 1.0] - random_actions = th.zeros_like(flat_policy_acts).uniform_(-1.0, 1.0) - random_values = self.critic(flat_obs_t, random_actions)[0] - random_values = random_values.view(n_critics, obs_t.shape[0], self.n_action_samples, 1) - - # get maximum value to avoid overflow - base = th.max(policy_values.max(), random_values.max()).detach() - - # compute logsumexp - policy_meanexp = (policy_values - base - log_probs).exp().mean(dim=2) - random_meanexp = (random_values - base).exp().mean(dim=2) / 0.5 - # small constant value seems to be necessary to avoid nan - logsumexp = (0.5 * random_meanexp + 0.5 * policy_meanexp + 1e-10).log() - logsumexp += base - - # estimate action-values for data actions - data_values = self.critic(obs_t, act_t)[0] - - element_wise_loss = logsumexp - data_values - self.alpha_threshold - - # this clipping seems to stabilize training - if self.alpha_coef == "auto": - clipped_alpha = self.log_alpha.clamp(-10.0, 2.0).exp() - else: - clipped_alpha = self.alpha_coef_tensor - - return (clipped_alpha * element_wise_loss).sum(dim=0).mean() - def train(self, gradient_steps: int, batch_size: int = 64) -> None: # Update optimizers learning rate optimizers = [self.actor.optimizer, self.critic.optimizer] if self.ent_coef_optimizer is not None: optimizers += [self.ent_coef_optimizer] - if self.use_cql and self.alpha_coef == "auto": - optimizers += [self.alpha_optimizer] - # Update learning rate according to lr schedule self._update_learning_rate(optimizers) ent_coef_losses, ent_coefs = [], [] actor_losses, critic_losses = [], [] - range_ = tqdm(range(gradient_steps)) if self.use_cql else range(gradient_steps) - - for gradient_step in range_: + for gradient_step in range(gradient_steps): # Sample replay buffer replay_data = self.replay_buffer.sample(batch_size, env=self._vec_normalize_env) @@ -349,10 +248,6 @@ class SAC(OffPolicyAlgorithm): critic_loss = 0.5 * sum([F.mse_loss(current_q, q_backup) for current_q in current_q_estimates]) critic_losses.append(critic_loss.item()) - # CQL - if self.use_cql: - critic_loss += self._compute_conservative_loss(replay_data) - # Optimize the critic self.critic.optimizer.zero_grad() critic_loss.backward() @@ -371,10 +266,6 @@ class SAC(OffPolicyAlgorithm): actor_loss.backward() self.actor.optimizer.step() - # CQL - if self.use_cql and self.alpha_coef == "auto": - self.update_alpha(replay_data) - # Update target networks if gradient_step % self.target_update_interval == 0: polyak_update(self.critic.parameters(), self.critic_target.parameters(), self.tau) @@ -387,150 +278,6 @@ class SAC(OffPolicyAlgorithm): logger.record("train/critic_loss", np.mean(critic_losses)) if len(ent_coef_losses) > 0: logger.record("train/ent_coef_loss", np.mean(ent_coef_losses)) - if self.use_cql: - logger.record("train/alpha", th.exp(self.log_alpha.detach()).item()) - - def pretrain( - self, - gradient_steps: int, - batch_size: int = 64, - n_action_samples: int = -1, - target_update_interval: int = 1, - tau: float = 0.005, - strategy: str = "exp", - reduce: str = "mean", - exp_temperature: float = 1.0, - off_policy_update_freq: int = -1, - ) -> None: - """ - Pretrain with Critic Regularized Regression (CRR) - Paper: https://arxiv.org/abs/2006.15134 - """ - # 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) - - actor_losses, critic_losses = [], [] - - for gradient_step in tqdm(range(gradient_steps)): - - if off_policy_update_freq > 0 and gradient_step % off_policy_update_freq == 0: - self.train(gradient_steps=1, batch_size=batch_size) - continue - - # 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() - else: - ent_coef = self.ent_coef_tensor - - self.replay_buffer.ent_coef = 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 target Q value: min over all critics targets - targets = th.cat(self.critic_target(replay_data.next_observations, next_actions), dim=1) - target_q, _ = th.min(targets, dim=1, keepdim=True) - # add entropy term - target_q = target_q - ent_coef * next_log_prob.reshape(-1, 1) - # td error + entropy term - q_backup = replay_data.rewards + (1 - replay_data.dones) * self.gamma * target_q - - # Get current Q estimates for each critic network - # using action from the replay buffer - current_q_estimates = self.critic(replay_data.observations, replay_data.actions) - - # Compute critic loss - critic_loss = 0.5 * sum([F.mse_loss(current_q, q_backup) for current_q in current_q_estimates]) - critic_losses.append(critic_loss.item()) - - # Optimize the critic - self.critic.optimizer.zero_grad() - critic_loss.backward() - self.critic.optimizer.step() - - if strategy == "bc": - # Behavior cloning - weight = 1 - else: - with th.no_grad(): - qf_buffer = th.min(*self.critic(replay_data.observations, replay_data.actions)) - - if n_action_samples == -1: - actions_pi = self.actor.forward(replay_data.observations, deterministic=True) - qf_agg = th.min(*self.critic(replay_data.observations, actions_pi)) - else: - qf_agg = None - for _ in range(n_action_samples): - if self.use_sde: - self.actor.reset_noise() - actions_pi, _ = self.actor.action_log_prob(replay_data.observations) - - qf_pi = th.min(*self.critic(replay_data.observations, actions_pi.detach())) - if qf_agg is None: - if reduce == "max": - qf_agg = qf_pi - else: - qf_agg = qf_pi / n_action_samples - else: - if reduce == "max": - qf_agg = th.max(qf_pi, qf_agg) - else: - qf_agg += qf_pi / n_action_samples - - advantage = qf_buffer - qf_agg - if strategy == "binary": - # binary advantage - weight = advantage > 0 - else: - # exp advantage - exp_clip = 20.0 - weight = th.clamp(th.exp(advantage / exp_temperature), 0.0, exp_clip) - - # Log prob by the current actor for the sampled state and action - log_prob = self.actor.evaluate_actions(replay_data.observations, replay_data.actions) - log_prob = log_prob.reshape(-1, 1) - - # weigthed regression loss (close to policy gradient loss) - actor_loss = (-log_prob * weight).mean() - # actor_loss = ((actions_pi - replay_data.actions * weight) ** 2).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 % target_update_interval == 0: - polyak_update(self.critic.parameters(), self.critic_target.parameters(), tau) def learn( self, @@ -567,10 +314,4 @@ class SAC(OffPolicyAlgorithm): state_dicts.append("ent_coef_optimizer") else: saved_pytorch_variables.append("ent_coef_tensor") - if self.use_cql: - if self.alpha_coef == "auto": - state_dicts.append("alpha_optimizer") - saved_pytorch_variables.append("log_alpha") - else: - saved_pytorch_variables.append("alpha_coef_tensor") return state_dicts, saved_pytorch_variables diff --git a/stable_baselines3/tqc/tqc.py b/stable_baselines3/tqc/tqc.py index cb6e305..2a017eb 100644 --- a/stable_baselines3/tqc/tqc.py +++ b/stable_baselines3/tqc/tqc.py @@ -2,7 +2,6 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union import numpy as np import torch as th -from tqdm import tqdm from stable_baselines3.common import logger from stable_baselines3.common.buffers import ReplayBuffer @@ -285,178 +284,6 @@ class TQC(OffPolicyAlgorithm): if hasattr(self.actor, "scale_rl"): logger.record("train/scale_rl", th.sigmoid(self.actor.scale_rl).mean().item()) - def pretrain( - self, - gradient_steps: int, - batch_size: int = 64, - n_action_samples: int = -1, - target_update_interval: int = 1, - tau: float = 0.005, - strategy: str = "exp", - reduce: str = "mean", - exp_temperature: float = 1.0, - off_policy_update_freq: int = -1, - normalize_advantage: bool = True, - ) -> None: - """ - Pretrain with Critic Regularized Regression (CRR) - Paper: https://arxiv.org/abs/2006.15134 - """ - # 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) - - actor_losses, critic_losses = [], [] - - for gradient_step in tqdm(range(gradient_steps)): - - if off_policy_update_freq > 0 and gradient_step % off_policy_update_freq == 0: - self.train(gradient_steps=1, batch_size=batch_size) - continue - - # 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 - _, 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() - else: - ent_coef = self.ent_coef_tensor - - self.replay_buffer.ent_coef = 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(): - top_quantiles_to_drop = self.top_quantiles_to_drop_per_net * self.critic.n_critics - # Select action according to policy - next_actions, next_log_prob = self.actor.action_log_prob(replay_data.next_observations) - # Compute and cut quantiles at the next state - # batch x nets x quantiles - next_z = self.critic_target(replay_data.next_observations, next_actions) - sorted_z, _ = th.sort(next_z.reshape(batch_size, -1)) - sorted_z_part = sorted_z[:, : self.critic.quantiles_total - top_quantiles_to_drop] - - target_q = sorted_z_part - ent_coef * next_log_prob.reshape(-1, 1) - # td error + entropy term - q_backup = replay_data.rewards + (1 - replay_data.dones) * self.gamma * target_q - - # Get current Q estimates - # using action from the replay buffer - current_z = self.critic(replay_data.observations, replay_data.actions) - # Compute critic loss - critic_loss = self.quantile_huber_loss(current_z, q_backup) - critic_losses.append(critic_loss.item()) - - # Optimize the critic - self.critic.optimizer.zero_grad() - critic_loss.backward() - self.critic.optimizer.step() - - if strategy == "bc": - # Behavior cloning - weight = 1 - else: - # Tensor version: TODO: check that the reshape works as expected - # cleaner but not faster on cpu for large batch size - # with th.no_grad(): - # # Q-value for the action in the buffer - # qf_buffer = self.critic(replay_data.observations, replay_data.actions).mean(2).mean(1, keepdim=True) - # # Create tensor to avoid loop - # # Note: For SDE, we need to sample several matrices - # obs_ = replay_data.observations.repeat(n_action_samples, 1) - # if self.use_sde: - # self.actor.reset_noise(batch_size * n_action_samples) - # actions_pi, _ = self.actor.action_log_prob(obs_) - # qf_pi = self.critic(obs_, actions_pi.detach()).mean(2).mean(1, keepdim=True) - # # Agregate: reduce mean or reduce max - # if reduce == "max": - # _, qf_agg = qf_pi.reshape(n_action_samples, batch_size, 1).max(axis=0) - # else: - # qf_agg = qf_pi.reshape(n_action_samples, batch_size, 1).mean(axis=0) - with th.no_grad(): - qf_buffer = self.critic(replay_data.observations, replay_data.actions).mean(2).mean(1, keepdim=True) - - # Use the mean (as done in AWAC, cf rlkit) - if n_action_samples == -1: - actions_pi = self.actor.forward(replay_data.observations, deterministic=True) - qf_agg = self.critic(replay_data.observations, actions_pi).mean(2).mean(1, keepdim=True) - else: - qf_agg = None - for _ in range(n_action_samples): - if self.use_sde: - self.actor.reset_noise() - actions_pi, _ = self.actor.action_log_prob(replay_data.observations) - - qf_pi = self.critic(replay_data.observations, actions_pi.detach()).mean(2).mean(1, keepdim=True) - if qf_agg is None: - if reduce == "max": - qf_agg = qf_pi - else: - qf_agg = qf_pi / n_action_samples - else: - if reduce == "max": - qf_agg = th.max(qf_pi, qf_agg) - else: - qf_agg += qf_pi / n_action_samples - - advantage = qf_buffer - qf_agg - - if normalize_advantage: - mean_values = advantage.mean(dim=0, keepdims=True) - std_values = advantage.std(dim=0, keepdims=True) + 1e-5 - advantage = (advantage - mean_values) / std_values - - if strategy == "binary": - # binary advantage - weight = advantage > 0 - else: - # exp advantage - exp_clip = 20.0 - weight = th.clamp(th.exp(advantage / exp_temperature), 0.0, exp_clip) - - # Log prob by the current actor for the sampled state and action - log_prob = self.actor.evaluate_actions(replay_data.observations, replay_data.actions) - log_prob = log_prob.reshape(-1, 1) - - # weigthed regression loss (close to policy gradient loss) - actor_loss = (-log_prob * weight).mean() - # actor_loss = ((actions_pi - replay_data.actions * weight) ** 2).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 % target_update_interval == 0: - polyak_update(self.critic.parameters(), self.critic_target.parameters(), tau) - - if self.use_sde: - print(f"std={(self.actor.get_std()).mean().item()}") - def learn( self, total_timesteps: int,