1.0.0-beta.4: Merge pull request #21 from Pabloo22/development

1.0.0-beta.4: Add `ResourceTaskGraphObservation` wrapper

Author: Pabloo22

Diff for: 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.0b3
+pip install job-shop-lib==1.0.0b4
 ```
 
 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.

Diff for: job_shop_lib/__init__.py

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

Diff for: job_shop_lib/reinforcement_learning/__init__.py

@@ -13,6 +13,8 @@
     RenderConfig
     add_padding
     create_edge_type_dict
+    ResourceTaskGraphObservation
+    ResourceTaskGraphObservationDict
 
 """
 
@@ -31,6 +33,7 @@
 from job_shop_lib.reinforcement_learning._utils import (
     add_padding,
     create_edge_type_dict,
+    map_values,
 )
 
 from job_shop_lib.reinforcement_learning._single_job_shop_graph_env import (
@@ -39,6 +42,9 @@
 from job_shop_lib.reinforcement_learning._multi_job_shop_graph_env import (
     MultiJobShopGraphEnv,
 )
+from ._resource_task_graph_observation import (
+    ResourceTaskGraphObservation, ResourceTaskGraphObservationDict
+)
 
 
 __all__ = [
@@ -52,4 +58,7 @@
     "add_padding",
     "MultiJobShopGraphEnv",
     "create_edge_type_dict",
+    "ResourceTaskGraphObservation",
+    "map_values",
+    "ResourceTaskGraphObservationDict",
 ]

Diff for: job_shop_lib/reinforcement_learning/_multi_job_shop_graph_env.py

@@ -235,13 +235,13 @@ def ready_operations_filter(
     @ready_operations_filter.setter
     def ready_operations_filter(
         self,
-        pruning_function: Callable[
+        ready_operations_filter: Callable[
             [Dispatcher, List[Operation]], List[Operation]
         ],
     ) -> None:
         """Sets the ready operations filter."""
         self.single_job_shop_graph_env.dispatcher.ready_operations_filter = (
-            pruning_function
+            ready_operations_filter
         )
 
     @property

Diff for: job_shop_lib/reinforcement_learning/_resource_task_graph_observation.py

@@ -0,0 +1,258 @@
+"""Contains wrappers for the environments."""
+
+from typing import TypeVar, TypedDict
+from gymnasium import ObservationWrapper
+import numpy as np
+from numpy.typing import NDArray
+
+from job_shop_lib.reinforcement_learning import (
+    ObservationDict,
+    SingleJobShopGraphEnv,
+    MultiJobShopGraphEnv,
+    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.dispatching.feature_observers import FeatureType
+
+T = TypeVar("T", bound=np.number)
+
+
+class ResourceTaskGraphObservationDict(TypedDict):
+    """Represents a dictionary for resource task graph observations."""
+
+    edge_index_dict: dict[str, NDArray[np.int64]]
+    node_features_dict: dict[str, NDArray[np.float32]]
+    original_ids_dict: dict[str, NDArray[np.int32]]
+
+
+# pylint: disable=line-too-long
+class ResourceTaskGraphObservation(ObservationWrapper):
+    """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).
+
+    In particular, the ``edge_index`` is converted into a ``edge_index_dict``
+    with keys ``(node_type_i, "to", node_type_j)``. The ``node_type_i`` and
+    ``node_type_j`` are the node types of the source and target nodes,
+    respectively.
+
+    Attributes:
+        global_to_local_id: A dictionary mapping global node IDs to local node
+            IDs for each node type.
+        type_ranges: A dictionary mapping node type names to (start, end) index
+            ranges.
+
+    Args:
+        env: The environment to wrap.
+    """
+
+    def __init__(self, env: SingleJobShopGraphEnv | MultiJobShopGraphEnv):
+        super().__init__(env)
+        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.
+
+        Args:
+            action:
+                The action to take. The action is a tuple of two integers
+                (job_id, machine_id):
+                the job ID and the machine ID in which to schedule the
+                operation.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - The reward obtained.
+            - Whether the environment is done.
+            - Whether the episode was truncated (always False).
+            - A dictionary with additional information. The dictionary
+              contains the following keys: "feature_names", the names of the
+              features in the observation; and "available_operations_with_ids",
+              a list of available actions in the form of (operation_id,
+              machine_id, job_id).
+        """
+        observation, reward, done, truncated, info = self.env.step(action)
+        return self.observation(observation), reward, done, truncated, info
+
+    def reset(self, *, seed: int | None = None, options: dict | None = None):
+        """Resets the environment.
+
+        Args:
+            seed:
+                Added to match the signature of the parent class. It is not
+                used in this method.
+            options:
+                Additional options to pass to the environment. Not used in
+                this method.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - A dictionary with additional information, keys
+                include: "feature_names", the names of the features in the
+                observation; and "available_operations_with_ids", a list of
+                available a list of available actions in the form of
+                (operation_id, machine_id, job_id).
+        """
+        observation, info = self.env.reset()
+        return self.observation(observation), info
+
+    def _compute_id_mappings(self) -> dict[int, int]:
+        """Computes mappings from global node IDs to type-local IDs.
+
+        Returns:
+            A dictionary mapping global node IDs to local node IDs for each
+            node type.
+        """
+        mappings = {}
+        for node_type in NodeType:
+            type_nodes = self.job_shop_graph.nodes_by_type[node_type]
+            if not type_nodes:
+                continue
+            # Create mapping from global ID to local ID
+            # (0 to len(type_nodes)-1)
+            type_mapping = {
+                node.node_id: local_id
+                for local_id, node in enumerate(type_nodes)
+            }
+            mappings.update(type_mapping)
+
+        return mappings
+
+    def _compute_node_type_ranges(self) -> dict[str, tuple[int, int]]:
+        """Computes index ranges for each node type.
+
+        Returns:
+            Dictionary mapping node type names to (start, end) index ranges
+        """
+        type_ranges = {}
+        for node_type in NodeType:
+            type_nodes = self.job_shop_graph.nodes_by_type[node_type]
+            if not type_nodes:
+                continue
+            start = min(node.node_id for node in type_nodes)
+            end = max(node.node_id for node in type_nodes) + 1
+            type_ranges[node_type.name.lower()] = (start, end)
+
+        return type_ranges
+
+    def observation(self, observation: ObservationDict):
+        edge_index_dict = create_edge_type_dict(
+            observation["edge_index"],
+            type_ranges=self.type_ranges,
+            relationship="to",
+        )
+        # mapping from global node ID to local node ID
+        for key, edge_index in edge_index_dict.items():
+            edge_index_dict[key] = map_values(
+                edge_index, self.global_to_local_id
+            )
+        node_features_dict = self._create_node_features_dict(observation)
+        node_features_dict, original_ids_dict = self._remove_nodes(
+            node_features_dict, observation["removed_nodes"]
+        )
+
+        return {
+            "edge_index_dict": edge_index_dict,
+            "node_features_dict": node_features_dict,
+            "original_ids_dict": original_ids_dict,
+        }
+
+    def _create_node_features_dict(
+        self, observation: ObservationDict
+    ) -> dict[str, NDArray]:
+        """Creates a dictionary of node features for each node type.
+
+        Args:
+            observation: The observation dictionary.
+
+        Returns:
+            Dictionary mapping node type names to node features.
+        """
+        node_type_to_feature_type = {
+            NodeType.OPERATION: FeatureType.OPERATIONS,
+            NodeType.MACHINE: FeatureType.MACHINES,
+            NodeType.JOB: FeatureType.JOBS,
+        }
+        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:
+                node_features_dict[feature_type.value] = observation[
+                    feature_type.value
+                ]
+                continue
+            if feature_type != FeatureType.JOBS:
+                continue
+            assert FeatureType.OPERATIONS.value in observation
+            job_features = observation[
+                feature_type.value  # type: ignore[literal-required]
+            ]
+            job_ids_of_ops = [
+                node.operation.job_id
+                for node in self.job_shop_graph.nodes_by_type[
+                    NodeType.OPERATION
+                ]
+            ]
+            job_features_expanded = job_features[job_ids_of_ops]
+            operation_features = observation[FeatureType.OPERATIONS.value]
+            node_features_dict[FeatureType.OPERATIONS.value] = np.concatenate(
+                (operation_features, job_features_expanded), axis=1
+            )
+        return node_features_dict
+
+    def _remove_nodes(
+        self,
+        node_features_dict: dict[str, NDArray[np.float32]],
+        removed_nodes: NDArray[np.bool_],
+    ) -> tuple[dict[str, NDArray[np.float32]], dict[str, NDArray[np.int32]]]:
+        """Removes nodes from the node features dictionary.
+
+        Args:
+            node_features_dict: The node features dictionary.
+
+        Returns:
+            The node features dictionary with the nodes removed and a
+            dictionary containing the original node ids.
+        """
+        removed_nodes_dict: dict[str, NDArray[np.float32]] = {}
+        original_ids_dict: dict[str, NDArray[np.int32]] = {}
+        feature_type_to_node_type = {
+            FeatureType.OPERATIONS.value: NodeType.OPERATION,
+            FeatureType.MACHINES.value: NodeType.MACHINE,
+            FeatureType.JOBS.value: NodeType.JOB,
+        }
+        for feature_type, features in node_features_dict.items():
+            node_type = feature_type_to_node_type[feature_type].name.lower()
+            if node_type not in self.type_ranges:
+                continue
+            start, end = self.type_ranges[node_type]
+            removed_nodes_of_this_type = removed_nodes[start:end]
+            removed_nodes_dict[node_type] = features[
+                ~removed_nodes_of_this_type
+            ]
+            original_ids_dict[node_type] = np.where(
+                ~removed_nodes_of_this_type
+            )[0]
+
+        return removed_nodes_dict, original_ids_dict

Diff for: job_shop_lib/reinforcement_learning/_single_job_shop_graph_env.py

@@ -243,8 +243,27 @@ def reset(
         *,
         seed: Optional[int] = None,
         options: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[ObservationDict, dict]:
-        """Resets the environment."""
+    ) -> Tuple[ObservationDict, dict[str, Any]]:
+        """Resets the environment.
+
+        Args:
+            seed:
+                Added to match the signature of the parent class. It is not
+                used in this method.
+            options:
+                Additional options to pass to the environment. Not used in
+                this method.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - A dictionary with additional information, keys
+                include: "feature_names", the names of the features in the
+                observation; and "available_operations_with_ids", a list of
+                available a list of available actions in the form of
+                (operation_id, machine_id, job_id).
+        """
         super().reset(seed=seed, options=options)
         self.dispatcher.reset()
         obs = self.get_observation()