job-shop-lib 0.5.1__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

job_shop_lib/reinforcement_learning/_single_job_shop_graph_env.py

@@ -0,0 +1,337 @@
+"""Home of the `SingleJobShopGraphEnv` class."""
+
+from copy import deepcopy
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import matplotlib.pyplot as plt
+import gymnasium as gym
+import numpy as np
+from numpy.typing import NDArray
+
+from job_shop_lib import JobShopInstance, Operation
+from job_shop_lib.graphs import JobShopGraph
+from job_shop_lib.graphs.graph_updaters import (
+    GraphUpdater,
+    ResidualGraphUpdater,
+)
+from job_shop_lib.dispatching import (
+    Dispatcher,
+    filter_dominated_operations,
+    DispatcherObserverConfig,
+)
+from job_shop_lib.dispatching.feature_observers import (
+    FeatureObserverConfig,
+    CompositeFeatureObserver,
+)
+from job_shop_lib.visualization import GanttChartCreator
+from job_shop_lib.reinforcement_learning import (
+    RewardObserver,
+    MakespanReward,
+    add_padding,
+    RenderConfig,
+    ObservationSpaceKey,
+    ObservationDict,
+)
+
+
+class SingleJobShopGraphEnv(gym.Env):
+    """A Gymnasium environment for solving a specific instance of the Job Shop
+    Scheduling Problem represented as a graph.
+
+    This environment manages the scheduling process for a single Job Shop
+    instance, using a graph representation and various observers to track the
+    state and compute rewards.
+
+    Observation Space:
+        A dictionary with the following keys:
+        - "removed_nodes": Binary vector indicating removed graph nodes.
+        - "edge_list": Matrix of graph edges in COO format.
+        - Feature matrices: Keys corresponding to the composite observer
+          features (e.g., "operations", "jobs", "machines").
+
+    Action Space:
+        MultiDiscrete space representing (job_id, machine_id) pairs.
+
+    Render Modes:
+        - "human": Displays the current Gantt chart.
+        - "save_video": Saves a video of the complete Gantt chart.
+        - "save_gif": Saves a GIF of the complete Gantt chart.
+
+    Attributes:
+        dispatcher:
+            Manages the scheduling process. See
+            :class:`~job_shop_lib.dispatching.Dispatcher`.
+        composite_observer:
+            A :class:`~job_shop_lib.dispatching.feature_observers.
+            CompositeFeatureObserver` which aggregates features from multiple
+            observers.
+        graph_updater:
+            Updates the graph representation after each action. See
+            :class:`~job_shop_lib.graphs.GraphUpdater`.
+        reward_function:
+            Computes rewards for actions taken. See
+            :class:`~job_shop_lib.reinforcement_learning.RewardObserver`.
+        action_space:
+            Defines the action space. The action is a tuple of two integers
+            (job_id, machine_id). The machine_id can be -1 if the selected
+            operation can only be scheduled in one machine.
+        observation_space:
+            Defines the observation space. The observation is a dictionary
+            with the following keys:
+            - "removed_nodes": Binary vector indicating removed graph nodes.
+            - "edge_list": Matrix of graph edges in COO format.
+            - Feature matrices: Keys corresponding to the composite observer
+                features (e.g., "operations", "jobs", "machines").
+        render_mode:
+            The mode for rendering the environment ("human", "save_video",
+            "save_gif").
+        gantt_chart_creator:
+            Creates Gantt chart visualizations. See
+            :class:`~job_shop_lib.visualization.GanttChartCreator`.
+        use_padding:
+            Whether to use padding in observations. Padding maintains the
+    """
+
+    metadata = {"render_modes": ["human", "save_video", "save_gif"]}
+
+    # I think the class is easier to use this way. We could initiliaze the
+    # class from Dispatcher or an already initialized RewardFunction. However,
+    # it would be impossible to add good default values.
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        job_shop_graph: JobShopGraph,
+        feature_observer_configs: Sequence[FeatureObserverConfig],
+        reward_function_config: DispatcherObserverConfig[
+            type[RewardObserver]
+        ] = DispatcherObserverConfig(class_type=MakespanReward),
+        graph_updater_config: DispatcherObserverConfig[
+            type[GraphUpdater]
+        ] = DispatcherObserverConfig(class_type=ResidualGraphUpdater),
+        ready_operations_filter: (
+            Callable[[Dispatcher, list[Operation]], list[Operation]] | None
+        ) = filter_dominated_operations,
+        render_mode: str | None = None,
+        render_config: RenderConfig | None = None,
+        use_padding: bool = True,
+    ) -> None:
+        """Initializes the SingleJobShopGraphEnv environment.
+
+        Args:
+            job_shop_graph:
+                The JobShopGraph instance representing the job shop problem.
+            feature_observer_configs:
+                A list of FeatureObserverConfig instances for the feature
+                observers.
+            reward_function_config:
+                The configuration for the reward function.
+            graph_updater_config:
+                The configuration for the graph updater.
+            ready_operations_filter:
+                The function to use for pruning dominated operations.
+            render_mode:
+                The mode for rendering the environment ("human", "save_video",
+                "save_gif").
+            render_config:
+                Configuration for rendering (e.g., paths for saving videos
+                or GIFs).
+            use_padding:
+                Whether to use padding for the edge index.
+        """
+        super().__init__()
+        # Used for resetting the environment
+        self.initial_job_shop_graph = deepcopy(job_shop_graph)
+
+        self.dispatcher = Dispatcher(
+            job_shop_graph.instance,
+            ready_operations_filter=ready_operations_filter,
+        )
+
+        # Observers added to track the environment state
+        self.composite_observer = (
+            CompositeFeatureObserver.from_feature_observer_configs(
+                self.dispatcher, feature_observer_configs
+            )
+        )
+        self.graph_updater = graph_updater_config.class_type(
+            dispatcher=self.dispatcher,
+            job_shop_graph=job_shop_graph,
+            **graph_updater_config.kwargs,
+        )
+        self.reward_function = reward_function_config.class_type(
+            dispatcher=self.dispatcher, **reward_function_config.kwargs
+        )
+        self.action_space = gym.spaces.MultiDiscrete(
+            [self.instance.num_jobs, self.instance.num_machines], start=[0, -1]
+        )
+        self.observation_space: gym.spaces.Dict = self._get_observation_space()
+        self.render_mode = render_mode
+        if render_config is None:
+            render_config = {}
+        self.gantt_chart_creator = GanttChartCreator(
+            dispatcher=self.dispatcher, **render_config
+        )
+        self.use_padding = use_padding
+
+    @property
+    def instance(self) -> JobShopInstance:
+        """Returns the instance the environment is working on."""
+        return self.job_shop_graph.instance
+
+    @property
+    def job_shop_graph(self) -> JobShopGraph:
+        """Returns the job shop graph."""
+        return self.graph_updater.job_shop_graph
+
+    def _get_observation_space(self) -> gym.spaces.Dict:
+        """Returns the observation space dictionary."""
+        num_edges = self.job_shop_graph.num_edges
+        dict_space: dict[str, gym.Space] = {
+            ObservationSpaceKey.REMOVED_NODES.value: gym.spaces.MultiBinary(
+                len(self.job_shop_graph.nodes)
+            ),
+            ObservationSpaceKey.EDGE_INDEX.value: gym.spaces.MultiDiscrete(
+                np.full(
+                    (2, num_edges),
+                    fill_value=len(self.job_shop_graph.nodes) + 1,
+                    dtype=np.int32,
+                ),
+                start=np.full(
+                    (2, num_edges),
+                    fill_value=-1,  # -1 is used for padding
+                    dtype=np.int32,
+                ),
+            ),
+        }
+        for feature_type, matrix in self.composite_observer.features.items():
+            dict_space[feature_type.value] = gym.spaces.Box(
+                low=-np.inf, high=np.inf, shape=matrix.shape
+            )
+        return gym.spaces.Dict(dict_space)
+
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[ObservationDict, dict]:
+        """Resets the environment."""
+        super().reset(seed=seed, options=options)
+        self.dispatcher.reset()
+        obs = self.get_observation()
+        return obs, {}
+
+    def step(
+        self, action: tuple[int, int]
+    ) -> tuple[ObservationDict, float, bool, bool, dict[str, Any]]:
+        """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.
+                - "available_operations": The operations that are ready to be
+                  scheduled.
+        """
+        job_id, machine_id = action
+        operation = self.dispatcher.next_operation(job_id)
+        if machine_id == -1:
+            machine_id = operation.machine_id
+
+        self.dispatcher.dispatch(operation, machine_id)
+
+        obs = self.get_observation()
+        reward = self.reward_function.last_reward
+        done = self.dispatcher.schedule.is_complete()
+        truncated = False
+        info: dict[str, Any] = {
+            "feature_names": self.composite_observer.column_names,
+            "available_operations": self.dispatcher.ready_operations(),
+        }
+        return obs, reward, done, truncated, info
+
+    def get_observation(self) -> ObservationDict:
+        """Returns the current observation of the environment."""
+        observation: ObservationDict = {
+            ObservationSpaceKey.REMOVED_NODES.value: np.array(
+                self.job_shop_graph.removed_nodes, dtype=bool
+            ),
+            ObservationSpaceKey.EDGE_INDEX.value: self._get_edge_index(),
+        }
+        for feature_type, matrix in self.composite_observer.features.items():
+            observation[feature_type.value] = matrix
+        return observation
+
+    def _get_edge_index(self) -> NDArray[np.int32]:
+        """Returns the edge index matrix."""
+        edge_index = np.array(
+            self.job_shop_graph.graph.edges(), dtype=np.int32
+        ).T
+
+        if self.use_padding:
+            output_shape = self.observation_space[
+                ObservationSpaceKey.EDGE_INDEX.value
+            ].shape
+            assert output_shape is not None  # For the type checker
+            edge_index = add_padding(
+                edge_index, output_shape=output_shape, dtype=np.int32
+            )
+        return edge_index
+
+    def render(self):
+        """Renders the environment.
+
+        The rendering mode is set by the `render_mode` attribute:
+
+        - human: Renders the current Gannt chart.
+        - save_video: Saves a video of the Gantt chart. Used only if the
+            schedule is completed.
+        - save_gif: Saves a GIF of the Gantt chart. Used only if the schedule
+            is completed.
+        """
+        if self.render_mode == "human":
+            self.gantt_chart_creator.plot_gantt_chart()
+            plt.show(block=False)
+        elif self.render_mode == "save_video":
+            self.gantt_chart_creator.create_video()
+        elif self.render_mode == "save_gif":
+            self.gantt_chart_creator.create_gif()
+
+
+if __name__ == "__main__":
+    from job_shop_lib.dispatching.feature_observers import (
+        FeatureObserverType,
+        FeatureType,
+    )
+    from job_shop_lib.graphs import build_disjunctive_graph
+    from job_shop_lib.benchmarking import load_benchmark_instance
+
+    instance = load_benchmark_instance("ft06")
+    job_shop_graph_ = build_disjunctive_graph(instance)
+    feature_observer_configs_ = [
+        DispatcherObserverConfig(
+            FeatureObserverType.IS_READY,
+            kwargs={"feature_types": [FeatureType.JOBS]},
+        )
+    ]
+
+    env = SingleJobShopGraphEnv(
+        job_shop_graph=job_shop_graph_,
+        feature_observer_configs=feature_observer_configs_,
+        render_mode="save_video",
+        render_config={"video_config": {"fps": 4}},
+    )

