instadeepai · djbyrne · Feb 6, 2024 · Jun 9, 2023 · Jun 12, 2023 · Jun 16, 2023
diff --git a/.gitignore b/.gitignore
@@ -154,3 +154,4 @@ cython_debug/
 jumanji_env/
 **/outputs/
 *.xml
+.sokoban_cache/
diff --git a/README.md b/README.md
@@ -47,6 +47,7 @@
   </div>
     <div class="row" align="center">
     <img src="docs/env_anim/pac_man.gif" alt="RobotWarehouse" width="16%">
+    <img src="docs/env_anim/sokoban.gif" alt="RobotWarehouse" width="16%">
   </div>
 </div>
 
@@ -112,6 +113,7 @@ problems.
 | 📬 TSP (Travelling Salesman Problem)           | Routing  | `TSP-v1`                                             | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/tsp/)       | [doc](https://instadeepai.github.io/jumanji/environments/tsp/)         |
 | Multi Minimum Spanning Tree Problem | Routing  | `MMST-v0`                                | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/mmst)    | [doc](https://instadeepai.github.io/jumanji/environments/mmst/)    |
 | ᗧ•••ᗣ•• PacMan   | Routing  | `PacMan-v0`                                            | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/pacman/)      | [doc](https://instadeepai.github.io/jumanji/environments/pacman/)
+| 👾 Sokoban                                                     | Routing  | `Sokoban-v0`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/sokoban/)          | [doc](https://instadeepai.github.io/jumanji/environments/sokoban/)         |
 
 <h2 name="install" id="install">Installation 🎬</h2>
 

diff --git a/docs/api/environments/sokoban.md b/docs/api/environments/sokoban.md
@@ -0,0 +1,9 @@
+::: jumanji.environments.routing.sokoban.env.Sokoban
+    selection:
+      members:
+        - __init__
+        - observation_spec
+        - action_spec
+        - reset
+        - step
+        - render
diff --git a/docs/env_anim/sokoban.gif b/docs/env_anim/sokoban.gif
diff --git a/docs/env_img/sokoban.png b/docs/env_img/sokoban.png
diff --git a/docs/environments/sokoban.md b/docs/environments/sokoban.md
@@ -0,0 +1,76 @@
+# Sokoban Environment 👾
+
+<p align="center">
+        <img src="../env_anim/sokoban.gif" width="500"/>
+</p>
+
+This is a Jax implementation of the _Sokoban_ puzzle, a dynamic box-pushing environment where the agent's goal is to place all boxes on their targets. This version follows the rules from the DeepMind paper on [Imagination Augmented Agents](https://arxiv.org/abs/1707.06203), with levels based on the Boxoban dataset from [Guez et al., 2018](https://github.com/deepmind/boxoban-levels)[[1]](#ref1). The graphical assets were taken from [gym-sokoban](https://github.com/mpSchrader/gym-sokoban) by Schrader, a diverse Sokoban library implementing many versions of the game in the OpenAI gym framework [[2]](#ref2).
+
+## Observation
+
+- `grid`: An Array (uint8) of shape `(10, 10, 2)`. It represents the variable grid (containing movable objects: boxes and the agent) and the fixed grid (containing fixed objects: walls and targets).
+- `step_count`: An Array (int32) of shape `()`, representing the current number of steps in the episode.
+
+## Object Encodings
+
+| Object       | Encoding |
+|--------------|----------|
+| Empty Space  | 0        |
+| Wall         | 1        |
+| Target       | 2        |
+| Agent        | 3        |
+| Box          | 4        |
+
+## Actions
+
+The agent's action space is an Array (int32) with potential values of `[0,1,2,3]` (corresponding to `[Up, Down, Left, Right]`). If the agent attempts to move into a wall, off the grid, or push a box into a wall or off the grid, the grid state remains unchanged; however, the step count is incremented by one. Chained box pushes are not allowed and will result in no action.
+
+## Reward
+
+The reward function comprises:
+- `-0.1` for each step taken in the environment.
+- `+1` for each box moved onto a target location and `-1` for each box moved off a target location.
+- `+10` upon successful placement of all four boxes on their targets.
+
+## Episode Termination
+
+The episode concludes when:
+- The step limit of 120 is reached.
+- All 4 boxes are placed on targets (i.e., the problem is solved).
+
+## Dataset
+
+The Boxoban dataset offers a collection of puzzle levels. Each level features four boxes and four targets. The dataset has three levels of difficulty: 'unfiltered', 'medium', and 'hard'.
+
+| Dataset Split | Number of Levels |
+|---------------|------------------|
+| Unfiltered (Training) | 900,000          |
+| Unfiltered (Validation) | 100,000           |
+| Unfiltered (Test) | 1000              |
+| Medium (Training) | 450,000          |
+| Medium (Validation) | 50,000           |
+| Hard | 3332             |
+
+
+The dataset generation procedure and more details can be found in Guez et al., 2018 <a href="#ref1">[1]</a>.
+
+## Graphics
+
+| Type             | Graphic                                                                                |
+|------------------|----------------------------------------------------------------------------------------|
+| Wall             | ![Wall](../../jumanji/environments/routing/sokoban/imgs/wall.png)                      |
+| Floor            | ![Floor](../../jumanji/environments/routing/sokoban/imgs/floor.png)                    |
+| Target    | ![BoxTarget](../../jumanji/environments/routing/sokoban/imgs/box_target.png)           |
+| Box on Target    | ![BoxTarget](../../jumanji/environments/routing/sokoban/imgs/box_on_target.png)        |
+| Box Off Target   | ![BoxOffTarget](../../jumanji/environments/routing/sokoban/imgs/box.png)               |
+| Agent Off Target | ![PlayerOffTarget](../../jumanji/environments/routing/sokoban/imgs/agent.png)          |
+| Agent On Target  | ![PlayerOnTarget](../../jumanji/environments/routing/sokoban/imgs/agent_on_target.png) |
+
+## Registered Versions 📖
+
+- `Sokoban-v0`: Sokoban game with levels generated using DeepMind Boxoban dataset (unfiltered train).
+
+## References
+<a id="ref1">[1]</a> Guez, A., Mirza, M., Gregor, K., Kabra, R., Racaniere, S., Weber, T., Raposo, D., Santoro, A., Orseau, L., Eccles, T., Wayne, G., Silver, D., Lillicrap, T., Valdes, V. (2018). An investigation of Model-free planning: boxoban levels. Available at [https://github.com/deepmind/boxoban-levels](https://github.com/deepmind/boxoban-levels)
+
+<a id="ref2">[2]</a> Schrader, M. (2018). Gym-sokoban. Available at [https://github.com/mpSchrader/gym-sokoban](https://github.com/mpSchrader/gym-sokoban)
diff --git a/jumanji/__init__.py b/jumanji/__init__.py
@@ -73,7 +73,6 @@
     kwargs={"generator": very_easy_sudoku_generator},
 )
 
-
 ###
 # Packing Environments
 ###
@@ -93,7 +92,6 @@
 # Tetris - the game of tetris with a grid size of 10x10 and a time limit of 400.
 register(id="Tetris-v0", entry_point="jumanji.environments:Tetris")
 
-
 ###
 # Routing Environments
 ###
@@ -128,5 +126,7 @@
 # TSP with 20 randomly generated cities and a dense reward function.
 register(id="TSP-v1", entry_point="jumanji.environments:TSP")
 
+# Sokoban with deepmind dataset generator
+register(id="Sokoban-v0", entry_point="jumanji.environments:Sokoban")
 # Pacman - minimal version of Atarti Pacman game
 register(id="PacMan-v0", entry_point="jumanji.environments:PacMan")
diff --git a/jumanji/environments/__init__.py b/jumanji/environments/__init__.py
@@ -35,6 +35,7 @@
     pac_man,
     robot_warehouse,
     snake,
+    sokoban,
     tsp,
 )
 from jumanji.environments.routing.cleaner.env import Cleaner
@@ -46,6 +47,7 @@
 from jumanji.environments.routing.pac_man.env import PacMan
 from jumanji.environments.routing.robot_warehouse.env import RobotWarehouse
 from jumanji.environments.routing.snake.env import Snake
+from jumanji.environments.routing.sokoban.env import Sokoban
 from jumanji.environments.routing.tsp.env import TSP
 
 

diff --git a/jumanji/environments/logic/game_2048/env.py b/jumanji/environments/logic/game_2048/env.py
@@ -66,7 +66,7 @@ class Game2048(Environment[State]):
     ```python
     from jumanji.environments import Game2048
     env = Game2048()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/logic/graph_coloring/env.py b/jumanji/environments/logic/graph_coloring/env.py
@@ -73,7 +73,7 @@ class GraphColoring(Environment[State]):
     ```python
     from jumanji.environments import GraphColoring
     env = GraphColoring()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/logic/minesweeper/env.py b/jumanji/environments/logic/minesweeper/env.py
@@ -78,7 +78,7 @@ class Minesweeper(Environment[State]):
     ```python
     from jumanji.environments import Minesweeper
     env = Minesweeper()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/logic/rubiks_cube/env.py b/jumanji/environments/logic/rubiks_cube/env.py
@@ -72,7 +72,7 @@ class RubiksCube(Environment[State]):
     ```python
     from jumanji.environments import RubiksCube
     env = RubiksCube()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/logic/sudoku/env.py b/jumanji/environments/logic/sudoku/env.py
@@ -63,7 +63,7 @@ class Sudoku(Environment[State]):
     ```python
     from jumanji.environments import Sudoku
     env = Sudoku()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/packing/bin_pack/env.py b/jumanji/environments/packing/bin_pack/env.py
@@ -103,7 +103,7 @@ class BinPack(Environment[State]):
     ```python
     from jumanji.environments import BinPack
     env = BinPack()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/packing/job_shop/env.py b/jumanji/environments/packing/job_shop/env.py
@@ -80,7 +80,7 @@ class JobShop(Environment[State]):
     ```python
     from jumanji.environments import JobShop
     env = JobShop()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/packing/knapsack/env.py b/jumanji/environments/packing/knapsack/env.py
@@ -73,7 +73,7 @@ class Knapsack(Environment[State]):
     ```python
     from jumanji.environments import Knapsack
     env = Knapsack()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/packing/tetris/env.py b/jumanji/environments/packing/tetris/env.py
@@ -66,7 +66,7 @@ class Tetris(Environment[State]):
     ```python
     from jumanji.environments import Tetris
     env = Tetris()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/cleaner/env.py b/jumanji/environments/routing/cleaner/env.py
@@ -71,7 +71,7 @@ class Cleaner(Environment[State]):
     ```python
     from jumanji.environments import Cleaner
     env = Cleaner()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/connector/env.py b/jumanji/environments/routing/connector/env.py
@@ -85,7 +85,7 @@ class Connector(Environment[State]):
     ```python
     from jumanji.environments import Connector
     env = Connector()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/cvrp/env.py b/jumanji/environments/routing/cvrp/env.py
@@ -86,7 +86,7 @@ class CVRP(Environment[State]):
     ```python
     from jumanji.environments import CVRP
     env = CVRP()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/maze/env.py b/jumanji/environments/routing/maze/env.py
@@ -68,7 +68,7 @@ class Maze(Environment[State]):
     ```python
     from jumanji.environments import Maze
     env = Maze()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/mmst/env.py b/jumanji/environments/routing/mmst/env.py
@@ -117,6 +117,17 @@ class MMST(Environment[State]):
             - INVALID_CHOICE = -1
             - INVALID_TIE_BREAK = -2
             - INVALID_ALREADY_TRAVERSED = -3
+
+    ```python
+    from jumanji.environments import MMST
+    env = MMST()
+    key = jax.random.PRNGKey(0)
+    state, timestep = jax.jit(env.reset)(key)
+    env.render(state)
+    action = env.action_spec().generate_value()
+    state, timestep = jax.jit(env.step)(state, action)
+    env.render(state)
+    ```
     """
 
     def __init__(

diff --git a/jumanji/environments/routing/multi_cvrp/env.py b/jumanji/environments/routing/multi_cvrp/env.py
@@ -64,6 +64,17 @@ class MultiCVRP(Environment[State]):
 
     [1] Zhang et al. (2020). "Multi-Vehicle Routing Problems with Soft Time Windows: A
     Multi-Agent Reinforcement Learning Approach".
+
+    ```python
+    from jumanji.environments import MultiCVRP
+    env = MultiCVRP()
+    key = jax.random.PRNGKey(0)
+    state, timestep = jax.jit(env.reset)(key)
+    env.render(state)
+    action = env.action_spec().generate_value()
+    state, timestep = jax.jit(env.step)(state, action)
+    env.render(state)
+    ```
     """
 
     def __init__(

diff --git a/jumanji/environments/routing/snake/env.py b/jumanji/environments/routing/snake/env.py
@@ -81,7 +81,7 @@ class Snake(Environment[State]):
     ```python
     from jumanji.environments import Snake
     env = Snake()
-    key = jax.random.key(0)
+    key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
     action = env.action_spec().generate_value()

diff --git a/jumanji/environments/routing/sokoban/__init__.py b/jumanji/environments/routing/sokoban/__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.routing.sokoban.env import Sokoban
+from jumanji.environments.routing.sokoban.types import Observation, State
diff --git a/jumanji/environments/routing/sokoban/constants.py b/jumanji/environments/routing/sokoban/constants.py
@@ -0,0 +1,37 @@
+# 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 jax.numpy as jnp
+
+# Translating actions to coordinate changes
+MOVES = jnp.array([[-1, 0], [0, 1], [1, 0], [0, -1]])
+NOOP = -1
+
+# Object encodings
+EMPTY = 0
+WALL = 1
+TARGET = 2
+AGENT = 3
+BOX = 4
+TARGET_AGENT = 5
+TARGET_BOX = 6
+
+# Environment Variables
+N_BOXES = 4
+GRID_SIZE = 10
+
+# Reward Function
+LEVEL_COMPLETE_BONUS = 10
+SINGLE_BOX_BONUS = 1
+STEP_BONUS = -0.1