feat: full lbf implementation

jumanji/__init__.py

@@ -128,7 +128,8 @@
 # TSP with 20 randomly generated cities and a dense reward function.
 register(id="TSP-v1", entry_point="jumanji.environments:TSP")
 
-# LevelBasedForaging on 10x10 grid with 3 agents and 3 foods
+# LevelBasedForaging  with a random generator with 8 grid size,
+# 2 agents and 2 foods and maximum agent's level is 2.
 register(
     id="LevelBasedForaging-v0", entry_point="jumanji.environments:LevelBasedForaging"
 )

jumanji/environments/routing/lbf/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 
 from jumanji.environments.routing.lbf.env import LevelBasedForaging
-from jumanji.environments.routing.lbf.types import Observation, State
+from jumanji.environments.routing.lbf.observer import GridObserver, VectorObserver
+from jumanji.environments.routing.lbf.types import Agent, Food, Observation, State

jumanji/environments/routing/lbf/constants.py

@@ -13,22 +13,20 @@
 # limitations under the License.
 import jax.numpy as jnp
 
-# Observations
-EMPTY = 0
-FOOD = 1
-AGENT = 2
-
 # Actions
 NOOP = 0
 UP = 1
-RIGHT = 2
-DOWN = 3
-LEFT = 4
+DOWN = 2
+LEFT = 3
+RIGHT = 4
 LOAD = 5
 
-# No-op, Up, Right, Down, Left, Load
-MOVES = jnp.array([[0, 0], [-1, 0], [0, 1], [1, 0], [0, -1], [0, 0]])
+# NOOP, UP, DOWN, LEFT, RIGHT, LOAD
+MOVES = jnp.array([[0, 0], [-1, 0], [1, 0], [0, -1], [0, 1], [0, 0]])
+
+# viewer constants
+_FIGURE_SIZE = (5, 5)
 
-# Rendering
-VIEWER_WIDTH = 1000
-VIEWER_HEIGHT = 1000
+# Define some colors for visualization.
+_GRID_COLOR = (0, 0, 0)  # black
+_LINE_COLOR = (1, 1, 1)  # white