Merge branch 'main' into main

.github/workflows/tests_linters.yml

@@ -1,34 +1,51 @@
 name: Tests and Linters 🧪
 
-on: [ push, pull_request ]
+on: [ pull_request ]
 
 jobs:
   tests-and-linters:
     name: "Python ${{ matrix.python-version }} on ${{ matrix.os }}"
     runs-on: "${{ matrix.os }}"
+    timeout-minutes: 20
 
     strategy:
       matrix:
-        python-version: ["3.8", "3.9"]
+        python-version: ["3.10", "3.11", "3.12"]
         os: [ubuntu-latest]
 
     steps:
       - name: Install dependencies for viewer test
         run: sudo apt-get update && sudo apt-get install -y xvfb
+
       - name: Checkout jumanji 🐍
-        uses: actions/checkout@v3
-      - uses: actions/setup-python@v4
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v3
+        with:
+          version: "0.4.26"
+          enable-cache: true
+          cache-dependency-glob: "requirements/requirements**.txt"  # invalidate cache when requirements file changes
+
+      - uses: actions/setup-python@v5
         with:
             python-version: "${{ matrix.python-version }}"
+
       - name: Install python dependencies 🔧
-        run: pip install .[dev,train]
+        run: uv pip install .[dev,train]
+        env:
+          UV_SYSTEM_PYTHON: 1
+
       - name: Run linters 🖌️
         run: pre-commit run --all-files --verbose
+
       - name: Run tests 🧪
         run: pytest -n 2 --cov=jumanji --cov-report=term-missing --junit-xml=test-results.xml -vv jumanji
+
       - name: Run coverage
         run: |
           coverage html --directory=coverage_html_report
           coverage report --fail-under=0.97
+
       - name: Test build docs 📖
         run: mkdocs build --verbose --site-dir docs_public

.gitignore

@@ -150,7 +150,7 @@ cython_debug/
 # MacBook Finder
 .DS_Store
 
-3.8/
+3.10/
 jumanji_env/
 **/outputs/
 *.xml

.pre-commit-config.yaml

@@ -33,7 +33,7 @@ repos:
         name: "Trailing whitespace fixer"
 
   - repo: https://github.com/PyCQA/flake8
-    rev: 3.9.2
+    rev: 7.1.1
     hooks:
       - id: flake8
         name: "Linter"

README.md

