Modify OptimFactory interface to take named_modules instead of parame…

…ters

d3rlpy/algos/qlearning/awac.py

@@ -124,10 +124,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
 
         dummy_log_temp = Parameter(torch.zeros(1, 1))

d3rlpy/algos/qlearning/bc.py

@@ -88,7 +88,7 @@ def inner_create_impl(
             raise ValueError(f"invalid policy_type: {self._config.policy_type}")
 
         optim = self._config.optim_factory.create(
-            imitator.parameters(), lr=self._config.learning_rate
+            imitator.named_modules(), lr=self._config.learning_rate
         )
 
         modules = BCModules(optim=optim, imitator=imitator)
@@ -159,7 +159,7 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            imitator.parameters(), lr=self._config.learning_rate
+            imitator.named_modules(), lr=self._config.learning_rate
         )
 
         modules = DiscreteBCModules(optim=optim, imitator=imitator)

d3rlpy/algos/qlearning/bcq.py

@@ -211,13 +211,13 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.imitator_optim_factory.create(
-            imitator.parameters(), lr=self._config.imitator_learning_rate
+            imitator.named_modules(), lr=self._config.imitator_learning_rate
         )
 
         modules = BCQModules(
@@ -373,16 +373,10 @@ def inner_create_impl(
                 device=self._device,
             )
 
-        # TODO: replace this with a cleaner way
-        # retrieve unique elements
-        q_func_params = list(q_funcs.parameters())
-        imitator_params = list(imitator.parameters())
-        unique_dict = {}
-        for param in q_func_params + imitator_params:
-            unique_dict[param] = param
-        unique_params = list(unique_dict.values())
+        q_func_params = list(q_funcs.named_modules())
+        imitator_params = list(imitator.named_modules())
         optim = self._config.optim_factory.create(
-            unique_params, lr=self._config.learning_rate
+            q_func_params + imitator_params, lr=self._config.learning_rate
         )
 
         modules = DiscreteBCQModules(

d3rlpy/algos/qlearning/bear.py

@@ -197,19 +197,19 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.imitator_optim_factory.create(
-            imitator.parameters(), lr=self._config.imitator_learning_rate
+            imitator.named_modules(), lr=self._config.imitator_learning_rate
         )
         temp_optim = self._config.temp_optim_factory.create(
-            log_temp.parameters(), lr=self._config.temp_learning_rate
+            log_temp.named_modules(), lr=self._config.temp_learning_rate
         )
         alpha_optim = self._config.alpha_optim_factory.create(
-            log_alpha.parameters(), lr=self._config.actor_learning_rate
+            log_alpha.named_modules(), lr=self._config.actor_learning_rate
         )
 
         modules = BEARModules(

d3rlpy/algos/qlearning/cql.py

@@ -166,20 +166,20 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), 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
+                log_temp.named_modules(), lr=self._config.temp_learning_rate
             )
         else:
             temp_optim = None
         if self._config.alpha_learning_rate > 0:
             alpha_optim = self._config.alpha_optim_factory.create(
-                log_alpha.parameters(), lr=self._config.alpha_learning_rate
+                log_alpha.named_modules(), lr=self._config.alpha_learning_rate
             )
         else:
             alpha_optim = None
@@ -296,7 +296,7 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            q_funcs.parameters(), lr=self._config.learning_rate
+            q_funcs.named_modules(), lr=self._config.learning_rate
         )
 
         modules = DQNModules(

d3rlpy/algos/qlearning/crr.py

@@ -162,10 +162,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
 
         modules = CRRModules(

d3rlpy/algos/qlearning/ddpg.py

@@ -123,10 +123,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
 
         modules = DDPGModules(

d3rlpy/algos/qlearning/dqn.py

@@ -84,7 +84,7 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            q_funcs.parameters(), lr=self._config.learning_rate
+            q_funcs.named_modules(), lr=self._config.learning_rate
         )
 
         modules = DQNModules(
@@ -186,7 +186,7 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            q_funcs.parameters(), lr=self._config.learning_rate
+            q_funcs.named_modules(), lr=self._config.learning_rate
         )
 
         modules = DQNModules(

d3rlpy/algos/qlearning/iql.py

@@ -140,10 +140,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
-        q_func_params = list(q_funcs.parameters())
-        v_func_params = list(value_func.parameters())
+        q_func_params = list(q_funcs.named_modules())
+        v_func_params = list(value_func.named_modules())
         critic_optim = self._config.critic_optim_factory.create(
             q_func_params + v_func_params, lr=self._config.critic_learning_rate
         )

d3rlpy/algos/qlearning/nfq.py

@@ -86,7 +86,7 @@ def inner_create_impl(
         )
 
         optim = self._config.optim_factory.create(
-            q_funcs.parameters(), lr=self._config.learning_rate
+            q_funcs.named_modules(), lr=self._config.learning_rate
         )
 
         modules = DQNModules(

d3rlpy/algos/qlearning/plas.py

@@ -146,13 +146,13 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.critic_optim_factory.create(
-            imitator.parameters(), lr=self._config.imitator_learning_rate
+            imitator.named_modules(), lr=self._config.imitator_learning_rate
         )
 
         modules = PLASModules(
@@ -296,16 +296,16 @@ def inner_create_impl(
             device=self._device,
         )
 
-        parameters = list(policy.parameters())
-        parameters += list(perturbation.parameters())
+        named_modules = list(policy.named_modules())
+        named_modules += list(perturbation.named_modules())
         actor_optim = self._config.actor_optim_factory.create(
-            params=parameters, lr=self._config.actor_learning_rate
+            named_modules, lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
         imitator_optim = self._config.critic_optim_factory.create(
-            imitator.parameters(), lr=self._config.imitator_learning_rate
+            imitator.named_modules(), lr=self._config.imitator_learning_rate
         )
 
         modules = PLASWithPerturbationModules(

d3rlpy/algos/qlearning/sac.py

@@ -151,14 +151,14 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), 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
+                log_temp.named_modules(), lr=self._config.temp_learning_rate
             )
         else:
             temp_optim = None
@@ -299,14 +299,14 @@ def inner_create_impl(
         )
 
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_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
+                log_temp.named_modules(), lr=self._config.temp_learning_rate
             )
         else:
             temp_optim = None

d3rlpy/algos/qlearning/td3.py

@@ -132,10 +132,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
 
         modules = DDPGModules(

d3rlpy/algos/qlearning/td3_plus_bc.py

@@ -124,10 +124,10 @@ def inner_create_impl(
         )
 
         actor_optim = self._config.actor_optim_factory.create(
-            policy.parameters(), lr=self._config.actor_learning_rate
+            policy.named_modules(), lr=self._config.actor_learning_rate
         )
         critic_optim = self._config.critic_optim_factory.create(
-            q_funcs.parameters(), lr=self._config.critic_learning_rate
+            q_funcs.named_modules(), lr=self._config.critic_learning_rate
         )
 
         modules = DDPGModules(

d3rlpy/algos/transformer/decision_transformer.py

@@ -114,7 +114,7 @@ def inner_create_impl(
             device=self._device,
         )
         optim = self._config.optim_factory.create(
-            transformer.parameters(), lr=self._config.learning_rate
+            transformer.named_modules(), lr=self._config.learning_rate
         )
         scheduler = torch.optim.lr_scheduler.LambdaLR(
             optim, lambda steps: min((steps + 1) / self._config.warmup_steps, 1)
@@ -234,7 +234,7 @@ def inner_create_impl(
             device=self._device,
         )
         optim = self._config.optim_factory.create(
-            transformer.parameters(), lr=self._config.learning_rate
+            transformer.named_modules(), lr=self._config.learning_rate
         )
         # JIT compile
         if self._config.compile: