From 153de76ab8ba2fb8680f8d1d6364083f64f2a6da Mon Sep 17 00:00:00 2001 From: asmith26 Date: Fri, 29 Sep 2023 14:59:39 +0100 Subject: [PATCH 1/2] Initial DroQ --- d3rlpy/algos/qlearning/__init__.py | 1 + d3rlpy/algos/qlearning/base.py | 29 +++--- d3rlpy/algos/qlearning/droq.py | 112 ++++++++++++++++++++++ d3rlpy/algos/qlearning/torch/droq_impl.py | 24 +++++ d3rlpy/models/encoders.py | 3 + d3rlpy/models/torch/encoders.py | 3 + 6 files changed, 159 insertions(+), 13 deletions(-) create mode 100644 d3rlpy/algos/qlearning/droq.py create mode 100644 d3rlpy/algos/qlearning/torch/droq_impl.py diff --git a/d3rlpy/algos/qlearning/__init__.py b/d3rlpy/algos/qlearning/__init__.py index 4ec9fbe6..d4fbf01f 100644 --- a/d3rlpy/algos/qlearning/__init__.py +++ b/d3rlpy/algos/qlearning/__init__.py @@ -7,6 +7,7 @@ from .crr import * from .ddpg import * from .dqn import * +from .droq import * from .explorers import * from .iql import * from .nfq import * diff --git a/d3rlpy/algos/qlearning/base.py b/d3rlpy/algos/qlearning/base.py index 480c89fe..e0bc3218 100644 --- a/d3rlpy/algos/qlearning/base.py +++ b/d3rlpy/algos/qlearning/base.py @@ -574,6 +574,7 @@ def fit_online( n_steps_per_epoch: int = 10000, update_interval: int = 1, update_start_step: int = 0, + utd_ratio: int = 1, random_steps: int = 0, eval_env: Optional[GymEnv] = None, eval_epsilon: float = 0.0, @@ -594,6 +595,7 @@ def fit_online( n_steps_per_epoch: Number of steps per epoch. update_interval: Number of steps per update. update_start_step: Steps before starting updates. + utd_ratio: UTD (update-to-data) ration, the number of updates taken by the agent compared to the number of actual interactions with the environment random_steps: Steps for the initial random explortion. eval_env: Gym-like environment. If None, evaluation is skipped. eval_epsilon: :math:`\\epsilon`-greedy factor during evaluation. @@ -691,19 +693,20 @@ def fit_online( and buffer.transition_count > self.batch_size ): if total_step % update_interval == 0: - # sample mini-batch - with logger.measure_time("sample_batch"): - batch = buffer.sample_transition_batch( - self.batch_size - ) - - # update parameters - with logger.measure_time("algorithm_update"): - loss = self.update(batch) - - # record metrics - for name, val in loss.items(): - logger.add_metric(name, val) + for _ in range(utd_ratio): + # sample mini-batch + with logger.measure_time("sample_batch"): + batch = buffer.sample_transition_batch( + self.batch_size + ) + + # update parameters + with logger.measure_time("algorithm_update"): + loss = self.update(batch) + + # record metrics + for name, val in loss.items(): + logger.add_metric(name, val) # call callback if given if callback: diff --git a/d3rlpy/algos/qlearning/droq.py b/d3rlpy/algos/qlearning/droq.py new file mode 100644 index 00000000..aacb9e9a --- /dev/null +++ b/d3rlpy/algos/qlearning/droq.py @@ -0,0 +1,112 @@ +import dataclasses +import math + +from .torch import SACModules +from .torch.droq_impl import DroQImpl +from .. import SACConfig +from ...base import DeviceArg, register_learnable +from ...constants import ActionSpace +from ...dataset import Shape +from ...models.builders import ( + create_continuous_q_function, + create_normal_policy, + create_parameter, +) +from ...models.encoders import EncoderFactory, VectorEncoderFactory +from .base import QLearningAlgoBase + + +__all__ = ["DroQConfig", "DroQ"] + + +@dataclasses.dataclass() +class DroQConfig(SACConfig): + r"""TODO + """ + critic_encoder_factory: EncoderFactory = VectorEncoderFactory( + dropout_rate=0.9, + use_layer_norm=True, + ) + dropout_rate: int = 0.9 + + def create(self, device: DeviceArg = False) -> "DroQ": + return DroQ(self, device) + + @staticmethod + def get_type() -> str: + return "droq" + + +class DroQ(QLearningAlgoBase[DroQImpl, DroQConfig]): + def inner_create_impl( + self, observation_shape: Shape, action_size: int + ) -> None: + policy = create_normal_policy( + observation_shape, + action_size, + self._config.actor_encoder_factory, + device=self._device, + ) + q_funcs, q_func_forwarder = create_continuous_q_function( + observation_shape, + action_size, + self._config.critic_encoder_factory, + self._config.q_func_factory, + n_ensembles=self._config.n_critics, + device=self._device, + ) + targ_q_funcs, targ_q_func_forwarder = create_continuous_q_function( + observation_shape, + action_size, + self._config.critic_encoder_factory, + self._config.q_func_factory, + n_ensembles=self._config.n_critics, + device=self._device, + ) + log_temp = create_parameter( + (1, 1), + math.log(self._config.initial_temperature), + device=self._device, + ) + + actor_optim = self._config.actor_optim_factory.create( + policy.parameters(), lr=self._config.actor_learning_rate + ) + critic_optim = self._config.critic_optim_factory.create( + q_funcs.parameters(), lr=self._config.critic_learning_rate + ) + if self._config.temp_learning_rate > 0: + temp_optim = self._config.temp_optim_factory.create( + log_temp.parameters(), lr=self._config.temp_learning_rate + ) + else: + temp_optim = None + + modules = SACModules( + policy=policy, + q_funcs=q_funcs, + targ_q_funcs=targ_q_funcs, + log_temp=log_temp, + actor_optim=actor_optim, + critic_optim=critic_optim, + temp_optim=temp_optim, + ) + + self._impl = DroQImpl( + observation_shape=observation_shape, + action_size=action_size, + modules=modules, + q_func_forwarder=q_func_forwarder, + targ_q_func_forwarder=targ_q_func_forwarder, + gamma=self._config.gamma, + tau=self._config.tau, + device=self._device, + ) + + def get_action_type(self) -> ActionSpace: + return ActionSpace.CONTINUOUS + + +# (TODO IF VALID) class DiscreteDroQConfig(LearnableConfig): + +register_learnable(DroQConfig) diff --git a/d3rlpy/algos/qlearning/torch/droq_impl.py b/d3rlpy/algos/qlearning/torch/droq_impl.py new file mode 100644 index 00000000..e0a29c77 --- /dev/null +++ b/d3rlpy/algos/qlearning/torch/droq_impl.py @@ -0,0 +1,24 @@ +import torch + +from . import SACImpl +from ....models.torch import build_squashed_gaussian_distribution +from ....torch_utility import TorchMiniBatch + +__all__ = ["DroQImpl"] + + +class DroQImpl(SACImpl): + def compute_actor_loss(self, batch: TorchMiniBatch) -> torch.Tensor: + dist = build_squashed_gaussian_distribution( + self._modules.policy(batch.observations) + ) + action, log_prob = dist.sample_with_log_prob() + entropy = self._modules.log_temp().exp() * log_prob + q_t = self._q_func_forwarder.compute_expected_q( + # Use "mean" (line 10 of Algorithm 2 in the paper) + batch.observations, action, "mean" + ) + return (entropy - q_t).mean() + + +# (TODO IF VALID) class DiscreteDroQImpl diff --git a/d3rlpy/models/encoders.py b/d3rlpy/models/encoders.py index f516c066..f040e60c 100644 --- a/d3rlpy/models/encoders.py +++ b/d3rlpy/models/encoders.py @@ -138,6 +138,7 @@ class VectorEncoderFactory(EncoderFactory): activation: str = "relu" use_batch_norm: bool = False dropout_rate: Optional[float] = None + use_layer_norm: bool = False exclude_last_activation: bool = False def create(self, observation_shape: Shape) -> VectorEncoder: @@ -147,6 +148,7 @@ def create(self, observation_shape: Shape) -> VectorEncoder: hidden_units=self.hidden_units, use_batch_norm=self.use_batch_norm, dropout_rate=self.dropout_rate, + # use_layer_norm=self.use_layer_norm, activation=create_activation(self.activation), exclude_last_activation=self.exclude_last_activation, ) @@ -164,6 +166,7 @@ def create_with_action( hidden_units=self.hidden_units, use_batch_norm=self.use_batch_norm, dropout_rate=self.dropout_rate, + use_layer_norm=self.use_layer_norm, discrete_action=discrete_action, activation=create_activation(self.activation), exclude_last_activation=self.exclude_last_activation, diff --git a/d3rlpy/models/torch/encoders.py b/d3rlpy/models/torch/encoders.py index bdc8cbfc..cca0e8de 100644 --- a/d3rlpy/models/torch/encoders.py +++ b/d3rlpy/models/torch/encoders.py @@ -226,6 +226,7 @@ def __init__( hidden_units: Optional[Sequence[int]] = None, use_batch_norm: bool = False, dropout_rate: Optional[float] = None, + use_layer_norm: bool = False, discrete_action: bool = False, activation: nn.Module = nn.ReLU(), exclude_last_activation: bool = False, @@ -251,6 +252,8 @@ def __init__( layers.append(nn.BatchNorm1d(out_unit)) if dropout_rate is not None: layers.append(nn.Dropout(dropout_rate)) + if use_layer_norm: + layers.append(nn.LayerNorm(out_unit)) self._layers = nn.Sequential(*layers) def forward(self, x: torch.Tensor, action: torch.Tensor) -> torch.Tensor: From 3099def09c546d02d932e0b63db2b16e4b18f678 Mon Sep 17 00:00:00 2001 From: asmith26 Date: Fri, 29 Sep 2023 15:19:43 +0100 Subject: [PATCH 2/2] Fix bug --- d3rlpy/algos/qlearning/droq.py | 27 ++++++++++++++++++--------- d3rlpy/models/encoders.py | 1 + 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/d3rlpy/algos/qlearning/droq.py b/d3rlpy/algos/qlearning/droq.py index aacb9e9a..4558df79 100644 --- a/d3rlpy/algos/qlearning/droq.py +++ b/d3rlpy/algos/qlearning/droq.py @@ -3,16 +3,16 @@ from .torch import SACModules from .torch.droq_impl import DroQImpl -from .. import SACConfig -from ...base import DeviceArg, register_learnable +from ...base import DeviceArg, register_learnable, LearnableConfig from ...constants import ActionSpace from ...dataset import Shape +from ...models import QFunctionFactory, make_q_func_field, make_optimizer_field, OptimizerFactory from ...models.builders import ( create_continuous_q_function, create_normal_policy, create_parameter, ) -from ...models.encoders import EncoderFactory, VectorEncoderFactory +from ...models.encoders import EncoderFactory, make_encoder_field from .base import QLearningAlgoBase @@ -20,14 +20,23 @@ @dataclasses.dataclass() -class DroQConfig(SACConfig): +class DroQConfig(LearnableConfig): r"""TODO """ - critic_encoder_factory: EncoderFactory = VectorEncoderFactory( - dropout_rate=0.9, - use_layer_norm=True, - ) - dropout_rate: int = 0.9 + actor_learning_rate: float = 3e-4 + critic_learning_rate: float = 3e-4 + temp_learning_rate: float = 3e-4 + actor_optim_factory: OptimizerFactory = make_optimizer_field() + critic_optim_factory: OptimizerFactory = make_optimizer_field() + temp_optim_factory: OptimizerFactory = make_optimizer_field() + actor_encoder_factory: EncoderFactory = make_encoder_field() + critic_encoder_factory: EncoderFactory = make_encoder_field() + q_func_factory: QFunctionFactory = make_q_func_field() + batch_size: int = 256 + gamma: float = 0.99 + tau: float = 0.005 + n_critics: int = 2 + initial_temperature: float = 1.0 def create(self, device: DeviceArg = False) -> "DroQ": return DroQ(self, device) diff --git a/d3rlpy/models/encoders.py b/d3rlpy/models/encoders.py index f040e60c..463aae18 100644 --- a/d3rlpy/models/encoders.py +++ b/d3rlpy/models/encoders.py @@ -192,6 +192,7 @@ class DefaultEncoderFactory(EncoderFactory): activation: str = "relu" use_batch_norm: bool = False dropout_rate: Optional[float] = None + use_layer_norm: bool = False def create(self, observation_shape: Shape) -> Encoder: factory: Union[PixelEncoderFactory, VectorEncoderFactory]