Don't artificially limit the rendering axes for CVRP environment

rl4co/envs/common/base.py

@@ -282,7 +282,9 @@
 
     def render(self, *args, **kwargs):
         """Render the environment"""
-        raise NotImplementedError
+        raise NotImplementedError(
+            f"Render is not implemented for {self.name} environment. Please implement the `render` method in the subclass."
+        )
 
     @staticmethod
     def load_data(fpath, batch_size=[]):

rl4co/envs/eda/dpp/render.py

@@ -1,10 +1,7 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
-from matplotlib import cm, colormaps
-
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
@@ -15,7 +12,6 @@ def render(self, decaps, probe, action_mask, ax=None, legend=True):
     Plot a grid of 1x1 squares representing the environment.
     The keepout regions are the action_mask - decaps - probe
     """
-    import matplotlib.pyplot as plt
 
     settings = {
         0: {"color": "white", "label": "available"},
@@ -58,9 +54,7 @@ def render(self, decaps, probe, action_mask, ax=None, legend=True):
             ax.add_patch(plt.Rectangle((x, y), 1, 1, color=color, linestyle="-"))
 
     # Add grid with 1x1 squares
-    ax.grid(
-        which="major", axis="both", linestyle="-", color="k", linewidth=1, alpha=0.5
-    )
+    ax.grid(which="major", axis="both", linestyle="-", color="k", linewidth=1, alpha=0.5)
     # set 10 ticks
     ax.set_xticks(np.arange(0, xdim, 1))
     ax.set_yticks(np.arange(0, ydim, 1))
@@ -82,3 +76,5 @@ def render(self, decaps, probe, action_mask, ax=None, legend=True):
             loc="upper center",
             bbox_to_anchor=(0.5, 1.1),
         )
+
+    return ax

rl4co/envs/eda/mdpp/render.py

@@ -1,10 +1,7 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
-from matplotlib import cm, colormaps
-
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
@@ -15,8 +12,6 @@ def render(self, td, actions=None, ax=None, legend=True, settings=None):
     The keepout regions are the action_mask - decaps - probe
     """
 
-    import matplotlib.pyplot as plt
-
     from matplotlib.lines import Line2D
     from matplotlib.patches import Annulus, Rectangle, RegularPolygon
 
@@ -87,9 +82,7 @@ def draw_keepout(ax, x, y, color="black"):
         # Backgrund rectangle: same as color but with alpha=0.5
         ax.add_patch(Rectangle((x, y), 1, 1, color=color, alpha=0.5))
         ax.add_patch(
-            RegularPolygon(
-                (x + 0.5, y + 0.5), numVertices=6, radius=0.45, color=color
-            )
+            RegularPolygon((x + 0.5, y + 0.5), numVertices=6, radius=0.45, color=color)
         )
 
     size = self.size
@@ -132,9 +125,7 @@ def draw_keepout(ax, x, y, color="black"):
             elif keepout[i, j] == 1:
                 draw_keepout(ax, x, y, color=settings["keepout"]["color"])
 
-    ax.grid(
-        which="major", axis="both", linestyle="-", color="k", linewidth=1, alpha=0.5
-    )
+    ax.grid(which="major", axis="both", linestyle="-", color="k", linewidth=1, alpha=0.5)
     # set 10 ticks
     ax.set_xticks(np.arange(0, xdim, 1))
     ax.set_yticks(np.arange(0, ydim, 1))
@@ -159,3 +150,5 @@ def draw_keepout(ax, x, y, color="black"):
             loc="upper center",
             bbox_to_anchor=(0.5, 1.1),
         )
+
+    return ax

rl4co/envs/routing/atsp/render.py

@@ -1,6 +1,6 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
 from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
@@ -41,10 +41,6 @@
 
     # Add arrows between visited nodes as a quiver plot
     dx, dy = np.diff(x), np.diff(y)
-    ax.quiver(
-        x[:-1], y[:-1], dx, dy, scale_units="xy", angles="xy", scale=1, color="k"
-    )
+    ax.quiver(x[:-1], y[:-1], dx, dy, scale_units="xy", angles="xy", scale=1, color="k")
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/cvrp/render.py

@@ -135,5 +135,4 @@
             annotation_clip=False,
         )
 
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/cvrptw/env.py

@@ -215,7 +215,7 @@
 
     @staticmethod
     def render(td: TensorDict, actions: torch.Tensor = None, ax=None):
-        render(td, actions, ax)
+        return render(td, actions, ax)
 
     @staticmethod
     def load_data(

rl4co/envs/routing/cvrptw/render.py

@@ -1,10 +1,9 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
 from matplotlib import cm, colormaps
 
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
@@ -129,5 +128,4 @@
             annotation_clip=False,
         )
 
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/mdcpdp/render.py

@@ -1,8 +1,11 @@
+import torch
+
 from tensordict.tensordict import TensorDict
 
 
 def render(td: TensorDict, actions=None, ax=None):
     import matplotlib.pyplot as plt
+
     markersize = 8
 
     td = td.detach().cpu()
@@ -28,10 +31,10 @@
 
     # Plot the actions in order
     last_depot = 0
-    for i in range(len(actions)-1):
-        if actions[i+1] < n_depots:
-            last_depot = actions[i+1]
-        if actions[i] < n_depots and actions[i+1] < n_depots:
+    for i in range(len(actions) - 1):
+        if actions[i + 1] < n_depots:
+            last_depot = actions[i + 1]
+        if actions[i] < n_depots and actions[i + 1] < n_depots:
             continue
         from_node = actions[i]
         to_node = (
@@ -115,6 +118,4 @@
             alpha=0.5,
         )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/mpdp/render.py

