diff --git a/rl4co/envs/__init__.py b/rl4co/envs/__init__.py
index f3b65b7d..1c43a7d1 100644
--- 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
index 40b5571e..febb183c 100644
--- 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
index 702ceda7..2381836a 100644
--- 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