Add DroQ

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 *

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:

d3rlpy/algos/qlearning/droq.py

@@ -0,0 +1,121 @@
+import dataclasses
+import math
+
+from .torch import SACModules
+from .torch.droq_impl import DroQImpl
+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, make_encoder_field
+from .base import QLearningAlgoBase
+
+
+__all__ = ["DroQConfig", "DroQ"]
+
+
+@dataclasses.dataclass()
+class DroQConfig(LearnableConfig):
+    r"""TODO
+    """
+    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)
+
+    @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)

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

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

d3rlpy/models/torch/encoders.py

@@ -226,6 +226,7 @@
         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 @@
                 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: