From 1ec880828a6bc72e31f8343f0dd4c924b506d945 Mon Sep 17 00:00:00 2001
From: Wiem Khlifi <w.khlifi@instadeep.com>
Date: Mon, 25 Nov 2024 16:07:24 +0100
Subject: [PATCH] fix: termination condition (#266)

---
 jumanji/wrappers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jumanji/wrappers.py b/jumanji/wrappers.py
index 82f757ef6..1c9b3d65c 100644
--- a/jumanji/wrappers.py
+++ b/jumanji/wrappers.py
@@ -617,7 +617,7 @@ def step(
         ) -> Tuple[State, Observation, chex.Array, chex.Array, chex.Array, Optional[Any]]:
             """Step function of a Jumanji environment to be jitted."""
             state, timestep = self._env.step(state, action)
-            term = timestep.discount.astype(bool)
+            term = ~timestep.discount.astype(bool)
             trunc = timestep.last().astype(bool)
             return state, timestep.observation, timestep.reward, term, trunc, timestep.extras