Simpify Q functions with forwarders

takuseno · Aug 14, 2023 · 3ec25d7 · 3ec25d7
1 parent b8263d4
commit 3ec25d7
Show file tree

Hide file tree

Showing 44 changed files with 1,614 additions and 682 deletions.
diff --git a/d3rlpy/algos/qlearning/awac.py b/d3rlpy/algos/qlearning/awac.py
@@ -106,7 +106,15 @@ def inner_create_impl(
             use_std_parameter=True,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        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,
@@ -119,13 +127,16 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
 
         self._impl = AWACImpl(
             observation_shape=observation_shape,
             action_size=action_size,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             policy=policy,
             actor_optim=actor_optim,
             critic_optim=critic_optim,

diff --git a/d3rlpy/algos/qlearning/base.py b/d3rlpy/algos/qlearning/base.py
@@ -14,6 +14,7 @@
 
 import numpy as np
 import torch
+from torch import nn
 from tqdm.auto import tqdm, trange
 from typing_extensions import Self
 
@@ -36,7 +37,7 @@
     LoggerAdapterFactory,
 )
 from ...metrics import EvaluatorProtocol, evaluate_qlearning_with_environment
-from ...models.torch import EnsembleQFunction, Policy
+from ...models.torch import Policy
 from ...torch_utility import (
     TorchMiniBatch,
     convert_to_torch,
@@ -119,15 +120,15 @@ def copy_policy_optim_from(self, impl: "QLearningAlgoImplBase") -> None:
         sync_optimizer_state(self.policy_optim, impl.policy_optim)
 
     @property
-    def q_function(self) -> EnsembleQFunction:
+    def q_function(self) -> nn.ModuleList:
         raise NotImplementedError
 
     def copy_q_function_from(self, impl: "QLearningAlgoImplBase") -> None:
-        q_func = self.q_function.q_funcs[0]
-        if not isinstance(impl.q_function.q_funcs[0], type(q_func)):
+        q_func = self.q_function[0]
+        if not isinstance(impl.q_function[0], type(q_func)):
             raise ValueError(
                 f"Invalid Q-function type: expected={type(q_func)},"
-                f"actual={type(impl.q_function.q_funcs[0])}"
+                f"actual={type(impl.q_function[0])}"
             )
         hard_sync(self.q_function, impl.q_function)
 

diff --git a/d3rlpy/algos/qlearning/bcq.py b/d3rlpy/algos/qlearning/bcq.py
@@ -174,7 +174,15 @@ def inner_create_impl(
             self._config.actor_encoder_factory,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        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,
@@ -196,7 +204,7 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.imitator_optim_factory.create(
             imitator.parameters(), lr=self._config.imitator_learning_rate
@@ -206,7 +214,10 @@ def inner_create_impl(
             observation_shape=observation_shape,
             action_size=action_size,
             policy=policy,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             imitator=imitator,
             actor_optim=actor_optim,
             critic_optim=critic_optim,
@@ -323,7 +334,15 @@ class DiscreteBCQ(QLearningAlgoBase[DiscreteBCQImpl, DiscreteBCQConfig]):
     def inner_create_impl(
         self, observation_shape: Shape, action_size: int
     ) -> None:
-        q_func = create_discrete_q_function(
+        q_funcs, q_func_forwarder = create_discrete_q_function(
+            observation_shape,
+            action_size,
+            self._config.encoder_factory,
+            self._config.q_func_factory,
+            n_ensembles=self._config.n_critics,
+            device=self._device,
+        )
+        targ_q_funcs, targ_q_func_forwarder = create_discrete_q_function(
             observation_shape,
             action_size,
             self._config.encoder_factory,
@@ -333,14 +352,14 @@ def inner_create_impl(
         )
 
         # share convolutional layers if observation is pixel
-        if isinstance(q_func.q_funcs[0].encoder, PixelEncoder):
+        if isinstance(q_funcs[0].encoder, PixelEncoder):
             hidden_size = compute_output_size(
                 [observation_shape],
-                q_func.q_funcs[0].encoder,
+                q_funcs[0].encoder,
                 device=self._device,
             )
             imitator = CategoricalPolicy(
-                encoder=q_func.q_funcs[0].encoder,
+                encoder=q_funcs[0].encoder,
                 hidden_size=hidden_size,
                 action_size=action_size,
             )
@@ -355,7 +374,7 @@ def inner_create_impl(
 
         # TODO: replace this with a cleaner way
         # retrieve unique elements
-        q_func_params = list(q_func.parameters())
+        q_func_params = list(q_funcs.parameters())
         imitator_params = list(imitator.parameters())
         unique_dict = {}
         for param in q_func_params + imitator_params:
@@ -368,7 +387,10 @@ def inner_create_impl(
         self._impl = DiscreteBCQImpl(
             observation_shape=observation_shape,
             action_size=action_size,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             imitator=imitator,
             optim=optim,
             gamma=self._config.gamma,

diff --git a/d3rlpy/algos/qlearning/bear.py b/d3rlpy/algos/qlearning/bear.py
@@ -164,7 +164,15 @@ def inner_create_impl(
             self._config.actor_encoder_factory,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        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,
@@ -194,7 +202,7 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.imitator_optim_factory.create(
             imitator.parameters(), lr=self._config.imitator_learning_rate
@@ -210,7 +218,10 @@ def inner_create_impl(
             observation_shape=observation_shape,
             action_size=action_size,
             policy=policy,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             imitator=imitator,
             log_temp=log_temp,
             log_alpha=log_alpha,

diff --git a/d3rlpy/algos/qlearning/cql.py b/d3rlpy/algos/qlearning/cql.py
@@ -141,7 +141,15 @@ def inner_create_impl(
             self._config.actor_encoder_factory,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        q_funcs, q_func_fowarder = 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,
@@ -162,7 +170,7 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
         temp_optim = self._config.temp_optim_factory.create(
             log_temp.parameters(), lr=self._config.temp_learning_rate
@@ -175,7 +183,10 @@ def inner_create_impl(
             observation_shape=observation_shape,
             action_size=action_size,
             policy=policy,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_fowarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             log_temp=log_temp,
             log_alpha=log_alpha,
             actor_optim=actor_optim,
@@ -279,7 +290,15 @@ class DiscreteCQL(QLearningAlgoBase[DiscreteCQLImpl, DiscreteCQLConfig]):
     def inner_create_impl(
         self, observation_shape: Shape, action_size: int
     ) -> None:
-        q_func = create_discrete_q_function(
+        q_funcs, q_func_forwarder = create_discrete_q_function(
+            observation_shape,
+            action_size,
+            self._config.encoder_factory,
+            self._config.q_func_factory,
+            n_ensembles=self._config.n_critics,
+            device=self._device,
+        )
+        targ_q_funcs, targ_q_func_forwarder = create_discrete_q_function(
             observation_shape,
             action_size,
             self._config.encoder_factory,
@@ -289,13 +308,16 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            q_func.parameters(), lr=self._config.learning_rate
+            q_funcs.parameters(), lr=self._config.learning_rate
         )
 
         self._impl = DiscreteCQLImpl(
             observation_shape=observation_shape,
             action_size=action_size,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             optim=optim,
             gamma=self._config.gamma,
             alpha=self._config.alpha,

diff --git a/d3rlpy/algos/qlearning/crr.py b/d3rlpy/algos/qlearning/crr.py
@@ -140,7 +140,15 @@ def inner_create_impl(
             self._config.actor_encoder_factory,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        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,
@@ -153,14 +161,17 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
 
         self._impl = CRRImpl(
             observation_shape=observation_shape,
             action_size=action_size,
             policy=policy,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             actor_optim=actor_optim,
             critic_optim=critic_optim,
             gamma=self._config.gamma,

diff --git a/d3rlpy/algos/qlearning/ddpg.py b/d3rlpy/algos/qlearning/ddpg.py
@@ -101,7 +101,15 @@ def inner_create_impl(
             self._config.actor_encoder_factory,
             device=self._device,
         )
-        q_func = create_continuous_q_function(
+        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,
@@ -114,14 +122,17 @@ def inner_create_impl(
             policy.parameters(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_func.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.parameters(), lr=self._config.critic_learning_rate
         )
 
         self._impl = DDPGImpl(
             observation_shape=observation_shape,
             action_size=action_size,
             policy=policy,
-            q_func=q_func,
+            q_funcs=q_funcs,
+            q_func_forwarder=q_func_forwarder,
+            targ_q_funcs=targ_q_funcs,
+            targ_q_func_forwarder=targ_q_func_forwarder,
             actor_optim=actor_optim,
             critic_optim=critic_optim,
             gamma=self._config.gamma,