[RLlib] LINT: Enable ruff imports for offline/, tests/, callbacks/ and env/ in rllib (#56744)

Author: czgdp1807

pyproject.toml

@@ -74,10 +74,10 @@ afterray = ["psutil", "setproctitle"]
 "rllib/algorithms/*" = ["I"]
 # "rllib/core/*" = ["I"]
 # "rllib/examples/*" = ["I"]
-"rllib/offline/*" = ["I"]
-"rllib/tests/*" = ["I"]
-"rllib/callbacks/*" = ["I"]
-"rllib/env/*" = ["I"]
+# "rllib/offline/*" = ["I"]
+# "rllib/tests/*" = ["I"]
+# "rllib/callbacks/*" = ["I"]
+# "rllib/env/*" = ["I"]
 # "rllib/execution/*" = ["I"]
 # "rllib/policy/*" = ["I"]
 # "rllib/tuned_examples/*" = ["I"]

rllib/callbacks/callbacks.py

@@ -17,9 +17,9 @@
 from ray.rllib.policy.sample_batch import SampleBatch
 from ray.rllib.utils.annotations import (
     OldAPIStack,
-    override,
     OverrideToImplementCustomLogic,
     PublicAPI,
+    override,
 )
 from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
 from ray.rllib.utils.typing import AgentID, EnvType, EpisodeType, PolicyID

rllib/callbacks/tests/test_callbacks_old_api_stack.py

@@ -1,5 +1,5 @@
-from collections import Counter
 import unittest
+from collections import Counter
 
 import ray
 from ray.rllib.algorithms.callbacks import DefaultCallbacks, make_multi_callbacks
@@ -209,7 +209,8 @@ def test_on_episode_created(self):
 
 
 if __name__ == "__main__":
-    import pytest
     import sys
 
+    import pytest
+
     sys.exit(pytest.main(["-v", __file__]))

rllib/callbacks/tests/test_callbacks_on_algorithm.py

@@ -4,8 +4,8 @@
 
 import ray
 from ray import tune
-from ray.rllib.callbacks.callbacks import RLlibCallback
 from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.callbacks.callbacks import RLlibCallback
 from ray.rllib.examples.envs.classes.cartpole_crashing import CartPoleCrashing
 from ray.rllib.utils.test_utils import check
 
@@ -108,7 +108,8 @@ def test_on_init_and_checkpoint_loaded(self):
 
 
 if __name__ == "__main__":
-    import pytest
     import sys
 
+    import pytest
+
     sys.exit(pytest.main(["-v", __file__]))

rllib/callbacks/tests/test_callbacks_on_env_runner.py

@@ -1,12 +1,12 @@
-from collections import Counter
 import unittest
+from collections import Counter
 
 import gymnasium as gym
 
 import ray
 from ray import tune
-from ray.rllib.callbacks.callbacks import RLlibCallback
 from ray.rllib.algorithms.ppo import PPOConfig
+from ray.rllib.callbacks.callbacks import RLlibCallback
 from ray.rllib.env.env_runner import EnvRunner
 from ray.rllib.env.vector.vector_multi_agent_env import VectorMultiAgentEnv
 from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole
@@ -237,7 +237,8 @@ def test_tune_trial_id_visible_in_callbacks(self):
 
 
 if __name__ == "__main__":
-    import pytest
     import sys
 
+    import pytest
+
     sys.exit(pytest.main(["-v", __file__]))

rllib/callbacks/tests/test_multicallback.py

@@ -1,4 +1,5 @@
 import unittest
+
 import ray
 from ray.rllib.algorithms import PPOConfig
 from ray.rllib.callbacks.callbacks import RLlibCallback
@@ -141,7 +142,8 @@ def test_single_callback_validation_error(self):
 
 
 if __name__ == "__main__":
-    import pytest
     import sys
 
+    import pytest
+
     sys.exit(pytest.main(["-v", __file__]))

rllib/env/__init__.py

@@ -6,12 +6,10 @@
 from ray.rllib.env.policy_client import PolicyClient
 from ray.rllib.env.remote_base_env import RemoteBaseEnv
 from ray.rllib.env.vector_env import VectorEnv
-
-from ray.rllib.env.wrappers.dm_env_wrapper import DMEnv
 from ray.rllib.env.wrappers.dm_control_wrapper import DMCEnv
+from ray.rllib.env.wrappers.dm_env_wrapper import DMEnv
 from ray.rllib.env.wrappers.group_agents_wrapper import GroupAgentsWrapper
-from ray.rllib.env.wrappers.pettingzoo_env import PettingZooEnv
-from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv
+from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv, PettingZooEnv
 from ray.rllib.env.wrappers.unity3d_env import Unity3DEnv
 
 INPUT_ENV_SPACES = "__env__"

rllib/env/base_env.py

@@ -1,7 +1,8 @@
 import logging
-from typing import Callable, Tuple, Optional, List, Dict, Any, TYPE_CHECKING, Union, Set
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
 import gymnasium as gym
+
 import ray
 from ray.rllib.utils.annotations import OldAPIStack
 from ray.rllib.utils.typing import AgentID, EnvID, EnvType, MultiEnvDict
@@ -368,9 +369,9 @@ def convert_to_base_env(
         The resulting BaseEnv object.
     """
 
-    from ray.rllib.env.remote_base_env import RemoteBaseEnv
     from ray.rllib.env.external_env import ExternalEnv
     from ray.rllib.env.multi_agent_env import MultiAgentEnv
+    from ray.rllib.env.remote_base_env import RemoteBaseEnv
     from ray.rllib.env.vector_env import VectorEnv, VectorEnvWrapper
 
     if remote_envs and num_envs == 1:

rllib/env/env_runner.py

@@ -1,6 +1,6 @@
 import abc
 import logging
-from typing import Any, Dict, Optional, Tuple, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
 
 import gymnasium as gym
 import tree  # pip install dm_tree
@@ -15,7 +15,7 @@
 from ray.rllib.utils.metrics.metrics_logger import MetricsLogger
 from ray.rllib.utils.torch_utils import convert_to_torch_tensor
 from ray.rllib.utils.typing import StateDict, TensorType
-from ray.util.annotations import PublicAPI, DeveloperAPI
+from ray.util.annotations import DeveloperAPI, PublicAPI
 from ray.util.metrics import Counter
 
 if TYPE_CHECKING:

rllib/env/env_runner_group.py

@@ -1,8 +1,8 @@
-import gymnasium as gym
-import logging
 import importlib.util
+import logging
 import os
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Collection,
@@ -11,12 +11,17 @@
     Optional,
     Tuple,
     Type,
-    TYPE_CHECKING,
     TypeVar,
     Union,
 )
 
+import gymnasium as gym
+
 import ray
+from ray._common.deprecation import (
+    DEPRECATED_VALUE,
+    deprecation_warning,
+)
 from ray.actor import ActorHandle
 from ray.exceptions import RayActorError
 from ray.rllib.core import (
@@ -28,18 +33,14 @@
 from ray.rllib.core.learner import LearnerGroup
 from ray.rllib.core.rl_module import validate_module_id
 from ray.rllib.core.rl_module.rl_module import RLModuleSpec
-from ray.rllib.evaluation.rollout_worker import RolloutWorker
 from ray.rllib.env.base_env import BaseEnv
 from ray.rllib.env.env_context import EnvContext
 from ray.rllib.env.env_runner import EnvRunner
+from ray.rllib.evaluation.rollout_worker import RolloutWorker
 from ray.rllib.offline import get_dataset_and_shards
 from ray.rllib.policy.policy import Policy, PolicyState
 from ray.rllib.utils.actor_manager import FaultTolerantActorManager
 from ray.rllib.utils.annotations import OldAPIStack
-from ray._common.deprecation import (
-    deprecation_warning,
-    DEPRECATED_VALUE,
-)
 from ray.rllib.utils.framework import try_import_tf
 from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME, WEIGHTS_SEQ_NO
 from ray.rllib.utils.typing import (