[Feat] Add JSSP environment #177

LTluttmann · 2024-05-16T08:35:47Z

Description

Added environment for the Job-Shop Scheduling Problem (JSSP). This implementation of JSSP treats JSSP as a special case of FJSP, where each operation can be processed by only one machine. As such, the environment is implemented as subclass of the FJSPEnv, only changing the action space (action space reduces to the selection of next job to execute) and the data generator.
In addition, the HetGNN policy has been restructured and renamed to L2D (for learning to dispatch) and is now applicable to both, FJSP and JSSP.

Motivation and Context

JSSP is a common CO problem and a widely used problem to benchmark new algorithms
I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds core functionality)
Breaking change (fix or feature that would cause existing functionality to change)
Documentation (update in the documentation)
Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

…ta generator

…ch) model

fedebotu · 2024-05-16T08:59:57Z

rl4co/envs/scheduling/jssp/env.py

+from .generator import JSSPFileGenerator, JSSPGenerator
+
+
+class JSSPEnv(FJSPEnv):


Awesome! I love that the code is being reused in such as smart way

Yesss, when see this inherit: wow nice.

fedebotu · 2024-05-16T09:00:46Z

rl4co/models/__init__.py

@@ -19,7 +19,7 @@
 from rl4co.models.rl.ppo.ppo import PPO
 from rl4co.models.rl.reinforce.baselines import REINFORCEBaseline, get_reinforce_baseline
 from rl4co.models.rl.reinforce.reinforce import REINFORCE
-from rl4co.models.zoo import HetGNNModel
+from rl4co.models.zoo import L2DModel


Good - let's make sure the baselines have their names! L2D is a very influential paper in the NCO community

fedebotu · 2024-05-16T09:02:32Z

rl4co/utils/ops.py

@@ -73,7 +73,8 @@ def gather_by_index(src, idx, dim=1, squeeze=True):
    expanded_shape = list(src.shape)
    expanded_shape[dim] = -1
    idx = idx.view(idx.shape + (1,) * (src.dim() - idx.dim())).expand(expanded_shape)
-    return src.gather(dim, idx).squeeze() if squeeze else src.gather(dim, idx)
+    squeeze = idx.size(dim) == 1 and squeeze
+    return src.gather(dim, idx).squeeze(dim) if squeeze else src.gather(dim, idx)


 def unbatchify_and_gather(x: Tensor, idx: Tensor, n: int):


[Minor] it might be faster if you put a @torch.jit decorator. Not 100% sure though

fedebotu · 2024-05-16T09:04:07Z

@Junyoungpark tagging you since you know this problem very well!

Do you think we have a chance to include ScheduleNet?

cbhua

Great job!

cbhua · 2024-05-26T14:46:12Z

rl4co/envs/scheduling/fjsp/env.py

