instadeepai · clement-bonnet · Mar 13, 2024 · May 21, 2023 · May 22, 2023 · May 22, 2023
diff --git a/README.md b/README.md
@@ -98,8 +98,9 @@ problems.
 | 🎨 GraphColoring                              | Logic  | `GraphColoring-v0`                                   | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/graph_coloring/)   | [doc](https://instadeepai.github.io/jumanji/environments/graph_coloring/)   |
 | 💣 Minesweeper                           | Logic    | `Minesweeper-v0`                                     | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/minesweeper/) | [doc](https://instadeepai.github.io/jumanji/environments/minesweeper/) |
 | 🎲 RubiksCube                            | Logic    | `RubiksCube-v0`<br/>`RubiksCube-partly-scrambled-v0` | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/rubiks_cube/) | [doc](https://instadeepai.github.io/jumanji/environments/rubiks_cube/) |
-| ✏️ Sudoku                       | Logic    | `Sudoku-v0` <br/>`Sudoku-very-easy-v0`               | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/sudoku/) | [doc](https://instadeepai.github.io/jumanji/environments/sudoku/) |
-| 📦 BinPack (3D BinPacking Problem)       | Packing  | `BinPack-v2`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/bin_pack/)  | [doc](https://instadeepai.github.io/jumanji/environments/bin_pack/)    |
+| ✏️ Sudoku                       | Logic    | `Sudoku-v0` <br/>`Sudoku-very-easy-v0`| [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/sudoku/) | [doc](https://instadeepai.github.io/jumanji/environments/sudoku/) |
+| 📦 BinPack (3D BinPacking Problem)       | Packing  | `BinPack-v1`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/bin_pack/)  | [doc](https://instadeepai.github.io/jumanji/environments/bin_pack/)    |
+| 🧩 FlatPack (2D Grid Filling Problem) | Packing  | `FlatPack-v0`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/flat_pack/)  | [doc](https://instadeepai.github.io/jumanji/environments/flat_pack/)    |
 | 🏭 JobShop (Job Shop Scheduling Problem) | Packing  | `JobShop-v0`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/job_shop/)  | [doc](https://instadeepai.github.io/jumanji/environments/job_shop/)    |
 | 🎒 Knapsack                              | Packing  | `Knapsack-v1`                                        | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/knapsack/)  | [doc](https://instadeepai.github.io/jumanji/environments/knapsack/)    |
 | ▒ Tetris                              | Packing  | `Tetris-v0`                                        | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/tetris/)  | [doc](https://instadeepai.github.io/jumanji/environments/tetris/)    |

diff --git a/docs/api/environments/flat_pack.md b/docs/api/environments/flat_pack.md
@@ -0,0 +1,8 @@
+::: jumanji.environments.packing.flat_pack.env.FlatPack
+    selection:
+      members:
+        - __init__
+        - reset
+        - step
+        - observation_spec
+        - action_spec
diff --git a/docs/env_anim/flat_pack.gif b/docs/env_anim/flat_pack.gif
diff --git a/docs/env_img/flat_pack.png b/docs/env_img/flat_pack.png
diff --git a/docs/environments/flat_pack.md b/docs/environments/flat_pack.md
@@ -0,0 +1,57 @@
+# FlatPack Environment
+
+<p align="center">
+        <img src="../env_anim/flat_pack.gif" width="500"/>
+</p>
+
+We provide here a Jax JIT-able implementation of a packing environment named _flat pack_. The goal of
+the agent is to place all the available blocks on an empty 2D grid.
+Each time an episode resets a new set of blocks is created and the grid is emptied. Blocks are randomly
+shuffled and rotated and all have shape (3, 3).
+
+## Observation
+The observation given to the agent gives a view of the current state of the grid as well as
+all blocks that can be placed.
+
+- `current_grid`: jax array (float32) of shape `(num_rows, num_cols)` with values in the range
+    `[0, num_blocks]` (corresponding to the number of each block). This grid will have zeros
+    where no blocks have been placed and numbers corresponding to each block where that particular
+    block has been placed.
+
+- `blocks`: jax array (float32) of shape `(num_blocks, 3, 3)` of all possible blocks in
+    that can fit in the current grid. These blocks are shuffled, rotated and will always have shape `(3, 3)`.
+
+- `action_mask`: jax array (bool) of shape `(num_blocks, 4, num_rows-2, num_cols-2)`, representing
+    which actions are possible given the current state of the grid. The first index indicates the
+    number of blocks associated with a given grid. The second index indicates the number of times a block may be rotated.
+    The third and fourth indices indicate the row and column coordinate of where a blocks top left-most corner may be placed
+    respectively. Blocks are placed by an agent by specifying the row and column coordinate on the grid where the top left corner
+    of the selected block should be placed. These values will always be `num_rows-2` and `num_cols-2`
+    respectively to make it impossible for an agent to place a block outside the current grid.
+
+
+## Action
+The action space is a `MultiDiscreteArray`, specifically a tuple of an index between 0 and `num_blocks - 1`,
+an index between 0 and 4 (since there are 4 possible rotations), an index between 0 and `num_rows-2`
+(the possible row coordinates for placing a block) and an index between 0 and `num_cols-2`
+(the possible column coordinates for placing a block). An action thus consists of four pieces of
+information:
+
+- Block to place,
+
+- Number of 90 degree rotations to make to a chosen block ({0, 90, 180, 270} degrees),
+
+- Row coordinate for placing the rotated block's top left corner,
+
+- Column coordinate for placing the rotated block's top left corner.
+
+
+## Reward
+The reward function is configurable, but by default is a fully dense reward giving the sum of the number of non-zero
+cells in a placed block normalised by the total number of cells in the grid at each timestep. The episode
+terminates if either the grid is filled or `num_blocks` steps have been taken by an agent.
+
+
+## Registered Versions 📖
+- `FlatPack-v0`, a flat pack environment grid with 11 rows and 11 columns containing 5 row blocks and 5 column blocks
+    for a total of 25 blocks that can be placed on the grid. This version has a dense reward.
diff --git a/examples/load_checkpoints.ipynb b/examples/load_checkpoints.ipynb
@@ -111,8 +111,11 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "collapsed": false
+   },
    "source": [
     "## Load configs"
    ]
@@ -194,6 +197,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -243,6 +247,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -279,6 +284,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [

diff --git a/jumanji/__init__.py b/jumanji/__init__.py
@@ -81,6 +81,10 @@
 # given in the observation.
 register(id="BinPack-v2", entry_point="jumanji.environments:BinPack")
 
+# 2D grid filling problem with 25 blocks, an 11x11 grid and a random grid generator.
+# The grid must be filled in `num_blocks` steps.
+register(id="FlatPack-v0", entry_point="jumanji.environments:FlatPack")
+
 # Job-shop scheduling problem with 20 jobs, 10 machines, at most
 # 8 operations per job, and a max operation duration of 6 timesteps.
 register(id="JobShop-v0", entry_point="jumanji.environments:JobShop")

diff --git a/jumanji/environments/__init__.py b/jumanji/environments/__init__.py
@@ -20,8 +20,9 @@
 from jumanji.environments.logic.minesweeper import Minesweeper
 from jumanji.environments.logic.rubiks_cube import RubiksCube
 from jumanji.environments.logic.sudoku import Sudoku
-from jumanji.environments.packing import bin_pack, job_shop, knapsack, tetris
+from jumanji.environments.packing import bin_pack, flat_pack, job_shop, knapsack, tetris
 from jumanji.environments.packing.bin_pack.env import BinPack
+from jumanji.environments.packing.flat_pack.env import FlatPack
 from jumanji.environments.packing.job_shop.env import JobShop
 from jumanji.environments.packing.knapsack.env import Knapsack
 from jumanji.environments.packing.tetris.env import Tetris

diff --git a/jumanji/environments/packing/flat_pack/__init__.py b/jumanji/environments/packing/flat_pack/__init__.py
@@ -0,0 +1,16 @@
+# Copyright 2022 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from jumanji.environments.packing.flat_pack.env import FlatPack
+from jumanji.environments.packing.flat_pack.types import Observation, State
diff --git a/jumanji/environments/packing/flat_pack/conftest.py b/jumanji/environments/packing/flat_pack/conftest.py
@@ -0,0 +1,162 @@
+# Copyright 2022 InstaDeep Ltd. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import chex
+import jax
+import jax.numpy as jnp
+import pytest
+
+
+@pytest.fixture
+def key() -> chex.PRNGKey:
+    """A determinstic key."""
+
+    return jax.random.PRNGKey(0)
+
+
+@pytest.fixture
+def block() -> chex.Array:
+    """A mock block for testing."""
+
+    return jnp.array(
+        [
+            [0, 1, 1],
+            [0, 1, 1],
+            [0, 0, 1],
+        ]
+    )
+
+
+@pytest.fixture
+def solved_grid() -> chex.Array:
+    """A mock solved grid for testing."""
+
+    return jnp.array(
+        [
+            [1, 1, 1, 2, 2],
+            [1, 1, 2, 2, 2],
+            [3, 1, 4, 4, 2],
+            [3, 3, 4, 4, 4],
+            [3, 3, 3, 4, 4],
+        ],
+    )
+
+
+@pytest.fixture
+def grid_with_block_one_placed() -> chex.Array:
+    """A grid with only block one placed."""
+
+    return jnp.array(
+        [
+            [1, 1, 1, 0, 0],
+            [1, 1, 0, 0, 0],
+            [0, 1, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0],
+        ],
+    )
+
+
+@pytest.fixture()
+def block_one_placed_at_0_0(grid_with_block_one_placed: chex.Array) -> chex.Array:
+    """A 2D array of zeros where block one has been placed with it left top-most
+    corner at position (0, 0).
+    """
+
+    return grid_with_block_one_placed
+
+
+@pytest.fixture()
+def block_one_placed_at_1_1(grid_with_block_one_placed: chex.Array) -> chex.Array:
+    """A 2D array of zeros where block one has been placed with it left top-most
+    corner at position (1, 1).
+    """
+
+    # Shift all elements in the array one down and one to the right
+    partially_placed_block = jnp.roll(grid_with_block_one_placed, shift=1, axis=0)
+    partially_placed_block = jnp.roll(partially_placed_block, shift=1, axis=1)
+
+    return partially_placed_block
+
+
+@pytest.fixture()
+def action_mask_with_block_1_placed() -> chex.Array:
+    """Action mask for a 4 piece grid where only block 1 has been placed with its
+    left top-most corner at (1, 1).
+    """
+
+    return jnp.array(
+        [
+            [
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+            ],
+            [
+                [[False, False, True], [False, False, True], [False, True, True]],
+                [[False, False, True], [False, True, True], [False, True, True]],
+                [[False, False, False], [False, False, True], [True, False, True]],
+                [[False, False, False], [False, False, True], [False, False, True]],
+            ],
+            [
+                [[False, False, False], [False, False, True], [True, False, True]],
+                [[False, False, False], [False, False, True], [False, False, True]],
+                [[False, False, False], [False, False, True], [False, False, True]],
+                [[False, False, True], [False, True, True], [True, True, True]],
+            ],
+            [
+                [[False, False, False], [False, False, True], [False, False, True]],
+                [[False, False, True], [False, False, True], [False, True, True]],
+                [[False, False, False], [False, False, True], [False, False, True]],
+                [[False, False, True], [False, False, True], [False, True, True]],
+            ],
+        ]
+    )
+
+
+@pytest.fixture()
+def action_mask_without_only_block_1_placed() -> chex.Array:
+    """Action mask for a 4 piece grid where only block 1 can be placed with its
+    left top-most corner at (1, 1).
+    """
+
+    return jnp.array(
+        [
+            [
+                [[True, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+            ],
+            [
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+            ],
+            [
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+            ],
+            [
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+                [[False, False, False], [False, False, False], [False, False, False]],
+            ],
+        ]
+    )