Implement SimbaV2 layers

d3rlpy/algos/qlearning/bcq.py

@@ -11,7 +11,11 @@
     create_vae_encoder,
 )
 from ...models.encoders import EncoderFactory, make_encoder_field
-from ...models.q_functions import QFunctionFactory, make_q_func_field
+from ...models.q_functions import (
+    MeanQFunctionFactory,
+    QFunctionFactory,
+    make_q_func_field,
+)
 from ...models.torch import CategoricalPolicy, compute_output_size
 from ...optimizers.optimizers import OptimizerFactory, make_optimizer_field
 from ...types import Shape
@@ -122,8 +126,6 @@ class BCQConfig(LearnableConfig):
             Encoder factory for the critic.
         imitator_encoder_factory (d3rlpy.models.encoders.EncoderFactory):
             Encoder factory for the conditional VAE.
-        q_func_factory (d3rlpy.models.q_functions.QFunctionFactory):
-            Q function factory.
         batch_size (int): Mini-batch size.
         gamma (float): Discount factor.
         tau (float): Target network synchronization coefficiency.
@@ -149,7 +151,6 @@ class BCQConfig(LearnableConfig):
     actor_encoder_factory: EncoderFactory = make_encoder_field()
     critic_encoder_factory: EncoderFactory = make_encoder_field()
     imitator_encoder_factory: EncoderFactory = make_encoder_field()
-    q_func_factory: QFunctionFactory = make_q_func_field()
     batch_size: int = 100
     gamma: float = 0.99
     tau: float = 0.005
