saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-28 20:11:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Remove CQL and pretrain

Browse source
This commit is contained in:
Antonin RAFFIN 2020-10-10 19:07:54 +02:00
parent 210d1c5699
commit 2614f44f3f
2 changed files with 2 additions and 434 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

263
stable_baselines3/sac/sac.py
View file

@ -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

173
stable_baselines3/tqc/tqc.py
View file

@ -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,