feat(flatpak): change flatpak specs to cached properties

jumanji/environments/packing/flat_pack/env.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from functools import cached_property
 from typing import Optional, Sequence, Tuple
 
 import chex
@@ -34,7 +35,7 @@
 from jumanji.viewer import Viewer
 
 
-class FlatPack(Environment[State]):
+class FlatPack(Environment[State, specs.MultiDiscreteArray, Observation]):
 
     """The FlatPack environment with a configurable number of row and column blocks.
     Here the goal of an agent is to completely fill an empty grid by placing all
@@ -129,6 +130,7 @@ def __init__(
         self.viewer = viewer or FlatPackViewer(
             "FlatPack", self.num_blocks, render_mode="human"
         )
+        super().__init__()
 
     def __repr__(self) -> str:
         return (
@@ -141,7 +143,6 @@ def reset(
         self,
         key: chex.PRNGKey,
     ) -> Tuple[State, TimeStep[Observation]]:
-
         """Resets the environment.
 
         Args:
@@ -259,6 +260,7 @@ def close(self) -> None:
 
         self.viewer.close()
 
+    @cached_property
     def observation_spec(self) -> specs.Spec[Observation]:
         """Returns the observation spec of the environment.
 
@@ -307,6 +309,7 @@ def observation_spec(self) -> specs.Spec[Observation]:
             action_mask=action_mask,
         )
 
+    @cached_property
     def action_spec(self) -> specs.MultiDiscreteArray:
         """Specifications of the action expected by the `FlatPack` environment.
 

jumanji/environments/packing/flat_pack/env_test.py

@@ -28,7 +28,10 @@
     CellDenseReward,
 )
 from jumanji.environments.packing.flat_pack.types import State
-from jumanji.testing.env_not_smoke import check_env_does_not_smoke
+from jumanji.testing.env_not_smoke import (
+    check_env_does_not_smoke,
+    check_env_specs_does_not_smoke,
+)
 from jumanji.testing.pytrees import assert_is_jax_array_tree
 from jumanji.types import StepType, TimeStep
 
@@ -182,6 +185,11 @@ def test_flat_pack__does_not_smoke(flat_pack: FlatPack) -> None:
     check_env_does_not_smoke(flat_pack)
 
 
+def test_flat_pack__specs_does_not_smoke(flat_pack: FlatPack) -> None:
+    """Test that we can access specs without any errors."""
+    check_env_specs_does_not_smoke(flat_pack)
+
+
 def test_flat_pack__is_done(flat_pack: FlatPack, key: chex.PRNGKey) -> None:
     """Test that the is_done method works as expected."""
 

jumanji/environments/routing/sokoban/env_test.py

@@ -26,7 +26,10 @@
     SimpleSolveGenerator,
 )
 from jumanji.environments.routing.sokoban.types import State
-from jumanji.testing.env_not_smoke import check_env_does_not_smoke
+from jumanji.testing.env_not_smoke import (
+    check_env_does_not_smoke,
+    check_env_specs_does_not_smoke,
+)
 from jumanji.types import TimeStep
 
 
@@ -215,3 +218,8 @@ def test_sokoban__reward_function_solved(sokoban_simple: Sokoban) -> None:
 def test_sokoban__does_not_smoke(sokoban: Sokoban) -> None:
     """Test that we can run an episode without any errors."""
     check_env_does_not_smoke(sokoban)
+
+
+def test_sokoban__specs_does_not_smoke(sokoban: Sokoban) -> None:
+    """Test that we can access specs without any errors."""
+    check_env_specs_does_not_smoke(sokoban)

jumanji/training/networks/flat_pack/actor_critic.py

@@ -40,7 +40,7 @@ def make_actor_critic_networks_flat_pack(
     hidden_size: int,
 ) -> ActorCriticNetworks:
     """Make actor-critic networks for the `FlatPack` environment."""
-    num_values = np.asarray(flat_pack.action_spec().num_values)
+    num_values = np.asarray(flat_pack.action_spec.num_values)
     parametric_action_distribution = FactorisedActionSpaceParametricDistribution(
         action_spec_num_values=num_values
     )
@@ -171,7 +171,6 @@ def __call__(self, observation: Observation) -> Tuple[chex.Array, chex.Array]:
         )  # (B, model_size), (B, num_rows-2, num_cols-2, hidden_size)
 
         for block_id in range(self.num_transformer_layers):
-
             (
                 self_attention_mask,  # (B, 1, num_blocks, num_blocks)
                 cross_attention_mask,  # (B, 1, num_blocks, 1)

jumanji/training/networks/flat_pack/random.py

@@ -21,7 +21,7 @@
 
 def make_random_policy_flat_pack(flat_pack: FlatPack) -> RandomPolicy:
     """Make random policy for FlatPack."""
-    action_spec_num_values = flat_pack.action_spec().num_values
+    action_spec_num_values = flat_pack.action_spec.num_values
 
     return make_masked_categorical_random_ndim(
         action_spec_num_values=action_spec_num_values