[Fix] Signatures of losses

Signed-off-by: Matteo Bettini <[email protected]>
facebookresearch · Oct 4, 2023 · 132e1b8 · 132e1b8
1 parent f46aaeb
commit 132e1b8
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 18 deletions.
diff --git a/benchmarl/algorithms/iddpg.py b/benchmarl/algorithms/iddpg.py
@@ -12,7 +12,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import AdditiveGaussianWrapper, ProbabilisticActor, TanhDelta
-from torchrl.objectives import ClipPPOLoss, DDPGLoss, LossModule, ValueEstimators
+from torchrl.objectives import DDPGLoss, LossModule, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -74,7 +74,7 @@ def _get_loss(
                 "Iddpg is not compatible with discrete actions yet"
             )
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
         return {
             "loss_actor": list(loss.actor_network_params.flatten_keys().values()),
             "loss_value": list(loss.value_network_params.flatten_keys().values()),

diff --git a/benchmarl/algorithms/iql.py b/benchmarl/algorithms/iql.py
@@ -11,7 +11,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import EGreedyModule, QValueModule
-from torchrl.objectives import ClipPPOLoss, DQNLoss, LossModule, ValueEstimators
+from torchrl.objectives import DQNLoss, LossModule, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -70,7 +70,7 @@ def _get_loss(
 
             return loss_module, True
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
         return {"loss": loss.parameters()}
 
     def _get_policy_for_loss(

diff --git a/benchmarl/algorithms/maddpg.py b/benchmarl/algorithms/maddpg.py
@@ -12,7 +12,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import AdditiveGaussianWrapper, ProbabilisticActor, TanhDelta
-from torchrl.objectives import ClipPPOLoss, DDPGLoss, LossModule, ValueEstimators
+from torchrl.objectives import DDPGLoss, LossModule, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -75,7 +75,7 @@ def _get_loss(
                 "MADDPG is not compatible with discrete actions yet"
             )
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
 
         return {
             "loss_actor": list(loss.actor_network_params.flatten_keys().values()),

diff --git a/benchmarl/algorithms/masac.py b/benchmarl/algorithms/masac.py
@@ -13,13 +13,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import MaskedCategorical, ProbabilisticActor, TanhNormal
-from torchrl.objectives import (
-    ClipPPOLoss,
-    DiscreteSACLoss,
-    LossModule,
-    SACLoss,
-    ValueEstimators,
-)
+from torchrl.objectives import DiscreteSACLoss, LossModule, SACLoss, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -126,7 +120,7 @@ def _get_loss(
 
         return loss_module, True
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
         return {
             "loss_actor": list(loss.actor_network_params.flatten_keys().values()),
             "loss_qvalue": list(loss.qvalue_network_params.flatten_keys().values()),

diff --git a/benchmarl/algorithms/qmix.py b/benchmarl/algorithms/qmix.py
@@ -11,7 +11,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import EGreedyModule, QMixer, QValueModule
-from torchrl.objectives import ClipPPOLoss, LossModule, QMixerLoss, ValueEstimators
+from torchrl.objectives import LossModule, QMixerLoss, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -75,7 +75,7 @@ def _get_loss(
 
             return loss_module, True
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
         return {"loss": loss.parameters()}
 
     def _get_policy_for_loss(

diff --git a/benchmarl/algorithms/vdn.py b/benchmarl/algorithms/vdn.py
@@ -11,7 +11,7 @@
 )
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.modules import EGreedyModule, QValueModule, VDNMixer
-from torchrl.objectives import ClipPPOLoss, LossModule, QMixerLoss, ValueEstimators
+from torchrl.objectives import LossModule, QMixerLoss, ValueEstimators
 
 from benchmarl.algorithms.common import Algorithm, AlgorithmConfig
 from benchmarl.models.common import ModelConfig
@@ -72,7 +72,7 @@ def _get_loss(
 
             return loss_module, True
 
-    def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
+    def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
 
         return {
             "loss": loss.parameters(),