+        # update adjacency matrices (remove edges)
+        td["proc_times"] = td["proc_times"].scatter(
+            2,
+            selected_op[:, None, None].expand(-1, self.num_mas, 1),


[Minor, Enhancement] Using einops.repeat could be "slightly" more efficient 😁:

repeat(selected_op, 'b -> b n d', n=self.num_mas, d=1)

Is it though? I think if you know already the dimensions, einops is slightly slower from what I know, but take this with a grain of salt. But I agree that it's more readable

Okay I did a trial and actually einops.repeat is way slower than tensor.expand 😂 (in large scale around 4x slower). Then I think it's good to keep using the tensor.expand.

yeah I think they use very clever optimizations in torch.expand() 😄

cbhua · 2024-05-26T14:57:27Z

rl4co/envs/scheduling/jssp/__init__.py

[Minor] Maybe we don't need this file for clean file structure.

cbhua · 2024-05-26T15:00:00Z

rl4co/envs/scheduling/jssp/env.py

+from .generator import JSSPFileGenerator, JSSPGenerator
+
+
+class JSSPEnv(FJSPEnv):


Yesss, when see this inherit: wow nice.

cbhua · 2024-05-26T15:29:12Z

rl4co/models/nn/graph/gcn.py

+    # self-loop is added by GCNConv layer
+    return get_full_graph_edge_index(td.device, num_nodes, self_loop=False)
+
+
 class GCNEncoder(nn.Module):


I like this clean refactoring, the logic is clearer. But will the get_full_graph_edge_index() be called at every forward step? i.e. in the previous version, if it's a fully connected connected graph, the edge_index will be saved as class variable, instead of regenerating every time.

thats true, but the result is cached so it should not be too slow. In fact, this implementation should be much faster as before (at least it was in my experiments), because it avoids the list comprehension over the batch data within the forward pass. But I agree, its still not optimal; I will revisit this in the near future

…and also adds self attn.

… attn class to nn/attention.

…etter training stability

LTluttmann · 2024-05-28T17:16:30Z

Thanks for reviewing guys. I will add a ton of changes here in a couple of minutes and I hope this PR will not get too messy. Let me know if we should go through it together.

Additional changes:

stepwise PPO for L2D
Some MatNet changes to better work for JSSP / FJSP
Running mean / variance class for reward / advantage scaling
stepwise L2D policy
attention based models for fjsp / jssp
minor bugfixes and improvements here and there

… to nn/ops

…by pytorch geometric

fedebotu · 2024-06-01T02:55:58Z

Wow, lots of changes here! 😁 Really curious about episodic / stepwise RL performances

Btw, feel free to merge anytime

fedebotu · 2024-06-01T07:19:35Z

rl4co/envs/routing/tsp/env.py

+# NOTE Experimental TSP class for stepwise PPO
+
+
+class TSPEnv4PPO(TSPEnv):


[Minor] this may be called DenseRewardTSPEnv or similar?

Sure! Btw., Stepwise PPO for TSP indeed converges to the nearest neighbor heuristic, at least with the stepwise reward as it is defined here (the distance added by the action):

Do you have a preferrence as to how to call the stepwise PPO in the paper (dense, stepwise something else)? And then we should probably adjust the description about PPO in the appendix

fedebotu · 2024-06-01T07:20:25Z

rl4co/envs/scheduling/ffsp/env.py

@@ -58,7 +57,7 @@ def __init__(
        generator_params: dict = {},
        **kwargs,
    ):
-        super().__init__(**kwargs)
+        super().__init__(check_solution=False, **kwargs)


Is this always the case (no solution check)?

uhm I think this is there just bc there is no check implemented for FFSP yet haha. Let me see if I can get one implemented

No worries, this is not a pressing issue! Actually, it's even better to keep it to False during training for efficiency

fedebotu · 2024-06-03T18:01:10Z

Great job!! This PR is truly huge

LTluttmann added 11 commits May 13, 2024 16:08

[Feat] added JSSP as a special case of FJSP -> same env, different da…

8a5d8a6

…ta generator

[Minor] added JSSP as valid environment for HetGNN model

5006253

[BugFix] in init of jssp generator

23f4a9f

[BugFix] in init of jssp generator

1c56df3

[Feat] added L2D model

89d6732

[Feat] integrated fjsp and jssp into a single l2d (learning to dispat…

14b9011

…ch) model

[Feat] integrated fjsp and jssp into a single l2d (learning to dispat…

55b7510

…ch) model

[Minor] custom load_data vor jssp and fjsp envs to load benchmark data

beb5b41

Merge branch 'main' into fjsp

614127b

[Minor] updated scheduling nb

44cfbdb

[Minor] renamed jssp generator

86759c3

LTluttmann requested review from cbhua and fedebotu May 16, 2024 08:35

fedebotu approved these changes May 16, 2024

View reviewed changes

fedebotu requested a review from Junyoungpark May 16, 2024 09:03

cbhua approved these changes May 26, 2024

View reviewed changes

LTluttmann added 8 commits May 28, 2024 18:57

[Minor] Merge branch 'main' into fjsp

9bebb69

[Minor] configs for scheduling experiments

f976485

[Minor] made jssp/fjsp envs applicable to stepwise ppo

7233994

[Feat] more efficient matnet encoder, computing cross attn only once …

ef45411

…and also adds self attn.

[BugFix] reordered args in FFSP matnet

5bdb722

[Minor] modularized MatNet encoder to be better reusable; added cross…

73e3fb1

… attn class to nn/attention.

[Minor] env embeddings for jssp / fjsp

be13834

[Feat] add running moments class for reward / advantage scaling and b…

35828d6

…etter training stability

LTluttmann added 4 commits May 28, 2024 19:17

[Feat] added the ffn sublayer from the transformer (skipcon+norm+mlp)…

08b45b0

… to nn/ops

[BugFix] also batchify TensorDicts in am decoder

6ffa3c4

[Feat] added stepwise PPO

076da30

[Feat] more models for L2D (FJSP/JSSP)

bd427f7

LTluttmann added 13 commits May 28, 2024 19:22

[Feat] more models for L2D (FJSP/JSSP)

ab5a80c

[Minor] allow for gradient loggin with wandb

a681bdc

[Minor] added fn to transform adj matrix to edge indices as required …

512dae2

…by pytorch geometric

[Minor] configs for scheduling experiments

c2c0b8f

[Minor] merge main

7fdf21c

[Minor] test for stepwise ppo / l2d

16a7974

[Minor] restructured scheduling experiments

9fea91a

[Minor] more tests for JSSP / FJSP

7c2ab66

[Minor] refactored actor / decoder logic in L2D

1d17981

[Minor] minor fixes

56c68f3

[Minor] experiment refinement

091f70d

[Feat] stepwise PPO for TSP

8cfb9a8

[Feat] stepwise PPO for TSP

23d0323

[Minor] merge main

60452fd

fedebotu removed the request for review from Junyoungpark June 1, 2024 07:17

fedebotu approved these changes Jun 1, 2024

View reviewed changes

LTluttmann added 4 commits June 3, 2024 18:44

[Config] full precision for ppo-tsp (works better somehow)

2ea9d3e

[Minor] Merge branch 'main' into fjsp

117d082

[Minor] rename TSP env with stepwise reward + add docstring

694b1b8

[Minor] rename TSP env with stepwise reward + add docstring

84f39fb

LTluttmann merged commit f4abe1b into main Jun 3, 2024
26 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feat] Add JSSP environment #177