job_shop_lib/reinforcement_learning/_types_and_constants.py

@@ -0,0 +1,61 @@
+"""Contains types and enumerations used in the reinforcement learning
+module."""
+
+from enum import Enum
+from typing import TypedDict
+
+import numpy as np
+
+from job_shop_lib.dispatching.feature_observers import FeatureType
+from job_shop_lib.visualization import (
+    GanttChartWrapperConfig,
+    GifConfig,
+    VideoConfig,
+)
+
+
+class RenderConfig(TypedDict, total=False):
+    """Configuration needed to initialize the `GanttChartCreator` class."""
+
+    gantt_chart_wrapper_config: GanttChartWrapperConfig
+    video_config: VideoConfig
+    gif_config: GifConfig
+
+
+class ObservationSpaceKey(str, Enum):
+    """Enumeration of the keys for the observation space dictionary."""
+
+    REMOVED_NODES = "removed_nodes"
+    EDGE_INDEX = "edge_index"
+    OPERATIONS = FeatureType.OPERATIONS.value
+    JOBS = FeatureType.JOBS.value
+    MACHINES = FeatureType.MACHINES.value
+
+
+class _ObservationDictRequired(TypedDict):
+    """Required fields for the observation dictionary."""
+
+    removed_nodes: np.ndarray
+    edge_index: np.ndarray
+
+
+class _ObservationDictOptional(TypedDict, total=False):
+    """Optional fields for the observation dictionary."""
+
+    operations: np.ndarray
+    jobs: np.ndarray
+    machines: np.ndarray
+
+
+class ObservationDict(_ObservationDictRequired, _ObservationDictOptional):
+    """A dictionary containing the observation of the environment.
+
+    Required fields:
+        removed_nodes (np.ndarray): Binary vector indicating removed nodes.
+        edge_index (np.ndarray): Edge list in COO format.
+
+    Optional fields:
+        operations (np.ndarray): Matrix of operation features.
+        jobs (np.ndarray): Matrix of job features.
+        machines (np.ndarray): Matrix of machine features.
+    """

