Adding JSSP environment #112
Conversation
Feture/unbatched jssp l2d
rl4co/envs/scheduling/jssp.py
Outdated
| init_quality_flag = False | ||
|
|
||
|
|
||
| configs = Configs() |
There was a problem hiding this comment.
About config: is this supposed not to be ever changed by the user?
Because if it can be changed, isn't it better to set these as hyperparameters to the JSSPEnv class?
rl4co/envs/scheduling/jssp.py
Outdated
| def __init__(self, num_jobs, num_machines, **kwargs): | ||
| super().__init__(**kwargs) | ||
| self.batch_size = torch.Size([1]) | ||
| adjacency_spec = DiscreteTensorSpec( |
There was a problem hiding this comment.
[Minor] For specs, I think it would be cleaner to move them to a _make_spec(...) class as done in the tutorial
rl4co/envs/scheduling/jssp.py
Outdated
| def done(self): | ||
| if len(self.partial_sol_sequence) == self.num_tasks: | ||
| return torch.tensor(True) | ||
| return torch.tensor(False) |
There was a problem hiding this comment.
Important note:
In general, whenever there is a torch.tensor(...), torch.ones(...) etc to be instantiated , it is better to do so on the device. Since we use GPU for training and we can move the envs there for offloading the CPU, it would be best to instantiate directly there.
For instance:
...
return torch.tensor(False, device=self.device)The same applies to other tensors instantiated this way - at the moment, trying to set the env on GPU raises some device mismatch error, but it can be easily fixed :)
| self.num_machines = num_machines | ||
| self.num_tasks = self.num_jobs * self.num_machines | ||
| # the task id for first column | ||
| self.first_col = torch.arange( |
There was a problem hiding this comment.
Here it seems that there are some tensors declared inside the __init__ signature; this makes the environment stateful (as explained here) .
I don't hold a grudge against stateful envs, but I wonder whether it would be better (i.e. more efficient) to have them as in the pendulum example, which is, simply passed upon generation? I.e., the TensorDict holds all of the parameters inside?
| self.positive_reward += reward | ||
| self.max_end_time = self.LBs.max() | ||
|
|
||
| tensordict = TensorDict( |
There was a problem hiding this comment.
| "cell_type": "code", | ||
| "outputs": [], | ||
| "source": [ | ||
| "from torchrl.envs.vec_envs import ParallelEnv\n", |
There was a problem hiding this comment.
Important
Check out env instantiation and stepping directly on GPU, not working in the current version (see other comments) - but can be easily fixed
fix: fixed gpu not working
|
Closing PR for now since we have added the batched JSSP. I'd keep an eye on this though since since |
Description
Added the JSSP env from Learning to Dispatch.
Motivation and Context
We wanted to have a JSSP environment from a published paper.
Types of changes
What types of changes does your code introduce? Remove all that do not apply:
Checklist
Go over all the following points, and put an
xin all the boxes that apply.If you are unsure about any of these, don't hesitate to ask. We are here to help!