[Feat] Add JSSP environment #177

LTluttmann commented May 16, 2024

fedebotu May 16, 2024

cbhua May 26, 2024

fedebotu May 16, 2024

fedebotu May 16, 2024

fedebotu commented May 16, 2024

cbhua left a comment

cbhua May 26, 2024

fedebotu May 26, 2024

cbhua May 27, 2024

LTluttmann May 28, 2024

cbhua May 26, 2024

cbhua May 26, 2024

cbhua May 26, 2024

LTluttmann May 28, 2024

LTluttmann commented May 28, 2024

fedebotu commented Jun 1, 2024

fedebotu Jun 1, 2024

LTluttmann Jun 1, 2024

fedebotu Jun 1, 2024

LTluttmann Jun 1, 2024

fedebotu Jun 1, 2024

fedebotu commented Jun 3, 2024

		from .generator import JSSPFileGenerator, JSSPGenerator


		class JSSPEnv(FJSPEnv):

		# NOTE Experimental TSP class for stepwise PPO


		class TSPEnv4PPO(TSPEnv):

[Feat] Add JSSP environment #177

[Feat] Add JSSP environment #177

Conversation

LTluttmann commented May 16, 2024

Description

Motivation and Context

Types of changes

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fedebotu commented May 16, 2024

cbhua left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LTluttmann commented May 28, 2024

fedebotu commented Jun 1, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fedebotu commented Jun 3, 2024