Pre-commit

lihuoran · lihuoran · commit b4db84214de1 · 2023-02-06T11:44:28.000+08:00
diff --git a/maro/rl/model/fc_block.py b/maro/rl/model/fc_block.py
@@ -72,9 +72,13 @@ def __init__(
 
         # build the net
         dims = [self._input_dim] + self._hidden_dims
-        layers = [self._build_layer(in_dim, out_dim, activation=self._activation) for in_dim, out_dim in zip(dims, dims[1:])]
+        layers = [
+            self._build_layer(in_dim, out_dim, activation=self._activation) for in_dim, out_dim in zip(dims, dims[1:])
+        ]
         # top layer
-        layers.append(self._build_layer(dims[-1], self._output_dim, head=self._head, activation=self._output_activation))
+        layers.append(
+            self._build_layer(dims[-1], self._output_dim, head=self._head, activation=self._output_activation),
+        )
 
         self._net = nn.Sequential(*layers)
 
@@ -103,7 +107,13 @@ def input_dim(self) -> int:
     def output_dim(self) -> int:
         return self._output_dim
 
-    def _build_layer(self, input_dim: int, output_dim: int, head: bool = False, activation: Type[torch.nn.Module] = None) -> nn.Module:
+    def _build_layer(
+        self,
+        input_dim: int,
+        output_dim: int,
+        head: bool = False,
+        activation: Type[torch.nn.Module] = None,
+    ) -> nn.Module:
         """Build a basic layer.
 
         BN -> Linear -> Activation -> Dropout
diff --git a/tests/rl/gym_wrapper/rl_component_bundle.py b/tests/rl/gym_wrapper/rl_component_bundle.py
@@ -6,11 +6,9 @@
 from maro.rl.rl_component.rl_component_bundle import RLComponentBundle
 from maro.simulator import Env
 
-from tests.rl.gym_wrapper.simulator.business_engine import GymBusinessEngine
-
 from .config import algorithm, env_conf
 from .env_sampler import GymEnvSampler
-
+from tests.rl.gym_wrapper.simulator.business_engine import GymBusinessEngine
 
 learn_env = Env(business_engine_cls=GymBusinessEngine, **env_conf)
 test_env = learn_env
@@ -44,7 +42,14 @@
     from tests.rl.algorithms.sac import get_sac_policy, get_sac_trainer
 
     policies = [
-        get_sac_policy(f"{algorithm}_{i}.policy", action_lower_bound, action_upper_bound, gym_state_dim, gym_action_dim, action_limit)
+        get_sac_policy(
+            f"{algorithm}_{i}.policy",
+            action_lower_bound,
+            action_upper_bound,
+            gym_state_dim,
+            gym_action_dim,
+            action_limit,
+        )
         for i in range(num_agents)
     ]
     trainers = [get_sac_trainer(f"{algorithm}_{i}", gym_state_dim, gym_action_dim) for i in range(num_agents)]
