Use functools.partial instead of lambdas for lax.cond branches (#245)

jumanji/environments/logic/game_2048/env.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import cached_property
+from functools import cached_property, partial
 from typing import Optional, Sequence, Tuple
 
 import chex
@@ -217,16 +217,10 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
         extras = {"highest_tile": highest_tile}
         timestep = jax.lax.cond(
             done,
-            lambda: termination(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
-            lambda: transition(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
+            partial(termination, extras=extras),
+            partial(transition, extras=extras),
+            reward,
+            observation,
         )
 
         return state, timestep

jumanji/environments/logic/sliding_tile_puzzle/env.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import cached_property
+from functools import cached_property, partial
 from typing import Dict, Optional, Sequence, Tuple
 
 import chex
@@ -144,16 +144,10 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
 
         timestep = jax.lax.cond(
             done | (next_state.step_count >= self.time_limit),
-            lambda: termination(
-                reward=reward,
-                observation=obs,
-                extras=extras,
-            ),
-            lambda: transition(
-                reward=reward,
-                observation=obs,
-                extras=extras,
-            ),
+            partial(termination, extras=extras),
+            partial(transition, extras=extras),
+            reward,
+            obs,
         )
 
         return next_state, timestep

jumanji/environments/packing/bin_pack/env.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import itertools
-from functools import cached_property
+from functools import cached_property, partial
 from typing import Dict, Optional, Sequence, Tuple
 
 import chex
@@ -334,16 +334,10 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
 
         timestep = jax.lax.cond(
             done,
-            lambda: termination(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
-            lambda: transition(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
+            partial(termination, extras=extras),
+            partial(transition, extras=extras),
+            reward,
+            observation,
         )
 
         return next_state, timestep

jumanji/environments/routing/sokoban/env.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from functools import cached_property
+from functools import cached_property, partial
 from typing import Dict, Optional, Sequence, Tuple
 
 import chex
@@ -243,16 +243,10 @@ def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observ
 
         timestep = jax.lax.cond(
             done,
-            lambda: termination(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
-            lambda: transition(
-                reward=reward,
-                observation=observation,
-                extras=extras,
-            ),
+            partial(termination, extras=extras),
+            partial(transition, extras=extras),
+            reward,
+            observation,
         )
 
         return next_state, timestep