@@ -1,10 +1,9 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
 from matplotlib import cm, colormaps
 
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
@@ -13,11 +12,6 @@
 def render(td, actions=None, ax=None):
     # TODO: color switch with new agents; add pickup and delivery nodes as in `PDPEnv.render`
 
-    import matplotlib.pyplot as plt
-    import numpy as np
-
-    from matplotlib import cm, colormaps
-
     num_routine = (actions == 0).sum().item() + 2
     base = colormaps["nipy_spectral"]
     color_list = base(np.linspace(0, 1, num_routine))
@@ -109,6 +103,4 @@ def render(td, actions=None, ax=None):
             annotation_clip=False,
         )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/mtsp/render.py

@@ -1,6 +1,6 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
 from matplotlib import colormaps
 
@@ -93,3 +93,5 @@
     ax.set_title("mTSP")
     ax.set_xlabel("x-coordinate")
     ax.set_ylabel("y-coordinate")
+
+    return ax

rl4co/envs/routing/mtvrp/render.py

@@ -8,7 +8,7 @@
 
 
 def render(
-    td: TensorDict, actions=None, ax=None, scale_xy: bool = True, vehicle_capacity=None
+    td: TensorDict, actions=None, ax=None, scale_xy: bool = False, vehicle_capacity=None
 ):
     import matplotlib.pyplot as plt
     import numpy as np
@@ -139,7 +139,4 @@ def render(
         ax.set_xlim(-0.05, 1.05)
         ax.set_ylim(-0.05, 1.05)
 
-    # Remove the ticks
-    ax.set_xticks([])
-    ax.set_yticks([])
-    plt.show()
+    return ax

rl4co/envs/routing/op/render.py

@@ -1,10 +1,7 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import torch
 
-from matplotlib import cm, colormaps
-
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
@@ -32,8 +29,7 @@
     customers = td["locs"][1:, :]
     prizes = td["prize"][1:]
     normalized_prizes = (
-        200 * (prizes - torch.min(prizes)) / (torch.max(prizes) - torch.min(prizes))
-        + 10
+        200 * (prizes - torch.min(prizes)) / (torch.max(prizes) - torch.min(prizes)) + 10
     )
 
     # Plot depot and customers with prize
@@ -81,6 +77,4 @@
             width=0.0035,
         )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/pctsp/render.py

@@ -29,8 +29,7 @@
     prizes = td["real_prize"][1:]
     penalties = td["penalty"][1:]
     normalized_prizes = (
-        200 * (prizes - torch.min(prizes)) / (torch.max(prizes) - torch.min(prizes))
-        + 10
+        200 * (prizes - torch.min(prizes)) / (torch.max(prizes) - torch.min(prizes)) + 10
     )
     normalized_penalties = (
         3
@@ -88,6 +87,4 @@
             width=0.0035,
         )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/pdp/render.py

@@ -75,9 +75,7 @@
             label="Delivery" if i == 0 else None,
         )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax
 
 
 def render_improvement(td, current_soltuion, best_soltuion):
@@ -140,3 +138,5 @@
                 )
             ax.set_title("Best Solution")
     plt.tight_layout()
+
+    return ax

rl4co/envs/routing/shpp/render.py

@@ -59,6 +59,4 @@
         headwidth=8,
     )
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax

rl4co/envs/routing/svrp/render.py

@@ -101,3 +101,4 @@ def render(td, actions=None, ax=None):
             size=15,
             annotation_clip=False,
         )
+    return ax

rl4co/envs/routing/tsp/render.py

@@ -43,9 +43,7 @@
     dx, dy = np.diff(x), np.diff(y)
     ax.quiver(x[:-1], y[:-1], dx, dy, scale_units="xy", angles="xy", scale=1, color="k")
 
-    # Setup limits and show
-    ax.set_xlim(-0.05, 1.05)
-    ax.set_ylim(-0.05, 1.05)
+    return ax
 
 
 def render_improvement(td, current_soltuion, best_soltuion):
@@ -108,3 +106,5 @@
                 )
             ax.set_title("Best Solution")
     plt.tight_layout()
+
+    return ax

rl4co/envs/scheduling/ffsp/render.py

@@ -1,19 +1,14 @@
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
 
-from matplotlib import cm, colormaps
 from tensordict.tensordict import TensorDict
 
-from rl4co.utils.ops import gather_by_index
 from rl4co.utils.pylogger import get_pylogger
 
 log = get_pylogger(__name__)
 
 
 def render(td: TensorDict, idx: int):
     import matplotlib.patches as patches
-    import matplotlib.pyplot as plt
 
     # TODO: fix this render function parameters
     num_machine_total = td["num_machine_total"][idx].item()
@@ -59,7 +54,8 @@ def render(td: TensorDict, idx: int):
 
     ax.grid()
     ax.set_axisbelow(True)
-    plt.show()
+    return ax
+
 
 def _get_cmap(color_cnt):
     from random import shuffle

rl4co/envs/scheduling/fjsp/render.py

@@ -68,5 +68,4 @@
     )
 
     plt.tight_layout()
-    # Show the Gantt chart
-    plt.show()
+    return ax

rl4co/envs/scheduling/smtwtp/env.py

@@ -195,4 +195,4 @@
 
     @staticmethod
     def render(td, actions=None, ax=None):
-        raise render(td, actions, ax)
+        return render(td, actions, ax)