Matplotlib upgrade

examples/visualize_random_agent.py

@@ -0,0 +1,118 @@
+# 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 argparse
+import os
+
+import jax
+import requests
+from hydra import compose, initialize
+
+from jumanji.training.setup_train import setup_agent, setup_env
+
+envs = [
+    "bin_pack",
+    "cleaner",
+    "connector",
+    "cvrp",
+    "flat_pack",
+    "game_2048",
+    "graph_coloring",
+    "job_shop",
+    "knapsack",
+    "lbf",
+    "maze",
+    "minesweeper",
+    "mmst",
+    "multi_cvrp",
+    "pac_man",
+    "robot_warehouse",
+    "rubiks_cube",
+    "search_and_rescue",
+    "sliding_tile_puzzle",
+    "snake",
+    "sokoban",
+    "sudoku",
+    "tetris",
+    "tsp",
+]
+
+
+def download_file(url: str, file_path: str) -> None:
+    # Send an HTTP GET request to the URL
+    response = requests.get(url)
+    # Check if the request was successful (status code 200)
+    if response.status_code == 200:
+        with open(file_path, "wb") as f:
+            f.write(response.content)
+    else:
+        print("Failed to download the file.")
+
+
+def create_animation(env_name: str, agent: str = "random", num_episodes: int = 2) -> None:
+    print(f"Animating {env_name}")
+
+    os.makedirs("configs", exist_ok=True)
+    config_url = "https://raw.githubusercontent.com/instadeepai/jumanji/main/jumanji/training/configs/config.yaml"
+    download_file(config_url, "configs/config.yaml")
+    env_url = f"https://raw.githubusercontent.com/instadeepai/jumanji/main/jumanji/training/configs/env/{env_name}.yaml"
+    os.makedirs("configs/env", exist_ok=True)
+    download_file(env_url, f"configs/env/{env_name}.yaml")
+    os.makedirs("animations", exist_ok=True)
+
+    with initialize(version_base=None, config_path="configs"):
+        cfg = compose(config_name="config.yaml", overrides=[f"env={env_name}", f"agent={agent}"])
+
+    env = setup_env(cfg).unwrapped
+    agent = setup_agent(cfg, env)
+    policy = jax.jit(agent.make_policy(stochastic=False))
+
+    reset_fn = jax.jit(env.reset)
+    step_fn = jax.jit(env.step)
+    states = []
+    key = jax.random.PRNGKey(cfg.seed)
+
+    for _ in range(num_episodes):
+        key, reset_key = jax.random.split(key)
+        state, timestep = reset_fn(reset_key)
+        states.append(state)
+
+        while not timestep.last():
+            key, action_key = jax.random.split(key)
+            observation = jax.tree_util.tree_map(lambda x: x[None], timestep.observation)
+            action = policy(observation, action_key)
+            state, timestep = step_fn(state, action.squeeze(axis=0))
+            states.append(state)
+
+    env.animate(states, 100, f"animations/{env_name}_animation.gif")
+
+
+if __name__ == "__main__":
+    cli = argparse.ArgumentParser()
+    cli.add_argument(
+        "envs",
+        nargs="*",
+        type=str,
+        default=None,
+    )
+
+    args = cli.parse_args()
+    arg_envs = args.envs
+    env_list = envs if len(arg_envs) == 0 else args.envs
+
+    for env in env_list:
+        try:
+            create_animation(env)
+        except Exception as e:
+            print(f"{env} failed", e)

jumanji/environments/__init__.py

@@ -82,7 +82,7 @@ def is_notebook() -> bool:
         if is_colab():
             backend = "inline"
         elif is_notebook():
-            backend = "notebook"
+            backend = "ipympl"
         else:
             backend = ""
         IPython.get_ipython().run_line_magic("matplotlib", backend)

jumanji/environments/commons/graph_view_utils.py

@@ -0,0 +1,99 @@
+# 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 typing import List, Tuple
+
+import chex
+import numpy as np
+
+
+def _compute_repulsive_forces(
+    repulsive_forces: np.ndarray, pos: np.ndarray, k: float, num_nodes: int
+) -> np.ndarray:
+    for i in range(num_nodes):
+        for j in range(i + 1, num_nodes):
+            delta = pos[i] - pos[j]
+            distance = np.linalg.norm(delta)
+            direction = delta / (distance + 1e-6)
+            force = k * k / (distance + 1e-6)
+            repulsive_forces[i] += direction * force
+            repulsive_forces[j] -= direction * force
+
+    return repulsive_forces
+
+
+def _compute_attractive_forces(
+    graph: chex.Array,
+    attractive_forces: np.ndarray,
+    pos: np.ndarray,
+    k: float,
+    num_nodes: int,
+) -> np.ndarray:
+    for i in range(num_nodes):
+        for j in range(num_nodes):
+            if graph[i, j]:
+                delta = pos[i] - pos[j]
+                distance = np.linalg.norm(delta)
+                direction = delta / (distance + 1e-6)
+                force = distance * distance / k
+                attractive_forces[i] -= direction * force
+                attractive_forces[j] += direction * force
+
+    return attractive_forces
+
+
+def spring_layout(
+    graph: chex.Array, num_nodes: int, seed: int = 42, iterations: int = 100
+) -> List[Tuple[float, float]]:
+    """
+    Compute a 2D spring layout for the given graph using
+    the Fruchterman-Reingold force-directed algorithm.
+
+    The algorithm computes a layout by simulating the graph as a physical system,
+    where nodes are repelling each other and edges are attracting connected nodes.
+    The method minimizes the energy of the system over several iterations.
+
+    Args:
+        graph: A Graph object representing the adjacency matrix of the graph.
+        num_nodes: Number of graph nodes.
+        seed: An integer used to seed the random number generator for reproducibility.
+        iterations: Number of layout refining iterations.
+
+    Returns:
+        A list of tuples representing the 2D positions of nodes in the graph.
+    """
+    rng = np.random.default_rng(seed)
+    pos = rng.random((num_nodes, 2)) * 2 - 1
+
+    k = np.sqrt(5 / num_nodes)
+    temperature = 2.0  # Added a temperature variable
+
+    for _ in range(iterations):
+        repulsive_forces = _compute_repulsive_forces(np.zeros((num_nodes, 2)), pos, k, num_nodes)
+        attractive_forces = _compute_attractive_forces(
+            graph, np.zeros((num_nodes, 2)), pos, k, num_nodes
+        )
+
+        pos += (repulsive_forces + attractive_forces) * temperature
+        # Reduce the temperature (cooling factor) to refine the layout.
+        temperature *= 0.9
+
+        pos = np.clip(pos, -1, 1)  # Keep positions within the [-1, 1] range
+
+    # Scale positions to fill figure
+    pos_max = np.max(pos, axis=0)
+    pos_min = np.min(pos, axis=0)
+    pos = 0.05 + (pos - pos_min) / (1.1 * (pos_max - pos_min))
+
+    return [(float(p[0]), float(p[1])) for p in pos]