job_shop_lib/reinforcement_learning/_utils.py

@@ -0,0 +1,96 @@
+"""Utility functions for reinforcement learning."""
+
+from typing import TypeVar, Any
+
+import numpy as np
+from numpy.typing import NDArray
+
+from job_shop_lib.exceptions import ValidationError
+
+T = TypeVar("T", bound=np.number)
+
+
+def add_padding(
+    array: NDArray[Any],
+    output_shape: tuple[int, ...],
+    padding_value: float = -1,
+    dtype: type[T] | None = None,
+) -> NDArray[T]:
+    """Adds padding to the array.
+
+    Pads the input array to the specified output shape with a given padding
+    value. If the ``dtype`` is not specified, the ``dtype`` of the input array
+    is used.
+
+    Args:
+        array:
+            The input array to be padded.
+        output_shape:
+            The desired shape of the output array.
+        padding_value:
+            The value to use for padding. Defaults to -1.
+        dtype:
+            The data type for the output array. Defaults to ``None``, in which
+            case the dtype of the input array is used.
+
+    Returns:
+        The padded array with the specified output shape.
+
+    Raises:
+        ValidationError:
+            If the output shape is smaller than the input shape.
+
+    Examples:
+
+    .. doctest::
+
+        >>> array = np.array([[1, 2], [3, 4]])
+        >>> add_padding(array, (3, 3))
+        array([[ 1,  2, -1],
+               [ 3,  4, -1],
+               [-1, -1, -1]])
+
+        >>> add_padding(array, (3, 3), padding_value=0)
+        array([[1, 2, 0],
+               [3, 4, 0],
+               [0, 0, 0]])
+
+        >>> bool_array = np.array([[True, False], [False, True]])
+        >>> add_padding(bool_array, (3, 3), padding_value=False, dtype=int)
+        array([[1, 0, 0],
+               [0, 1, 0],
+               [0, 0, 0]])
+
+        >>> add_padding(bool_array, (3, 3), dtype=int)
+        array([[ 1,  0, -1],
+               [ 0,  1, -1],
+               [-1, -1, -1]])
+    """
+
+    if np.any(np.less(output_shape, array.shape)):
+        raise ValidationError(
+            "Output shape must be greater than the input shape. "
+            f"Got output shape: {output_shape}, input shape: {array.shape}."
+        )
+
+    if dtype is None:
+        dtype = array.dtype.type
+
+    padded_array = np.full(
+        output_shape,
+        fill_value=padding_value,
+        dtype=dtype,
+    )
+
+    if array.size == 0:
+        return padded_array
+
+    slices = tuple(slice(0, dim) for dim in array.shape)
+    padded_array[slices] = array
+    return padded_array
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