@@ -138,7 +138,7 @@ Alternatively, you can install the latest development version directly from GitH
 pip install git+https://github.com/instadeepai/jumanji.git
 ```
 
-Jumanji has been tested on Python 3.8 and 3.9.
+Jumanji has been tested on Python 3.10, 3.11 and 3.12.
 Note that because the installation of JAX differs depending on your hardware accelerator,
 we advise users to explicitly install the correct JAX version (see the
 [official installation guide](https://github.com/google/jax#installation)).

jumanji/env.py

@@ -110,7 +110,7 @@ def reward_spec(self) -> specs.Array:
 
     @cached_property
     def discount_spec(self) -> specs.BoundedArray:
-        """Returns the discount spec. By default, this is assumed to be a single float between 0 and 1.
+        """Returns the discount spec. By default, this is assumed to be a float between 0 and 1.
 
         Returns:
             discount_spec: a `specs.BoundedArray` spec.

jumanji/environments/logic/game_2048/viewer.py

@@ -123,7 +123,7 @@ def make_frame(state_index: int) -> None:
         return self._animation
 
     def get_fig_ax(self) -> Tuple[plt.Figure, plt.Axes]:
-        """This function returns a `Matplotlib` figure and axes object for displaying the 2048 game board.
+        """This function returns a `Matplotlib` figure and axes for displaying the 2048 game board.
 
         Returns:
             A tuple containing the figure and axes objects.

jumanji/environments/logic/graph_coloring/env.py

@@ -296,7 +296,7 @@ def animate(
         interval: int = 200,
         save_path: Optional[str] = None,
     ) -> animation.FuncAnimation:
-        """Creates an animated gif of the `GraphColoring` environment based on the sequence of game states.
+        """Creates an animated gif of the `GraphColoring` environment based on a sequence of states.
 
         Args:
             states: is a list of `State` objects representing the sequence of game states.

jumanji/environments/logic/sliding_tile_puzzle/viewer.py

@@ -71,7 +71,7 @@ def animate(
         interval: int = 200,
         save_path: Optional[str] = None,
     ) -> matplotlib.animation.FuncAnimation:
-        """Creates an animated gif of the sliding tiles puzzle game based on the sequence of game states.
+        """Creates an animated gif of the sliding tiles puzzle game based on a sequence of states.
 
         Args:
             states: is a list of `State` objects representing the sequence of game states.
@@ -101,7 +101,7 @@ def make_frame(state_index: int) -> None:
         return self._animation
 
     def get_fig_ax(self) -> Tuple[plt.Figure, plt.Axes]:
-        """This function returns a `Matplotlib` figure and axes object for displaying the game puzzle.
+        """This function returns a `Matplotlib` figure and axes for displaying the puzzle.
 
         Returns:
             A tuple containing the figure and axes objects.

jumanji/environments/packing/bin_pack/env.py

@@ -397,9 +397,9 @@ def close(self) -> None:
     def _make_observation_and_extras(
         self, state: State
     ) -> Tuple[State, Observation, Dict]:
-        """Computes the observation and the environment metrics to include in `timestep.extras`. Also
-        updates the `action_mask` and `sorted_ems_indexes` in the state. The observation is obtained
-        by selecting a subset of all EMSs, namely the `obs_num_ems` largest ones.
+        """Computes the observation and the environment metrics to include in `timestep.extras`.
+        Also updates the `action_mask` and `sorted_ems_indexes` in the state. The observation is
+        obtained by selecting a subset of all EMSs, namely the `obs_num_ems` largest ones.
         """
         obs_ems, obs_ems_mask, sorted_ems_indexes = self._get_set_of_largest_ems(
             state.ems, state.ems_mask

jumanji/environments/packing/flat_pack/generator.py

@@ -28,8 +28,8 @@
 
 
 class InstanceGenerator(abc.ABC):
-    """Base class for generators for the flat_pack environment. An `InstanceGenerator` is responsible
-    for generating a problem instance when the environment is reset.
+    """Base class for generators for the flat_pack environment. An `InstanceGenerator` is
+    responsible for generating a problem instance when the environment is reset.
     """
 
     def __init__(

jumanji/environments/routing/mmst/generator.py

@@ -88,7 +88,7 @@ def __call__(self, key: chex.PRNGKey) -> State:
 
 
 class SplitRandomGenerator(Generator):
-    """Generates a random environments that is solvable by spliting the graph into multiple sub graphs.
+    """Generates a random environments that is solvable by spliting the graph into sub graphs.
 
     Returns a graph and with a desired number of edges and nodes to connect per agent.
     """

jumanji/environments/routing/multi_cvrp/viewer.py

@@ -210,10 +210,10 @@ def _draw_route(self, ax: plt.Axes, coords: chex.Array, col_id: int) -> None:
         ax.scatter(x, y, s=self.NODE_SIZE, color=self._cmap(col_id))
 
     def _add_tour(self, ax: plt.Axes, state: State) -> None:
-        """Add the customers and the depot to the plot, and draw each route in the tour in a different
-        colour. The tour is the entire trajectory between the visited customers and a route is a
-        trajectory either starting and ending at the depot or starting at the depot and ending at
-        the current city."""
+        """Add the customers and the depot to the plot, and draw each route in the tour in a
+        different colour. The tour is the entire trajectory between the visited customers and a
+        route is a trajectory either starting and ending at the depot or starting at the depot
+        and ending at the current city."""
         x_coords, y_coords = (
             state.nodes.coordinates[:, 0] / self._map_max,
             state.nodes.coordinates[:, 1] / self._map_max,

jumanji/environments/routing/robot_warehouse/conftest.py

@@ -31,8 +31,8 @@
 
 @pytest.fixture(scope="module")
 def robot_warehouse_env() -> RobotWarehouse:
-    """Instantiates a default RobotWarehouse environment with 2 agents, 1 shelf row, 3 shelf columns,
-    a column height of 2, sensor range of 1 and a request queue size of 4."""
+    """Instantiates a default RobotWarehouse environment with 2 agents, 1 shelf row, 3 shelf
+    columns, a column height of 2, sensor range of 1 and a request queue size of 4."""
     generator = RandomGenerator(
         shelf_rows=1,
         shelf_columns=3,

jumanji/environments/routing/robot_warehouse/env.py

@@ -362,7 +362,9 @@ def observation_spec(self) -> specs.Spec[Observation]:
 
     @cached_property
     def action_spec(self) -> specs.MultiDiscreteArray:
-        """Returns the action spec. 5 actions: [0,1,2,3,4] -> [No Op, Forward, Left, Right, Toggle_load].
+        """Returns the action spec. 5 actions:
+        [0,1,2,3,4] -> [No Op, Forward, Left, Right, Toggle_load].
+
         Since this is a multi-agent environment, the environment expects an array of actions.
         This array is of shape (num_agents,).
         """

jumanji/specs.py

@@ -44,9 +44,9 @@
 
 
 class Spec(abc.ABC, Generic[T]):
-    """Adapted from `dm_env.spec.Array`. This is an augmentation of the `Array` spec to allow for nested
-    specs. `self.name`, `self.generate_value` and `self.validate` methods are adapted from the
-    `dm_env` object."""
+    """Adapted from `dm_env.spec.Array`. This is an augmentation of the `Array` spec to allow for
+    nested specs. `self.name`, `self.generate_value` and `self.validate` methods are adapted from
+    the `dm_env` object."""
 
     def __init__(
         self,
@@ -139,7 +139,7 @@ def __getitem__(self, item: str) -> "Spec":
 
 
 class Array(Spec[chex.Array]):
-    """Describes a jax array spec. This is adapted from `dm_env.specs.Array` to suit Jax environments.
+    """Describes a jax array spec. This is adapted from `dm_env.specs.Array` for Jax environments.
 
     An `Array` spec allows an API to describe the arrays that it accepts or returns, before that
     array exists.

jumanji/specs_test.py

@@ -589,7 +589,7 @@ def test_array(self) -> None:
         converted_spec: dm_env.specs.Array = specs.jumanji_specs_to_dm_env_specs(
             jumanji_spec
         )
-        assert type(converted_spec) == type(dm_env_spec)
+        assert type(converted_spec) is type(dm_env_spec)
         assert converted_spec.shape == dm_env_spec.shape
         assert converted_spec.dtype == dm_env_spec.dtype
         assert converted_spec.name == dm_env_spec.name
@@ -602,7 +602,7 @@ def test_bounded_array(self) -> None:
         converted_spec: dm_env.specs.BoundedArray = specs.jumanji_specs_to_dm_env_specs(
             jumanji_spec
         )
-        assert type(converted_spec) == type(dm_env_spec)
+        assert type(converted_spec) is type(dm_env_spec)
         assert converted_spec.shape == dm_env_spec.shape
         assert converted_spec.dtype == dm_env_spec.dtype
         assert converted_spec.name == dm_env_spec.name
@@ -615,7 +615,7 @@ def test_discrete_array(self) -> None:
         converted_spec: dm_env.specs.DiscreteArray = (
             specs.jumanji_specs_to_dm_env_specs(jumanji_spec)
         )
-        assert type(converted_spec) == type(dm_env_spec)
+        assert type(converted_spec) is type(dm_env_spec)
         assert converted_spec.shape == dm_env_spec.shape
         assert converted_spec.dtype == dm_env_spec.dtype
         assert converted_spec.name == dm_env_spec.name
@@ -675,7 +675,7 @@ def test_array(self) -> None:
         jumanji_spec = specs.Array((1, 2), jnp.int32)
         gym_space = gym.spaces.Box(-np.inf, np.inf, (1, 2), jnp.int32)
         converted_spec = specs.jumanji_specs_to_gym_spaces(jumanji_spec)
-        assert type(converted_spec) == type(gym_space)
+        assert type(converted_spec) is type(gym_space)
         assert_trees_all_equal(converted_spec.low, gym_space.low)
         assert_trees_all_equal(converted_spec.high, gym_space.high)
         assert converted_spec.shape == gym_space.shape
@@ -687,7 +687,7 @@ def test_bounded_array(self) -> None:
         )
         gym_space = gym.spaces.Box(low=0.0, high=1.0, shape=(1, 2), dtype=jnp.float32)
         converted_spec = specs.jumanji_specs_to_gym_spaces(jumanji_spec)
-        assert type(converted_spec) == type(gym_space)
+        assert type(converted_spec) is type(gym_space)
         assert converted_spec.shape == gym_space.shape
         assert converted_spec.dtype == gym_space.dtype
         assert_trees_all_equal(converted_spec.low, gym_space.low)
@@ -697,7 +697,7 @@ def test_discrete_array(self) -> None:
         jumanji_spec = specs.DiscreteArray(num_values=5, dtype=jnp.int32)
         gym_space = gym.spaces.Discrete(n=5)
         converted_spec = specs.jumanji_specs_to_gym_spaces(jumanji_spec)
-        assert type(converted_spec) == type(gym_space)
+        assert type(converted_spec) is type(gym_space)
         assert converted_spec.shape == gym_space.shape
         assert converted_spec.dtype == gym_space.dtype
         assert converted_spec.n == gym_space.n
@@ -708,7 +708,7 @@ def test_multi_discrete_array(self) -> None:
         )
         gym_space = gym.spaces.MultiDiscrete(nvec=[5, 6])
         converted_spec = specs.jumanji_specs_to_gym_spaces(jumanji_spec)
-        assert type(converted_spec) == type(gym_space)
+        assert type(converted_spec) is type(gym_space)
         assert converted_spec.shape == gym_space.shape
         assert converted_spec.dtype == gym_space.dtype
         assert jnp.array_equal(converted_spec.nvec, gym_space.nvec)

jumanji/training/networks/graph_coloring/actor_critic.py

@@ -219,7 +219,7 @@ def __call__(self, observation: Observation) -> chex.Array:
             mlp_units=self.transformer_mlp_units,
             w_init_scale=2 / self.num_transformer_layers,
             model_size=self.model_size,
-            name=f"cross_attention_color_node_block_{block_id+1}",
+            name=f"cross_attention_color_node_block_{block_id + 1}",
         )(color_embeddings, current_node_embeddings, current_node_embeddings)
 
         return new_embedding

jumanji/training/networks/mmst/actor_critic.py

@@ -207,7 +207,7 @@ def __call__(self, observation: Observation) -> chex.Array:
             mlp_units=self.transformer_mlp_units,
             w_init_scale=2 / self.num_transformer_layers,
             model_size=self.model_size,
-            name=f"cross_attention_agent_node_block_{block_id+1}",
+            name=f"cross_attention_agent_node_block_{block_id + 1}",
         )(agents_embeddings, current_node_embeddings, current_node_embeddings)
 
         return new_embedding

jumanji/training/networks/tsp/actor_critic.py

@@ -72,8 +72,8 @@ def __init__(
         transformer_mlp_units: Sequence[int],
         name: Optional[str] = None,
     ):
-        """Linear embedding of all cities' coordinates followed by `transformer_num_blocks` blocks of self
-        attention.
+        """Linear embedding of all cities' coordinates followed by `transformer_num_blocks` blocks
+        of self attention.
         """
         super().__init__(name=name)
         self.transformer_num_blocks = transformer_num_blocks

pyproject.toml

@@ -1,8 +1,51 @@
-[tool.isort]
-profile = "black"
+[build-system]
+requires=["setuptools>=62.6"]
+build-backend="setuptools.build_meta"
+
+[project]
+name="jumanji"
+authors=[{name="InstaDeep Ltd", email="[email protected]"}]
+dynamic=["version", "dependencies", "optional-dependencies"]
+license={file="LICENSE"}
+description="A diverse suite of scalable reinforcement learning environments in JAX"
+readme ="README.md"
+requires-python=">=3.10"
+keywords=["reinforcement-learning", "python", "jax"]
+classifiers=[
+    "Development Status :: 5 - Production/Stable",
+    "Environment :: Console",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "License :: OSI Approved :: Apache Software License",
+]
+
+[tool.setuptools.packages.find]
+include=["jumanji*"]
+
+[tool.setuptools.package-data]
+"jumanji" = ["py.typed"]
+
+[tool.setuptools.dynamic]
+version={attr="jumanji.version.__version__"}
+dependencies={file="requirements/requirements.txt"}
+optional-dependencies.dev={file=["requirements/requirements-dev.txt"]}
+optional-dependencies.train={file=["requirements/requirements-train.txt"]}
+
+
+[project.urls]
+"Homepage"="https://github.com/instadeep/jumanji"
+"Bug Tracker"="https://github.com/instadeep/jumanji/issues"
+"Documentation"="https://instadeepai.github.io/jumanji"
 
 [tool.mypy]
-python_version = 3.8
+python_version = "3.10"
 namespace_packages = true
 incremental = false
 cache_dir = ""
@@ -47,3 +90,6 @@ module = [
     "PIL.*",
 ]
 ignore_missing_imports = true
+
+[tool.isort]
+profile = "black"

requirements/requirements-dev.txt

@@ -1,7 +1,6 @@
 black==22.3.0
 coverage
-flake8==3.9.2
-importlib-metadata<5.0
+flake8
 isort==5.11.5
 livereload
 mkdocs==1.2.3
@@ -22,9 +21,5 @@ pytest-cov
 pytest-mock
 pytest-parallel
 pytest-xdist
-pytype
 scipy>=1.7.3
 testfixtures
-types-Pillow
-types-requests<1.27
-types-setuptools