[Fix] EndOfLifeTransform fix in end of life detection (#1705)

pytorch · Nov 19, 2023 · 5cac16a · 5cac16a
1 parent 44bd026
commit 5cac16a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/envs/transforms/gym_transforms.py b/torchrl/envs/transforms/gym_transforms.py
@@ -148,7 +148,7 @@ def _step(self, tensordict, next_tensordict):
 
         lives = self._get_lives()
         end_of_life = torch.tensor(
-            tensordict.get(self.lives_key) < lives, device=self.parent.device
+            tensordict.get(self.lives_key) > lives, device=self.parent.device
         )
         try:
             done = next_tensordict.get(self.done_key)