[Fix] fixed JSSP import issues

ai4co · Sep 12, 2024 · b516158 · b516158
1 parent 09f91ba
commit b516158
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/rl4co/envs/__init__.py b/rl4co/envs/__init__.py
@@ -4,6 +4,9 @@
 # EDA
 from rl4co.envs.eda import DPPEnv, MDPPEnv
 
+# Graph
+from rl4co.envs.graph import FLPEnv, MCPEnv
+
 # Routing
 from rl4co.envs.routing import (
     ATSPEnv,
@@ -25,11 +28,7 @@
 )
 
 # Scheduling
-from rl4co.envs.scheduling import FFSPEnv, FJSPEnv, SMTWTPEnv
-from rl4co.envs.scheduling.jssp.env import JSSPEnv
-
-# Graph
-from rl4co.envs.graph import MCPEnv, FLPEnv
+from rl4co.envs.scheduling import FFSPEnv, FJSPEnv, JSSPEnv, SMTWTPEnv
 
 # Register environments
 ENV_REGISTRY = {

diff --git a/rl4co/envs/scheduling/__init__.py b/rl4co/envs/scheduling/__init__.py
@@ -1,5 +1,4 @@
-from rl4co.envs.scheduling.ffsp.env import FFSPEnv
-from rl4co.envs.scheduling.fjsp.env import FJSPEnv
-
-# from rl4co.envs.scheduling.jssp.env import JSSPEnv
-from rl4co.envs.scheduling.smtwtp.env import SMTWTPEnv
+from .ffsp.env import FFSPEnv
+from .fjsp.env import FJSPEnv
+from .jssp.env import JSSPEnv
+from .smtwtp.env import SMTWTPEnv
diff --git a/rl4co/envs/scheduling/jssp/env.py b/rl4co/envs/scheduling/jssp/env.py
@@ -4,7 +4,7 @@
 from tensordict import TensorDict
 from torch._tensor import Tensor
 
-from rl4co.envs import FJSPEnv
+from rl4co.envs.scheduling.fjsp.env import FJSPEnv
 from rl4co.utils.ops import gather_by_index
 
 from .generator import JSSPFileGenerator, JSSPGenerator