job_shop_lib/visualization/__init__.py

@@ -1,25 +1,41 @@
 """Package for visualization."""
 
-from job_shop_lib.visualization.gantt_chart import plot_gantt_chart
-from job_shop_lib.visualization.create_gif import (
+from job_shop_lib.visualization._gantt_chart import plot_gantt_chart
+from job_shop_lib.visualization._gantt_chart_video_and_gif_creation import (
     create_gif,
+    create_gantt_chart_video,
     create_gantt_chart_frames,
     plot_gantt_chart_wrapper,
+    create_video_from_frames,
     create_gif_from_frames,
 )
-from job_shop_lib.visualization.disjunctive_graph import plot_disjunctive_graph
-from job_shop_lib.visualization.agent_task_graph import (
+from job_shop_lib.visualization._disjunctive_graph import (
+    plot_disjunctive_graph,
+)
+from job_shop_lib.visualization._agent_task_graph import (
     plot_agent_task_graph,
     three_columns_layout,
 )
+from job_shop_lib.visualization._gantt_chart_creator import (
+    GanttChartCreator,
+    GanttChartWrapperConfig,
+    GifConfig,
+    VideoConfig,
+)
 
 __all__ = [
     "plot_gantt_chart",
+    "create_gantt_chart_video",
     "create_gif",
     "create_gantt_chart_frames",
     "plot_gantt_chart_wrapper",
     "create_gif_from_frames",
+    "create_video_from_frames",
     "plot_disjunctive_graph",
     "plot_agent_task_graph",
     "three_columns_layout",
+    "GanttChartCreator",
+    "GanttChartWrapperConfig",
+    "GifConfig",
+    "VideoConfig",
 ]

job_shop_lib/visualization/{agent_task_graph.py → _agent_task_graph.py}

@@ -51,6 +51,7 @@ def plot_agent_task_graph(
 
     # Create the networkx graph
     graph = job_shop_graph.graph
+    nodes = job_shop_graph.non_removed_nodes()
 
     # Create the layout if it was not provided
     if layout is None:
@@ -64,16 +65,17 @@ def plot_agent_task_graph(
             job_shop_graph.nodes_by_type[NodeType.MACHINE]
         )
     }
+
     node_colors = [
         _get_node_color(node, machine_colors) for node in job_shop_graph.nodes
-    ]
+    ]  # We need to get the color of all nodes to avoid an index error
     node_shapes = {"machine": "s", "job": "d", "operation": "o", "global": "o"}
 
     # Draw nodes with different shapes based on their type
     for node_type, shape in node_shapes.items():
         current_nodes = [
             node.node_id
-            for node in job_shop_graph.nodes
+            for node in nodes
             if node.node_type.name.lower() == node_type
         ]
         nx.draw_networkx_nodes(
@@ -90,9 +92,7 @@ def plot_agent_task_graph(
     # Draw edges
     nx.draw_networkx_edges(graph, layout, ax=ax)
 
-    node_labels = {
-        node.node_id: _get_node_label(node) for node in job_shop_graph.nodes
-    }
+    node_labels = {node.node_id: _get_node_label(node) for node in nodes}
     nx.draw_networkx_labels(graph, layout, node_labels, ax=ax)
 
     ax.set_axis_off()
@@ -181,10 +181,29 @@ def three_columns_layout(
 
     x_positions = _get_x_positions(leftmost_position, rightmost_position)
 
-    operation_nodes = job_shop_graph.nodes_by_type[NodeType.OPERATION]
-    machine_nodes = job_shop_graph.nodes_by_type[NodeType.MACHINE]
-    job_nodes = job_shop_graph.nodes_by_type[NodeType.JOB]
-    global_nodes = job_shop_graph.nodes_by_type[NodeType.GLOBAL]
+    operation_nodes = [
+        node
+        for node in job_shop_graph.nodes_by_type[NodeType.OPERATION]
+        if not job_shop_graph.is_removed(node)
+    ]
+    machine_nodes = [
+        node
+        for node in job_shop_graph.nodes_by_type[NodeType.MACHINE]
+        if not job_shop_graph.is_removed(node)
+    ]
+    job_nodes = [
+        node
+        for node in job_shop_graph.nodes_by_type[NodeType.JOB]
+        if not job_shop_graph.is_removed(node)
+    ]
+    global_nodes = [
+        node
+        for node in job_shop_graph.nodes_by_type[NodeType.GLOBAL]
+        if not job_shop_graph.is_removed(node)
+    ]
+
+    # job_nodes = job_shop_graph.nodes_by_type[NodeType.JOB]
+    # global_nodes = job_shop_graph.nodes_by_type[NodeType.GLOBAL]
 
     total_positions = len(operation_nodes) + len(global_nodes) * 2
     y_spacing = (topmost_position - bottommost_position) / total_positions