Source code for d3rlpy.algos.cql

from .base import AlgoBase
from .dqn import DoubleDQN
from .torch.cql_impl import CQLImpl, DiscreteCQLImpl
from ..optimizers import AdamFactory
from ..argument_utils import check_encoder
from ..argument_utils import check_use_gpu
from ..argument_utils import check_augmentation
from ..argument_utils import check_q_func


[docs]class CQL(AlgoBase): r""" Conservative Q-Learning algorithm. CQL is a SAC-based data-driven deep reinforcement learning algorithm, which achieves state-of-the-art performance in offline RL problems. CQL mitigates overestimation error by minimizing action-values under the current policy and maximizing values under data distribution for underestimation issue. .. math:: L(\theta_i) = \alpha \mathbb{E}_{s_t \sim D} [\log{\sum_a \exp{Q_{\theta_i}(s_t, a)}} - \mathbb{E}_{a \sim D} [Q_{\theta_i}(s, a)] - \tau] + L_{SAC}(\theta_i) where :math:`\alpha` is an automatically adjustable value via Lagrangian dual gradient descent and :math:`\tau` is a threshold value. If the action-value difference is smaller than :math:`\tau`, the :math:`\alpha` will become smaller. Otherwise, the :math:`\alpha` will become larger to aggressively penalize action-values. In continuous control, :math:`\log{\sum_a \exp{Q(s, a)}}` is computed as follows. .. math:: \log{\sum_a \exp{Q(s, a)}} \approx \log{( \frac{1}{2N} \sum_{a_i \sim \text{Unif}(a)}^N [\frac{\exp{Q(s, a_i)}}{\text{Unif}(a)}] + \frac{1}{2N} \sum_{a_i \sim \pi_\phi(a|s)}^N [\frac{\exp{Q(s, a_i)}}{\pi_\phi(a_i|s)}])} where :math:`N` is the number of sampled actions. The rest of optimization is exactly same as :class:`d3rlpy.algos.SAC`. References: * `Kumar et al., Conservative Q-Learning for Offline Reinforcement Learning. <https://arxiv.org/abs/2006.04779>`_ Args: actor_learning_rate (float): learning rate for policy function. critic_learning_rate (float): learning rate for Q functions. temp_learning_rate (float): learning rate for temperature parameter of SAC. alpha_learning_rate (float): learning rate for :math:`\alpha`. actor_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the actor. critic_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the critic. temp_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the temperature. alpha_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for :math:`\alpha`. actor_encoder_factory (d3rlpy.encoders.EncoderFactory or str): encoder factory for the actor. critic_encoder_factory (d3rlpy.encoders.EncoderFactory or str): encoder factory for the critic. q_func_factory (d3rlpy.q_functions.QFunctionFactory or str): Q function factory. batch_size (int): mini-batch size. n_frames (int): the number of frames to stack for image observation. n_steps (int): N-step TD calculation. gamma (float): discount factor. tau (float): target network synchronization coefficiency. n_critics (int): the number of Q functions for ensemble. bootstrap (bool): flag to bootstrap Q functions. share_encoder (bool): flag to share encoder network. update_actor_interval (int): interval to update policy function. initial_temperature (float): initial temperature value. initial_alpha (float): initial :math:`\alpha` value. alpha_threshold (float): threshold value described as :math:`\tau`. n_action_samples (int): the number of sampled actions to compute :math:`\log{\sum_a \exp{Q(s, a)}}`. use_gpu (bool, int or d3rlpy.gpu.Device): flag to use GPU, device ID or device. scaler (d3rlpy.preprocessing.Scaler or str): preprocessor. The available options are `['pixel', 'min_max', 'standard']` augmentation (d3rlpy.augmentation.AugmentationPipeline or list(str)): augmentation pipeline. dynamics (d3rlpy.dynamics.base.DynamicsBase): dynamics model for data augmentation. impl (d3rlpy.algos.torch.cql_impl.CQLImpl): algorithm implementation. Attributes: actor_learning_rate (float): learning rate for policy function. critic_learning_rate (float): learning rate for Q functions. temp_learning_rate (float): learning rate for temperature parameter of SAC. alpha_learning_rate (float): learning rate for :math:`\alpha`. actor_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the actor. critic_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the critic. temp_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for the temperature. alpha_optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory for :math:`\alpha`. actor_encoder_factory (d3rlpy.encoders.EncoderFactory): encoder factory for the actor. critic_encoder_factory (d3rlpy.encoders.EncoderFactory): encoder factory for the critic. q_func_factory (d3rlpy.q_functions.QFunctionFactory): Q function factory. batch_size (int): mini-batch size. n_frames (int): the number of frames to stack for image observation. n_steps (int): N-step TD calculation. gamma (float): discount factor. tau (float): target network synchronization coefficiency. n_critics (int): the number of Q functions for ensemble. bootstrap (bool): flag to bootstrap Q functions. share_encoder (bool): flag to share encoder network. update_actor_interval (int): interval to update policy function. initial_temperature (float): initial temperature value. initial_alpha (float): initial :math:`\alpha` value. alpha_threshold (float): threshold value described as :math:`\tau`. n_action_samples (int): the number of sampled actions to compute :math:`\log{\sum_a \exp{Q(s, a)}}`. use_gpu (d3rlpy.gpu.Device): GPU device. scaler (d3rlpy.preprocessing.Scaler): preprocessor. augmentation (d3rlpy.augmentation.AugmentationPipeline): augmentation pipeline. dynamics (d3rlpy.dynamics.base.DynamicsBase): dynamics model. impl (d3rlpy.algos.torch.cql_impl.CQLImpl): algorithm implementation. eval_results_ (dict): evaluation results. """ def __init__(self, *, actor_learning_rate=3e-5, critic_learning_rate=3e-4, temp_learning_rate=3e-5, alpha_learning_rate=3e-4, actor_optim_factory=AdamFactory(), critic_optim_factory=AdamFactory(), temp_optim_factory=AdamFactory(), alpha_optim_factory=AdamFactory(), actor_encoder_factory='default', critic_encoder_factory='default', q_func_factory='mean', batch_size=100, n_frames=1, n_steps=1, gamma=0.99, tau=0.005, n_critics=2, bootstrap=False, share_encoder=False, update_actor_interval=1, initial_temperature=1.0, initial_alpha=5.0, alpha_threshold=10.0, n_action_samples=10, use_gpu=False, scaler=None, augmentation=None, dynamics=None, impl=None, **kwargs): super().__init__(batch_size=batch_size, n_frames=n_frames, n_steps=n_steps, gamma=gamma, scaler=scaler, dynamics=dynamics) self.actor_learning_rate = actor_learning_rate self.critic_learning_rate = critic_learning_rate self.temp_learning_rate = temp_learning_rate self.alpha_learning_rate = alpha_learning_rate self.actor_optim_factory = actor_optim_factory self.critic_optim_factory = critic_optim_factory self.temp_optim_factory = temp_optim_factory self.alpha_optim_factory = alpha_optim_factory self.actor_encoder_factory = check_encoder(actor_encoder_factory) self.critic_encoder_factory = check_encoder(critic_encoder_factory) self.q_func_factory = check_q_func(q_func_factory) self.tau = tau self.n_critics = n_critics self.bootstrap = bootstrap self.share_encoder = share_encoder self.update_actor_interval = update_actor_interval self.initial_temperature = initial_temperature self.initial_alpha = initial_alpha self.alpha_threshold = alpha_threshold self.n_action_samples = n_action_samples self.augmentation = check_augmentation(augmentation) self.use_gpu = check_use_gpu(use_gpu) self.impl = impl
[docs] def create_impl(self, observation_shape, action_size): self.impl = CQLImpl(observation_shape=observation_shape, action_size=action_size, actor_learning_rate=self.actor_learning_rate, critic_learning_rate=self.critic_learning_rate, temp_learning_rate=self.temp_learning_rate, alpha_learning_rate=self.alpha_learning_rate, actor_optim_factory=self.actor_optim_factory, critic_optim_factory=self.critic_optim_factory, temp_optim_factory=self.temp_optim_factory, alpha_optim_factory=self.alpha_optim_factory, actor_encoder_factory=self.actor_encoder_factory, critic_encoder_factory=self.critic_encoder_factory, q_func_factory=self.q_func_factory, gamma=self.gamma, tau=self.tau, n_critics=self.n_critics, bootstrap=self.bootstrap, share_encoder=self.share_encoder, initial_temperature=self.initial_temperature, initial_alpha=self.initial_alpha, alpha_threshold=self.alpha_threshold, n_action_samples=self.n_action_samples, use_gpu=self.use_gpu, scaler=self.scaler, augmentation=self.augmentation) self.impl.build()
[docs] def update(self, epoch, total_step, batch): critic_loss = self.impl.update_critic(batch.observations, batch.actions, batch.next_rewards, batch.next_observations, batch.terminals, batch.n_steps) if total_step % self.update_actor_interval == 0: actor_loss = self.impl.update_actor(batch.observations) temp_loss, temp = self.impl.update_temp(batch.observations) alpha_loss, alpha = self.impl.update_alpha(batch.observations, batch.actions) self.impl.update_critic_target() self.impl.update_actor_target() else: actor_loss = None temp_loss = None temp = None alpha_loss = None alpha = None return critic_loss, actor_loss, temp_loss, temp, alpha_loss, alpha
def _get_loss_labels(self): return [ 'critic_loss', 'actor_loss', 'temp_loss', 'temp', 'alpha_loss', 'alpha' ]
[docs]class DiscreteCQL(DoubleDQN): r""" Discrete version of Conservative Q-Learning algorithm. Discrete version of CQL is a DoubleDQN-based data-driven deep reinforcement learning algorithm (the original paper uses DQN), which achieves state-of-the-art performance in offline RL problems. CQL mitigates overestimation error by minimizing action-values under the current policy and maximizing values under data distribution for underestimation issue. .. math:: L(\theta) = \mathbb{E}_{s_t \sim D} [\log{\sum_a \exp{Q_{\theta}(s_t, a)}} - \mathbb{E}_{a \sim D} [Q_{\theta}(s, a)]] + L_{DoubleDQN}(\theta) References: * `Kumar et al., Conservative Q-Learning for Offline Reinforcement Learning. <https://arxiv.org/abs/2006.04779>`_ Args: learning_rate (float): learning rate. optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory. encoder_factory (d3rlpy.encoders.EncoderFactory or str): encoder factory. q_func_factory (d3rlpy.q_functions.QFunctionFactory or str): Q function factory. batch_size (int): mini-batch size. n_frames (int): the number of frames to stack for image observation. n_steps (int): N-step TD calculation. gamma (float): discount factor. n_critics (int): the number of Q functions for ensemble. bootstrap (bool): flag to bootstrap Q functions. target_update_interval (int): interval to synchronize the target network. use_gpu (bool, int or d3rlpy.gpu.Device): flag to use GPU, device ID or device. scaler (d3rlpy.preprocessing.Scaler or str): preprocessor. The available options are `['pixel', 'min_max', 'standard']` augmentation (d3rlpy.augmentation.AugmentationPipeline or list(str)): augmentation pipeline. dynamics (d3rlpy.dynamics.base.DynamicsBase): dynamics model for data augmentation. impl (d3rlpy.algos.torch.cql_impl.DiscreteCQLImpl): algorithm implementation. Attributes: learning_rate (float): learning rate. optim_factory (d3rlpy.optimizers.OptimizerFactory): optimizer factory. encoder_factory (d3rlpy.encoders.EncoderFactory): encoder factory. q_func_factory (d3rlpy.q_functions.QFunctionFactory): Q function factory. batch_size (int): mini-batch size. n_frames (int): the number of frames to stack for image observation. n_steps (int): N-step TD calculation. gamma (float): discount factor. n_critics (int): the number of Q functions for ensemble. bootstrap (bool): flag to bootstrap Q functions. target_update_interval (int): interval to synchronize the target network. use_gpu (d3rlpy.gpu.Device): GPU device. scaler (d3rlpy.preprocessing.Scaler): preprocessor. augmentation (d3rlpy.augmentation.AugmentationPipeline): augmentation pipeline. dynamics (d3rlpy.dynamics.base.DynamicsBase): dynamics model. impl (d3rlpy.algos.torch.CQLImpl.DiscreteCQLImpl): algorithm implementation. eval_results_ (dict): evaluation results. """
[docs] def create_impl(self, observation_shape, action_size): self.impl = DiscreteCQLImpl(observation_shape=observation_shape, action_size=action_size, learning_rate=self.learning_rate, optim_factory=self.optim_factory, encoder_factory=self.encoder_factory, q_func_factory=self.q_func_factory, gamma=self.gamma, n_critics=self.n_critics, bootstrap=self.bootstrap, share_encoder=self.share_encoder, use_gpu=self.use_gpu, scaler=self.scaler, augmentation=self.augmentation) self.impl.build()