[Feature] Introduced terminated

Signed-off-by: Matteo Bettini <[email protected]>

benchmarl/algorithms/iddpg.py

@@ -62,6 +62,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
             loss_module.make_value_estimator(
@@ -150,13 +151,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/ippo.py

@@ -79,6 +79,7 @@ def _get_loss(
             reward=(group, "reward"),
             action=(group, "action"),
             done=(group, "done"),
+            terminated=(group, "terminated"),
             advantage=(group, "advantage"),
             value_target=(group, "value_target"),
             value=(group, "state_value"),
@@ -199,13 +200,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/iql.py

@@ -59,6 +59,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 action=(group, "action"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
                 action_value=(group, "action_value"),
                 value=(group, "chosen_action_value"),
                 priority=(group, "td_error"),
@@ -155,13 +156,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/isac.py

@@ -93,6 +93,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
         else:
@@ -116,6 +117,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
         loss_module.make_value_estimator(
@@ -232,13 +234,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/maddpg.py

@@ -62,6 +62,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
             loss_module.make_value_estimator(
@@ -152,13 +153,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/mappo.py

@@ -78,6 +78,7 @@ def _get_loss(
             reward=(group, "reward"),
             action=(group, "action"),
             done=(group, "done"),
+            terminated=(group, "terminated"),
             advantage=(group, "advantage"),
             value_target=(group, "value_target"),
             value=(group, "state_value"),
@@ -199,13 +200,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/masac.py

@@ -93,6 +93,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
         else:
@@ -116,6 +117,7 @@ def _get_loss(
                 reward=(group, "reward"),
                 priority=(group, "td_error"),
                 done=(group, "done"),
+                terminated=(group, "terminated"),
             )
 
         loss_module.make_value_estimator(
@@ -233,13 +235,21 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         group_shape = batch.get(group).shape
 
         nested_done_key = ("next", group, "done")
+        nested_terminated_key = ("next", group, "terminated")
         nested_reward_key = ("next", group, "reward")
 
         if nested_done_key not in keys:
             batch.set(
                 nested_done_key,
                 batch.get(("next", "done")).unsqueeze(-1).expand((*group_shape, 1)),
             )
+        if nested_terminated_key not in keys:
+            batch.set(
+                nested_terminated_key,
+                batch.get(("next", "terminated"))
+                .unsqueeze(-1)
+                .expand((*group_shape, 1)),
+            )
 
         if nested_reward_key not in keys:
             batch.set(

benchmarl/algorithms/qmix.py

@@ -63,6 +63,7 @@ def _get_loss(
                 reward="reward",
                 action=(group, "action"),
                 done="done",
+                terminated="terminated",
                 action_value=(group, "action_value"),
                 local_value=(group, "chosen_action_value"),
                 global_value="chosen_action_value",
@@ -159,13 +160,19 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         keys = list(batch.keys(True, True))
 
         done_key = ("next", "done")
+        terminated_key = ("next", "terminated")
         reward_key = ("next", "reward")
 
         if done_key not in keys:
             batch.set(
                 done_key,
                 batch.get(("next", group, "done")).any(-2),
             )
+        if terminated_key not in keys:
+            batch.set(
+                terminated_key,
+                batch.get(("next", group, "terminated")).any(-2),
+            )
 
         if reward_key not in keys:
             batch.set(

benchmarl/algorithms/vdn.py

@@ -60,6 +60,7 @@ def _get_loss(
                 reward="reward",
                 action=(group, "action"),
                 done="done",
+                terminated="terminated",
                 action_value=(group, "action_value"),
                 local_value=(group, "chosen_action_value"),
                 global_value="chosen_action_value",
@@ -159,13 +160,19 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
         keys = list(batch.keys(True, True))
 
         done_key = ("next", "done")
+        terminated_key = ("next", "terminated")
         reward_key = ("next", "reward")
 
         if done_key not in keys:
             batch.set(
                 done_key,
                 batch.get(("next", group, "done")).any(-2),
             )
+        if terminated_key not in keys:
+            batch.set(
+                terminated_key,
+                batch.get(("next", group, "terminated")).any(-2),
+            )
 
         if reward_key not in keys:
             batch.set(