feat(sudoku): implement environment

.pre-commit-config.yaml

@@ -74,3 +74,4 @@ repos:
         args:
           - --license-filepath=license_header.txt
           - --comment-style=#
+        exclude: .npy$

README.md

@@ -30,6 +30,7 @@
   <img src="docs/env_anim/rubiks_cube.gif" alt="RubiksCube" width="30%" />
   <img src="docs/env_anim/game_2048.gif" alt="Game2048" width="30%" />
   <img src="docs/env_anim/minesweeper.gif" alt="Minesweeper" width="30%" />
+  <img src="docs/env_anim/sudoku.gif" alt="Sudoku" width="30%" />
 </p>
 
 
@@ -81,6 +82,7 @@ problems.
 | 🔢 Game2048                              | Logic  | `Game2048-v0`                                        | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/logic/game_2048/)   | [doc](https://instadeepai.github.io/jumanji/environments/game_2048/)   |
 | 💣 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-v1`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/packing/bin_pack/)  | [doc](https://instadeepai.github.io/jumanji/environments/bin_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/)    |

docs/api/environments/sudoku.md

@@ -0,0 +1,8 @@
+::: jumanji.environments.logic.sudoku.env.Sudoku
+    selection:
+      members:
+        - __init__
+        - reset
+        - step
+        - observation_spec
+        - action_spec

docs/environments/sudoku.md

@@ -0,0 +1,43 @@
+# Sudoku Environment
+
+<p align="center">
+        <img src="../env_anim/sudoku.gif" width="500"/>
+</p>
+
+We provide here a Jax JIT-able implementation of the
+[Sudoku](https://en.wikipedia.org/wiki/Sudoku) puzzle game.
+
+
+## Observation
+The observation given to the agent consists of:
+
+
+- `board`: jax array (int32) of shape (9,9):
+    empty cells are represented by -1, and filled cells are represented by 0-8.
+- `action_mask`: jax array (bool) of shape (9,9,9):
+    indicates which actions are valid.
+
+
+
+## Action
+The action space is a `MultiDiscreteArray` of integer values representing coordinates of the square
+to explore and the digits to write in the cell, e.g. `[3, 6, 8]` for writing the digit `9`
+the cell located on the third row and sixth column.
+
+
+## Reward
+The reward is `1` at the end of the episode if the board is correctly solved, and `0` in every
+other case.
+
+
+## Registered Versions 📖
+- `Sudoku-v0`, the classic [game](https://en.wikipedia.org/wiki/Sudoku) on
+a 9x9 grid, 10000 random puzzles with mixed difficulty are included by default.
+- `Sudoku-very-easy-v0`, the classic [game](https://en.wikipedia.org/wiki/Sudoku) on
+a 9x9 grid, only 1000 very-easy random puzzles (>46 clues) included by default.
+
+## Using custom puzzle instances
+If one wants to include its own database of puzzles, the `DatabaseGenerator` can be initialized with any collection of puzzles using the argument `custom_boards`.
+Some references for databases of puzzle of various difficulties:
+ - https://www.kaggle.com/datasets/rohanrao/sudoku
+ - https://www.kaggle.com/datasets/informoney/4-million-sudoku-puzzles-easytohard

examples/training.ipynb

@@ -59,7 +59,7 @@
    },
    "outputs": [],
    "source": [
-    "env = \"maze\"  # @param ['bin_pack', 'cleaner', 'connector', 'cvrp', 'game_2048', 'job_shop', 'knapsack', 'maze', 'minesweeper', 'rubiks_cube', 'snake', 'tsp']\n",
+    "env = \"maze\"  # @param ['bin_pack', 'cleaner', 'connector', 'cvrp', 'game_2048', 'job_shop', 'knapsack', 'maze', 'minesweeper', 'rubiks_cube', 'snake', 'sudoku', 'tsp']\n",
     "agent = \"random\"  # @param ['random', 'a2c']"
    ]
   },
@@ -105,7 +105,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.15"
+   "version": "3.10.9"
   }
  },
  "nbformat": 4,

jumanji/__init__.py

@@ -14,6 +14,7 @@
 
 from jumanji.env import Environment
 from jumanji.environments.logic.rubiks_cube import generator as rubik_generator
+from jumanji.environments.logic.sudoku import generator as sudoku_generator
 from jumanji.registration import make, register, registered_environments
 from jumanji.version import __version__
 
@@ -42,6 +43,19 @@
     kwargs={"time_limit": 20, "generator": partly_scrambled_rubiks_cube_generator},
 )
 
+# Sudoku - the standard Sudoku puzzle with grid of size 9x9. By default 10000 puzzles
+# of mixed difficulties are set for reset.
+register(id="Sudoku-v0", entry_point="jumanji.environments:Sudoku")
+
+# Sudoku - the standard Sudoku puzzle with grid of size 9x9. 1000 puzzles of very-easy
+# difficulty (>46 clues) are set for reset.
+very_easy_sudoku_generator = sudoku_generator.DatabaseGenerator(level="very-easy")
+register(
+    id="Sudoku-very-easy-v0",
+    entry_point="jumanji.environments:Sudoku",
+    kwargs={"generator": very_easy_sudoku_generator},
+)
+
 
 ###
 # Packing Environments

jumanji/environments/__init__.py

@@ -18,6 +18,7 @@
 from jumanji.environments.logic.game_2048.env import Game2048
 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
 from jumanji.environments.packing.bin_pack.env import BinPack
 from jumanji.environments.packing.job_shop.env import JobShop

jumanji/environments/logic/sudoku/__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.logic.sudoku.env import Sudoku
+from jumanji.environments.logic.sudoku.types import Observation, State

jumanji/environments/logic/sudoku/conftest.py

@@ -0,0 +1,24 @@
+# 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 pytest
+
+from jumanji.environments.logic.sudoku.env import Sudoku
+from jumanji.environments.logic.sudoku.generator import DummyGenerator
+
+
+@pytest.fixture
+def sudoku_env() -> Sudoku:
+    """Fixture for a default sudoku environment."""
+    return Sudoku(generator=DummyGenerator())

jumanji/environments/logic/sudoku/constants.py

@@ -0,0 +1,28 @@
+# 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.
+
+BOARD_WIDTH: int = 9
+
+
+BOX_IDX = [
+    [0, 1, 2, 9, 10, 11, 18, 19, 20],
+    [27, 28, 29, 36, 37, 38, 45, 46, 47],
+    [54, 55, 56, 63, 64, 65, 72, 73, 74],
+    [3, 4, 5, 12, 13, 14, 21, 22, 23],
+    [30, 31, 32, 39, 40, 41, 48, 49, 50],
+    [57, 58, 59, 66, 67, 68, 75, 76, 77],
+    [6, 7, 8, 15, 16, 17, 24, 25, 26],
+    [33, 34, 35, 42, 43, 44, 51, 52, 53],
+    [60, 61, 62, 69, 70, 71, 78, 79, 80],
+]