@@ -195,7 +196,7 @@ def inner_create_impl(
             observation_shape,
             action_size,
             self._config.critic_encoder_factory,
-            self._config.q_func_factory,
+            MeanQFunctionFactory(),
             n_ensembles=self._config.n_critics,
             device=self._device,
             enable_ddp=self._enable_ddp,
@@ -204,7 +205,7 @@ def inner_create_impl(
             observation_shape,
             action_size,
             self._config.critic_encoder_factory,
-            self._config.q_func_factory,
+            MeanQFunctionFactory(),
             n_ensembles=self._config.n_critics,
             device=self._device,
             enable_ddp=self._enable_ddp,

d3rlpy/algos/qlearning/bear.py

@@ -11,7 +11,7 @@
     create_vae_encoder,
 )
 from ...models.encoders import EncoderFactory, make_encoder_field
-from ...models.q_functions import QFunctionFactory, make_q_func_field
+from ...models.q_functions import MeanQFunctionFactory
 from ...optimizers.optimizers import OptimizerFactory, make_optimizer_field
 from ...types import Shape
 from .base import QLearningAlgoBase
@@ -90,8 +90,6 @@ class BEARConfig(LearnableConfig):
             Encoder factory for the critic.
         imitator_encoder_factory (d3rlpy.models.encoders.EncoderFactory):
             Encoder factory for the behavior policy.
-        q_func_factory (d3rlpy.models.q_functions.QFunctionFactory):
-            Q function factory.
         batch_size (int): Mini-batch size.
         gamma (float): Discount factor.
         tau (float): Target network synchronization coefficiency.
@@ -130,7 +128,6 @@ class BEARConfig(LearnableConfig):
     actor_encoder_factory: EncoderFactory = make_encoder_field()
     critic_encoder_factory: EncoderFactory = make_encoder_field()
     imitator_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
@@ -172,7 +169,7 @@ def inner_create_impl(
             observation_shape,
             action_size,
             self._config.critic_encoder_factory,
-            self._config.q_func_factory,
+            MeanQFunctionFactory(),
             n_ensembles=self._config.n_critics,
             device=self._device,
             enable_ddp=self._enable_ddp,
@@ -181,7 +178,7 @@ def inner_create_impl(
             observation_shape,
             action_size,
             self._config.critic_encoder_factory,
-            self._config.q_func_factory,
+            MeanQFunctionFactory(),
             n_ensembles=self._config.n_critics,
             device=self._device,
             enable_ddp=self._enable_ddp,

d3rlpy/algos/qlearning/torch/bcq_impl.py

@@ -7,6 +7,7 @@
 
 from ....models.torch import (
     ActionOutput,
+    TargetOutput,
     CategoricalPolicy,
     ContinuousEnsembleQFunctionForwarder,
     DeterministicResidualPolicy,
@@ -196,18 +197,17 @@ def inner_predict_best_action(self, x: TorchObservation) -> torch.Tensor:
     def inner_sample_action(self, x: TorchObservation) -> torch.Tensor:
         return self.inner_predict_best_action(x)
 
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         # TODO: this seems to be slow with image observation
         with torch.no_grad():
             repeated_x = self._repeat_observation(batch.next_observations)
             actions = self._sample_repeated_action(repeated_x, True)
-            values = compute_max_with_n_actions(
+            return compute_max_with_n_actions(
                 batch.next_observations,
                 actions,
                 self._targ_q_func_forwarder,
                 self._lam,
             )
-            return values
 
     def update_actor_target(self) -> None:
         soft_sync(self._modules.targ_policy, self._modules.policy, self._tau)

d3rlpy/algos/qlearning/torch/bear_impl.py

@@ -5,6 +5,7 @@
 
 from ....models.torch import (
     ActionOutput,
+    TargetOutput,
     ContinuousEnsembleQFunctionForwarder,
     Parameter,
     VAEDecoder,
@@ -245,7 +246,7 @@ def _compute_mmd(self, x: TorchObservation) -> torch.Tensor:
 
         return (mmd + 1e-6).sqrt().view(-1, 1)
 
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         with torch.no_grad():
             # BCQ-like target computation
             dist = build_squashed_gaussian_distribution(
@@ -254,7 +255,7 @@ def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
             actions, log_probs = dist.sample_n_with_log_prob(
                 self._n_target_samples
             )
-            values, indices = compute_max_with_n_actions_and_indices(
+            target, indices = compute_max_with_n_actions_and_indices(
                 batch.next_observations,
                 actions,
                 self._targ_q_func_forwarder,
@@ -266,7 +267,7 @@ def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
             max_log_prob = log_probs[torch.arange(batch_size), indices]
 
             log_temp = get_parameter(self._modules.log_temp)
-            return values - log_temp.exp() * max_log_prob
+            return TargetOutput(target.q_value - log_temp.exp() * max_log_prob)
 
     def inner_predict_best_action(self, x: TorchObservation) -> torch.Tensor:
         batch_size = (

d3rlpy/algos/qlearning/torch/ddpg_impl.py

@@ -11,6 +11,7 @@
     ActionOutput,
     ContinuousEnsembleQFunctionForwarder,
     Policy,
+    TargetOutput,
 )
 from ....optimizers.optimizers import OptimizerWrapper
 from ....torch_utility import (
@@ -111,13 +112,13 @@ def update_critic(self, batch: TorchMiniBatch) -> dict[str, float]:
         return asdict_as_float(loss)
 
     def compute_critic_loss(
-        self, batch: TorchMiniBatch, q_tpn: torch.Tensor
+        self, batch: TorchMiniBatch, target: TargetOutput
     ) -> DDPGBaseCriticLoss:
         loss = self._q_func_forwarder.compute_error(
             observations=batch.observations,
             actions=batch.actions,
             rewards=batch.rewards,
-            target=q_tpn,
+            target=target,
             terminals=batch.terminals,
             gamma=self._gamma**batch.intervals,
         )
@@ -153,7 +154,7 @@ def compute_actor_loss(
         pass
 
     @abstractmethod
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         pass
 
     def inner_predict_best_action(self, x: TorchObservation) -> torch.Tensor:
@@ -224,7 +225,7 @@ def compute_actor_loss(
         )[0]
         return DDPGBaseActorLoss(-q_t.mean())
 
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         with torch.no_grad():
             action = self._modules.targ_policy(batch.next_observations)
             return self._targ_q_func_forwarder.compute_target(

d3rlpy/algos/qlearning/torch/sac_impl.py

@@ -16,6 +16,7 @@
     NormalPolicy,
     Parameter,
     Policy,
+    TargetOutput,
     build_squashed_gaussian_distribution,
     get_parameter,
 )
@@ -113,7 +114,7 @@ def update_temp(self, log_prob: torch.Tensor) -> torch.Tensor:
         self._modules.temp_optim.step()
         return loss
 
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         with torch.no_grad():
             dist = build_squashed_gaussian_distribution(
                 self._modules.policy(batch.next_observations)
@@ -125,7 +126,7 @@ def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
                 action,
                 reduction="min",
             )
-            return target - entropy
+            return target.add(-entropy)
 
     def inner_sample_action(self, x: TorchObservation) -> torch.Tensor:
         dist = build_squashed_gaussian_distribution(self._modules.policy(x))
@@ -213,7 +214,7 @@ def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
                 batch.next_observations
             )
             keepdims = True
-            if target.dim() == 3:
+            if target.q_value.dim() == 3:
                 entropy = entropy.unsqueeze(-1)
                 probs = probs.unsqueeze(-1)
                 keepdims = False

d3rlpy/algos/qlearning/torch/td3_impl.py

@@ -1,6 +1,6 @@
 import torch
 
-from ....models.torch import ContinuousEnsembleQFunctionForwarder
+from ....models.torch import ContinuousEnsembleQFunctionForwarder, TargetOutput
 from ....torch_utility import TorchMiniBatch
 from ....types import Shape
 from .ddpg_impl import DDPGImpl, DDPGModules
@@ -43,7 +43,7 @@ def __init__(
         self._target_smoothing_clip = target_smoothing_clip
         self._update_actor_interval = update_actor_interval
 
-    def compute_target(self, batch: TorchMiniBatch) -> torch.Tensor:
+    def compute_target(self, batch: TorchMiniBatch) -> TargetOutput:
         with torch.no_grad():
             action = self._modules.targ_policy(batch.next_observations)
             # smoothing target

d3rlpy/models/encoders.py

@@ -1,3 +1,4 @@
+import math
 from dataclasses import dataclass, field
 from typing import Optional, Union
 
@@ -9,10 +10,13 @@
     EncoderWithAction,
     PixelEncoder,
     PixelEncoderWithAction,
+    SimBaEncoder,
+    SimBaEncoderWithAction,
+    SimbaV2Encoder,
+    SimbaV2EncoderWithAction,
     VectorEncoder,
     VectorEncoderWithAction,
 )
-from .torch.encoders import SimBaEncoder, SimBaEncoderWithAction
 from .utility import create_activation
 
 __all__ = [
@@ -21,6 +25,7 @@
     "VectorEncoderFactory",
     "DefaultEncoderFactory",
     "SimBaEncoderFactory",
+    "SimbaV2EncoderFactory",
     "register_encoder_factory",
     "make_encoder_field",
 ]
@@ -315,6 +320,80 @@ def get_type() -> str:
         return "simba"
 
 
+@dataclass()
+class SimbaV2EncoderFactory(EncoderFactory):
+    """SimbaV2 encoder factory class.
+
+    This class implements SimbaV2 encoder architecture.
+
+    References:
+        * `Lee et al., Hyperspherical Normalization for Scalable Deep
+          Reinforcement Learning, <https://arxiv.org/abs/2502.15280>`_
+
+    Args:
+        feature_size (int): Feature unit size.
+        hidden_size (int): HIdden expansion layer unit size.
+        n_blocks (int): Number of SimBa blocks.
+    """
+
+    feature_size: int = 256
+    n_blocks: int = 1
+    c_shift: float = 3
+
+    def create(self, observation_shape: Shape) -> SimbaV2Encoder:
+        assert len(observation_shape) == 1
+        return SimbaV2Encoder(
+            observation_shape=cast_flat_shape(observation_shape),
+            hidden_size=self.feature_size,
+            n_blocks=self.n_blocks,
+            scaler_init=self.scaler_init,
+            scaler_scale=self.scaler_scale,
+            alpha_init=self.alpha_init,
+            alpha_scale=self.alpha_scale,
+            c_shift=self.c_shift,
+        )
+
+    def create_with_action(
+        self,
+        observation_shape: Shape,
+        action_size: int,
+        discrete_action: bool = False,
+    ) -> SimbaV2EncoderWithAction:
+        assert len(observation_shape) == 1
+        return SimbaV2EncoderWithAction(
+            observation_shape=cast_flat_shape(observation_shape),
+            action_size=action_size,
+            hidden_size=self.feature_size,
+            n_blocks=self.n_blocks,
+            scaler_init=self.scaler_init,
+            scaler_scale=self.scaler_scale,
+            alpha_init=self.alpha_init,
+            alpha_scale=self.alpha_scale,
+            c_shift=self.c_shift,
+            discrete_action=discrete_action,
+        )
+
+    @staticmethod
+    def get_type() -> str:
+        return "simba_v2"
+
+    @property
+    def scaler_init(self) -> float:
+        return math.sqrt(2 / self.feature_size)
+
+    @property
+    def scaler_scale(self) -> float:
+        return math.sqrt(2 / self.feature_size)
+
+    @property
+    def alpha_init(self) -> float:
+        return 1 / (self.n_blocks + 1)
+
+    @property
+    def alpha_scale(self) -> float:
+        return math.sqrt(1 / self.feature_size)
+
+
 register_encoder_factory, make_encoder_field = generate_config_registration(
     EncoderFactory, lambda: DefaultEncoderFactory()
 )
@@ -324,3 +403,4 @@ def get_type() -> str:
 register_encoder_factory(PixelEncoderFactory)
 register_encoder_factory(DefaultEncoderFactory)
 register_encoder_factory(SimBaEncoderFactory)
+register_encoder_factory(SimbaV2EncoderFactory)