1.0.0-beta.5: Merge pull request #22 from Pabloo22/development

1.0.0-beta.5: Enhance `ResourceTaskGraphObservation` with generic environment type

Author: Pabloo22

README.md

@@ -36,7 +36,7 @@ See [this](https://colab.research.google.com/drive/1XV_Rvq1F2ns6DFG8uNj66q_rcoww
 Version 1.0.0 is currently in beta stage and can be installed with:
 
 ```bash
-pip install job-shop-lib==1.0.0b4
+pip install job-shop-lib==1.0.0b5
 ```
 
 Although this version is not stable and may contain breaking changes in subsequent releases, it is recommended to install it to access the new reinforcement learning environments and familiarize yourself with new changes (see the [latest pull requests](https://github.com/Pabloo22/job_shop_lib/pulls?q=is%3Apr+is%3Aclosed)). There is a [documentation page](https://job-shop-lib.readthedocs.io/en/latest/) for versions 1.0.0a3 and onward.

job_shop_lib/__init__.py

@@ -19,7 +19,7 @@
 from job_shop_lib._base_solver import BaseSolver, Solver
 
 
-__version__ = "1.0.0-b.4"
+__version__ = "1.0.0-b.5"
 
 __all__ = [
     "Operation",

job_shop_lib/reinforcement_learning/_resource_task_graph_observation.py

@@ -1,6 +1,6 @@
 """Contains wrappers for the environments."""
 
-from typing import TypeVar, TypedDict
+from typing import TypeVar, TypedDict, Generic
 from gymnasium import ObservationWrapper
 import numpy as np
 from numpy.typing import NDArray
@@ -12,11 +12,13 @@
     create_edge_type_dict,
     map_values,
 )
-from job_shop_lib.graphs import NodeType, JobShopGraph
-from job_shop_lib.exceptions import ValidationError
+from job_shop_lib.graphs import NodeType
 from job_shop_lib.dispatching.feature_observers import FeatureType
 
 T = TypeVar("T", bound=np.number)
+EnvType = TypeVar(  # pylint: disable=invalid-name
+    "EnvType", bound=SingleJobShopGraphEnv | MultiJobShopGraphEnv
+)
 
 
 class ResourceTaskGraphObservationDict(TypedDict):
@@ -28,7 +30,7 @@ class ResourceTaskGraphObservationDict(TypedDict):
 
 
 # pylint: disable=line-too-long
-class ResourceTaskGraphObservation(ObservationWrapper):
+class ResourceTaskGraphObservation(ObservationWrapper, Generic[EnvType]):
     """Observation wrapper that converts an observation following the
     :class:`ObservationDict` format to a format suitable to PyG's
     [`HeteroData`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.data.HeteroData.html).
@@ -48,26 +50,12 @@ class ResourceTaskGraphObservation(ObservationWrapper):
         env: The environment to wrap.
     """
 
-    def __init__(self, env: SingleJobShopGraphEnv | MultiJobShopGraphEnv):
+    def __init__(self, env: EnvType):
         super().__init__(env)
+        self.env = env  # Unnecessary, but makes mypy happy
         self.global_to_local_id = self._compute_id_mappings()
         self.type_ranges = self._compute_node_type_ranges()
 
-    @property
-    def job_shop_graph(self) -> JobShopGraph:
-        """Returns the job shop graph from the environment.
-
-        Raises:
-            ValidationError: If the environment is not an instance of
-                ``SingleJobShopGraphEnv`` or ``MultiJobShopGraphEnv``.
-        """
-        if isinstance(self.env, (SingleJobShopGraphEnv, MultiJobShopGraphEnv)):
-            return self.env.job_shop_graph
-        raise ValidationError(
-            "The environment must be an instance of "
-            "SingleJobShopGraphEnv or MultiJobShopGraphEnv"
-        )
-
     def step(self, action: tuple[int, int]):
         """Takes a step in the environment.
 
@@ -127,7 +115,7 @@ def _compute_id_mappings(self) -> dict[int, int]:
         """
         mappings = {}
         for node_type in NodeType:
-            type_nodes = self.job_shop_graph.nodes_by_type[node_type]
+            type_nodes = self.unwrapped.job_shop_graph.nodes_by_type[node_type]
             if not type_nodes:
                 continue
             # Create mapping from global ID to local ID
@@ -148,7 +136,7 @@ def _compute_node_type_ranges(self) -> dict[str, tuple[int, int]]:
         """
         type_ranges = {}
         for node_type in NodeType:
-            type_nodes = self.job_shop_graph.nodes_by_type[node_type]
+            type_nodes = self.unwrapped.job_shop_graph.nodes_by_type[node_type]
             if not type_nodes:
                 continue
             start = min(node.node_id for node in type_nodes)
@@ -197,7 +185,7 @@ def _create_node_features_dict(
         }
         node_features_dict = {}
         for node_type, feature_type in node_type_to_feature_type.items():
-            if node_type in self.job_shop_graph.nodes_by_type:
+            if node_type in self.unwrapped.job_shop_graph.nodes_by_type:
                 node_features_dict[feature_type.value] = observation[
                     feature_type.value
                 ]
@@ -210,7 +198,7 @@ def _create_node_features_dict(
             ]
             job_ids_of_ops = [
                 node.operation.job_id
-                for node in self.job_shop_graph.nodes_by_type[
+                for node in self.unwrapped.job_shop_graph.nodes_by_type[
                     NodeType.OPERATION
                 ]
             ]
@@ -256,3 +244,8 @@ def _remove_nodes(
             )[0]
 
         return removed_nodes_dict, original_ids_dict
+
+    @property
+    def unwrapped(self) -> EnvType:
+        """Returns the unwrapped environment."""
+        return self.env  # type: ignore[return-value]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "job-shop-lib"
-version = "1.0.0-beta.4"
+version = "1.0.0-beta.5"
 description = "An easy-to-use and modular Python library for the Job Shop Scheduling Problem (JSSP)"
 authors = ["Pabloo22 <[email protected]>"]
 license = "MIT"

tests/reinforcement_learning/test_resource_task_graph_observation.py

@@ -14,13 +14,13 @@ def test_edge_index_dict(
         single_env_ft06_resource_task_graph_with_all_features
     )
     obs, info = env.reset()
-    max_index = env.job_shop_graph.instance.num_operations
+    max_index = env.unwrapped.job_shop_graph.instance.num_operations
     edge_index_dict = obs["edge_index_dict"]
     _check_that_edge_index_has_been_reindexed(edge_index_dict, max_index)
 
     done = False
     _, machine_id, job_id = info["available_operations_with_ids"][0]
-    removed_nodes = env.job_shop_graph.removed_nodes
+    removed_nodes = env.unwrapped.job_shop_graph.removed_nodes
     _check_count_of_unique_ids(edge_index_dict, removed_nodes)
     while not done:
         obs, _, done, _, info = env.step((job_id, machine_id))
@@ -43,7 +43,7 @@ def test_node_features_dict(
     obs, info = env.reset()
     done = False
     _, machine_id, job_id = info["available_operations_with_ids"][0]
-    removed_nodes = env.job_shop_graph.removed_nodes
+    removed_nodes = env.unwrapped.job_shop_graph.removed_nodes
     _check_number_of_nodes(obs["node_features_dict"], removed_nodes)
     while not done:
         obs, _, done, _, info = env.step((job_id, machine_id))
@@ -64,7 +64,7 @@ def test_original_ids_dict(
     obs, info = env.reset()
     done = False
     _, machine_id, job_id = info["available_operations_with_ids"][0]
-    removed_nodes = env.job_shop_graph.removed_nodes
+    removed_nodes = env.unwrapped.job_shop_graph.removed_nodes
     _check_original_ids_dict(obs["original_ids_dict"], removed_nodes)
     while not done:
         obs, _, done, _, info = env.step((job_id, machine_id))