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