jumanji/environments/commons/maze_utils/maze_rendering.py

@@ -12,25 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, ClassVar, Dict, List, Optional, Sequence, Tuple
+from typing import ClassVar, Dict, List, Optional, Sequence, Tuple
 
 import chex
 import matplotlib.animation
 import matplotlib.cm
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib import image
+from matplotlib.artist import Artist
 from matplotlib.axes import Axes
 from numpy.typing import NDArray
 
-import jumanji.environments
 from jumanji.environments.commons.maze_utils.maze_generation import EMPTY, WALL
-from jumanji.viewer import Viewer
+from jumanji.viewer import MatplotlibViewer
 
 
-class MazeViewer(Viewer):
+class MazeViewer(MatplotlibViewer):
     FONT_STYLE = "monospace"
-    FIGURE_SIZE = (10.0, 10.0)
     # EMPTY is white, WALL is black
     COLORS: ClassVar[Dict[int, List[int]]] = {EMPTY: [1, 1, 1], WALL: [0, 0, 0]}
 
@@ -43,18 +42,7 @@ def __init__(self, name: str, render_mode: str = "human") -> None:
                 - "human": render the environment on screen.
                 - "rgb_array": return a numpy array frame representing the environment.
         """
-        self._name = name
-        # The animation must be stored in a variable that lives as long as the
-        # animation should run. Otherwise, the animation will get garbage-collected.
-        self._animation: Optional[matplotlib.animation.Animation] = None
-
-        self._display: Callable[[plt.Figure], Optional[NDArray]]
-        if render_mode == "rgb_array":
-            self._display = self._display_rgb_array
-        elif render_mode == "human":
-            self._display = self._display_human
-        else:
-            raise ValueError(f"Invalid render mode: {render_mode}")
+        super().__init__(name, render_mode)
 
     def render(self, maze: chex.Array) -> Optional[NDArray]:
         """
@@ -89,19 +77,19 @@ def animate(
         Returns:
             Animation that can be saved as a GIF, MP4, or rendered with HTML.
         """
-        fig, ax = plt.subplots(num=f"{self._name}Animation", figsize=self.FIGURE_SIZE)
-        plt.close(fig)
+        fig, ax = self._get_fig_ax(name_suffix="_animation", show=False)
+        plt.close(fig=fig)
 
-        def make_frame(maze_index: int) -> None:
+        def make_frame(maze: chex.Array) -> Tuple[Artist]:
             ax.clear()
-            maze = mazes[maze_index]
             self._add_grid_image(maze, ax)
+            return (ax,)
 
         # Create the animation object.
         self._animation = matplotlib.animation.FuncAnimation(
             fig,
             make_frame,
-            frames=len(mazes),
+            frames=mazes,
             interval=interval,
         )
 
@@ -111,20 +99,6 @@ def make_frame(maze_index: int) -> None:
 
         return self._animation
 
-    def close(self) -> None:
-        plt.close(self._name)
-
-    def _get_fig_ax(self) -> Tuple[plt.Figure, plt.Axes]:
-        recreate = not plt.fignum_exists(self._name)
-        fig = plt.figure(self._name, figsize=self.FIGURE_SIZE)
-        if recreate:
-            if not plt.isinteractive():
-                fig.show()
-            ax = fig.add_subplot()
-        else:
-            ax = fig.get_axes()[0]
-        return fig, ax
-
     def _add_grid_image(self, maze: chex.Array, ax: Axes) -> image.AxesImage:
         img = self._create_grid_image(maze)
         ax.set_axis_off()
@@ -137,24 +111,3 @@ def _create_grid_image(self, maze: chex.Array) -> NDArray:
         # Draw black frame around maze by padding axis 0 and 1
         img = np.pad(img, ((1, 1), (1, 1), (0, 0)))  # type: ignore
         return img
-
-    def _display_human(self, fig: plt.Figure) -> None:
-        if plt.isinteractive():
-            # Required to update render when using Jupyter Notebook.
-            fig.canvas.draw()
-            if jumanji.environments.is_colab():
-                plt.show(self._name)
-        else:
-            # Required to update render when not using Jupyter Notebook.
-            fig.canvas.draw_idle()
-            fig.canvas.flush_events()
-
-    def _display_rgb_array(self, fig: plt.Figure) -> NDArray:
-        fig.canvas.draw()
-        return np.asarray(fig.canvas.buffer_rgba())
-
-    def _clear_display(self) -> None:
-        if jumanji.environments.is_colab():
-            import IPython.display
-
-            IPython.display.clear_output(True)