Merge branch 'instadeepai:main' into main

MANIFEST.in

@@ -2,6 +2,7 @@ include LICENSE
 include requirements/*
 recursive-include * *.npy
 recursive-include jumanji *.png
+recursive-include jumanji *.jpeg
 
 # remove the test specific files
 recursive-exclude * *_test.py

README.md

@@ -39,7 +39,7 @@
   </div>
   <div class="row" align="center">
     <img src="docs/env_anim/multi_cvrp.gif" alt="MultiCVRP" width="16%">
-    <img src="docs/env_anim/pac_man.gif" alt="PacMan" width="16%">
+    <img src="docs/env_anim/pac_man.gif" alt="PacMan" width="12.9%">
     <img src="docs/env_anim/robot_warehouse.gif" alt="RobotWarehouse" width="16%">
     <img src="docs/env_anim/rubiks_cube.gif" alt="RubiksCube" width="16%">
     <img src="docs/env_anim/sliding_tile_puzzle.gif" alt="SlidingTilePuzzle" width="16%">
@@ -50,6 +50,7 @@
     <img src="docs/env_anim/sudoku.gif" alt="Sudoku" width="16%">
     <img src="docs/env_anim/tetris.gif" alt="Tetris" width="16%">
     <img src="docs/env_anim/tsp.gif" alt="Tetris" width="16%">
+    <img src="docs/env_anim/lbf.gif" alt="Level-Based Foraging" width="16%">
   </div>
 </div>
 
@@ -121,6 +122,7 @@ problems.
 | 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-v1`                                            | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/pac_man/)      | [doc](https://instadeepai.github.io/jumanji/environments/pac_man/)
 | 👾 Sokoban                                                     | Routing  | `Sokoban-v0`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/sokoban/)          | [doc](https://instadeepai.github.io/jumanji/environments/sokoban/)         |
+| 🍎 Level-Based Foraging                                                     | Routing  | `LevelBasedForaging-v0`                                         | [code](https://github.com/instadeepai/jumanji/tree/main/jumanji/environments/routing/lbf/)          | [doc](https://instadeepai.github.io/jumanji/environments/lbf/)         |
 
 <h2 name="install" id="install">Installation 🎬</h2>
 

docs/api/environments/cleaner.md

@@ -1 +1,9 @@
 ::: jumanji.environments.routing.cleaner.env.Cleaner
+    selection:
+      members:
+        - __init__
+        - reset
+        - step
+        - observation_spec
+        - action_spec
+        - render

docs/api/environments/connector.md

@@ -1,7 +1,7 @@
 ::: jumanji.environments.routing.connector.env.Connector
     selection:
       members:
-        - init
+        - __init__
         - observation_spec
         - action_spec
         - reset

docs/api/environments/lbf.md

@@ -0,0 +1,9 @@
+::: jumanji.environments.routing.lbf.env.LevelBasedForaging
+    selection:
+      members:
+        - __init__
+        - reset
+        - step
+        - observation_spec
+        - action_spec
+        - render

docs/environments/lbf.md

@@ -0,0 +1,43 @@
+# # Level-Based Foraging Environment
+
+<p align="center">
+        <img src="../env_anim/lbf.gif" width="600"/>
+</p>
+
+We provide a JAX jit-able implementation of the [Level-Based Foraging](https://github.com/semitable/lb-foraging/tree/master)
+environment.
+
+The Level-Based Foraging (LBF) represents a mixed cooperative-competitive environment that emphasises coordination between agents. As illustrated above, agents are placed within a grid world and assigned different levels.
+
+To collect food, agents must be adjacent to it and the cumulative level of participating agents must meet or exceed the food's designated level. Agents receive points based on the level of the collected food and their own level.
+
+## Observation
+
+The **observation** seen by the agent is a `NamedTuple` containing the following:
+
+- `agents_view`: jax array (int32) of shape `(num_agents, num_obs_features)`, array representing the agent's view of other agents
+    and food.
+
+- `action_mask`: jax array (bool) of shape `(num_agents, 6)`, array specifying, for each agent,
+    which action (noop, up, down, left, right, load) is legal.
+
+- `step_count`: jax array (int32) of shape `()`, number of steps elapsed in the current episode.
+
+## Action
+
+The action space is a `MultiDiscreteArray` containing an integer value in `[0, 1, 2, 3, 4, 5]` for each
+agent. Each agent can take one of five actions: noop (`0`), up (`1`), down (`2`), turn left (`3`), turn right (`4`), or pick up food (`5`).
+
+The episode terminates under the following conditions:
+
+- An invalid action is taken, or
+
+- An agent collides with another agent.
+
+## Reward
+
+ The reward is equal to the sum of the levels of collected food divided by the level of the agents that collected them.
+
+## Registered Versions 📖
+
+- `LevelBasedForaging-v0`, a grid with 2 agents each with a field of view equal to the grid size (full observation case), with 2 food items and forcing the cooperation between agents.

jumanji/__init__.py

@@ -140,3 +140,9 @@
 register(
     id="SlidingTilePuzzle-v0", entry_point="jumanji.environments:SlidingTilePuzzle"
 )
+
+# LevelBasedForaging  with a random generator with 8 grid size,
+# 2 agents and 2 food items and the maximum agent's level is 2.
+register(
+    id="LevelBasedForaging-v0", entry_point="jumanji.environments:LevelBasedForaging"
+)

jumanji/environments/__init__.py

@@ -50,6 +50,7 @@
 from jumanji.environments.routing.cleaner.env import Cleaner
 from jumanji.environments.routing.connector.env import Connector
 from jumanji.environments.routing.cvrp.env import CVRP
+from jumanji.environments.routing.lbf.env import LevelBasedForaging
 from jumanji.environments.routing.maze.env import Maze
 from jumanji.environments.routing.mmst.env import MMST
 from jumanji.environments.routing.multi_cvrp.env import MultiCVRP

jumanji/environments/packing/bin_pack/env.py

@@ -86,20 +86,18 @@ class BinPack(Environment[State, specs.MultiDiscreteArray, Observation]):
             already packed.
 
     - state: `State`
-        - coordinates: jax array (float) of shape (num_nodes + 1, 2)
-            the coordinates of each node and the depot.
-        - demands: jax array (int32) of shape (num_nodes + 1,)
-            the associated cost of each node and the depot (0.0 for the depot).
-        - position: jax array (int32)
-            the index of the last visited node.
-        - capacity: jax array (int32)
-            the current capacity of the vehicle.
-        - visited_mask: jax array (bool) of shape (num_nodes + 1,)
-            binary mask (False/True <--> not visited/visited).
-        - trajectory: jax array (int32) of shape (2 * num_nodes,)
-            identifiers of the nodes that have been visited (set to DEPOT_IDX if not filled yet).
-        - num_visits: int32
-            number of actions that have been taken (i.e., unique visits).
+        - container: space defined by 2 points, i.e. 6 coordinates.
+        - ems: empty maximal spaces (EMSs) in the container, each defined by 2 points
+            (6 coordinates).
+        - ems_mask: array of booleans that indicate the EMSs that are valid.
+        - items: defined by 3 attributes (x, y, z).
+        - items_mask: array of booleans that indicate the items that can be packed.
+        - items_placed: array of booleans that indicate the items that have been placed so far.
+        - items_location: locations of items in the container, defined by 3 coordinates (x, y, x).
+        - action_mask: array of booleans that indicate the valid actions,
+            i.e. EMSs and items that can be chosen.
+        - sorted_ems_indexes: EMS indexes that are sorted by decreasing volume order.
+        - key: random key used for auto-reset.
 
     ```python
     from jumanji.environments import BinPack

jumanji/environments/routing/cleaner/env.py

@@ -91,14 +91,13 @@ def __init__(
         """Instantiates a `Cleaner` environment.
 
         Args:
-            num_agents: number of agents. Defaults to 3.
-            time_limit: max number of steps in an episode. Defaults to `num_rows * num_cols`.
             generator: `Generator` whose `__call__` instantiates an environment instance.
                 Implemented options are [`RandomGenerator`]. Defaults to `RandomGenerator` with
                 `num_rows=10`, `num_cols=10` and `num_agents=3`.
+            time_limit: max number of steps in an episode. Defaults to `num_rows * num_cols`.
+            penalty_per_timestep: the penalty returned at each timestep in the reward.
             viewer: `Viewer` used for rendering. Defaults to `CleanerViewer` with "human" render
                 mode.
-            penalty_per_timestep: the penalty returned at each timestep in the reward.
         """
         self.generator = generator or RandomGenerator(
             num_rows=10, num_cols=10, num_agents=3

jumanji/environments/routing/connector/env.py

@@ -89,7 +89,7 @@ class Connector(Environment[State, specs.MultiDiscreteArray, Observation]):
     key = jax.random.PRNGKey(0)
     state, timestep = jax.jit(env.reset)(key)
     env.render(state)
-    action = env.action_specc.generate_value()
+    action = env.action_spec.generate_value()
     state, timestep = jax.jit(env.step)(state, action)
     env.render(state)
     ```

jumanji/environments/routing/cvrp/viewer.py

@@ -18,6 +18,7 @@
 import matplotlib.animation
 import matplotlib.pyplot as plt
 import numpy as np
+import pkg_resources
 from chex import Array
 from numpy.typing import NDArray
 
@@ -146,7 +147,10 @@ def _prepare_figure(self, ax: plt.Axes) -> None:
         ax.set_ylim(0, 1)
         ax.get_xaxis().set_visible(False)
         ax.get_yaxis().set_visible(False)
-        map_img = plt.imread("docs/img/city_map.jpeg")
+        img_path = pkg_resources.resource_filename(
+            "jumanji", "environments/routing/cvrp/img/city_map.jpeg"
+        )
+        map_img = plt.imread(img_path)
         ax.imshow(map_img, extent=[0, 1, 0, 1])
 
     def _group_tour(self, tour: Array) -> list:

jumanji/environments/routing/lbf/__init__.py

@@ -0,0 +1,17 @@
+# 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.lbf.env import LevelBasedForaging
+from jumanji.environments.routing.lbf.observer import GridObserver, VectorObserver
+from jumanji.environments.routing.lbf.types import Agent, Food, Observation, State