Merge pull request #161 from Replicable-MARL/fix_critic

fix critic paramters
Replicable-MARL · Jul 25, 2023 · 1ec9660 · 1ec9660
2 parents 521316d + 17ef5b3
commit 1ec9660
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/marllib/marl/models/zoo/mlp/base_mlp.py b/marllib/marl/models/zoo/mlp/base_mlp.py
@@ -80,6 +80,7 @@ def __init__(
         self.q_flag = False
 
         self.actors = [self.p_encoder, self.p_branch]
+        self.critics = [self.vf_encoder, self.vf_branch]
         self.actor_initialized_parameters = self.actor_parameters()
 
     @override(TorchModelV2)
@@ -121,4 +122,4 @@ def actor_parameters(self):
         return reduce(lambda x, y: x + y, map(lambda p: list(p.parameters()), self.actors))
 
     def critic_parameters(self):
-        return list(self.vf_branch.parameters())
+        return reduce(lambda x, y: x + y, map(lambda p: list(p.parameters()), self.critics))