+ flake8

rl_toolkit/__init__.py

@@ -1 +1 @@
-from .__main__ import main
+from .__main__ import main  # noqa

rl_toolkit/__main__.py

@@ -119,7 +119,9 @@ def main():
                 clip_mean_max=config["Model"]["Actor"]["clip_mean_max"],
                 n_quantiles=config["Model"]["Critic"]["n_quantiles"],
                 merge_index=config["Model"]["Critic"]["merge_index"],
-                top_quantiles_to_drop=config["Model"]["Critic"]["top_quantiles_to_drop"],
+                top_quantiles_to_drop=config["Model"]["Critic"][
+                    "top_quantiles_to_drop"
+                ],
                 n_critics=config["Model"]["Critic"]["count"],
                 gamma=config["Learner"]["gamma"],
                 tau=config["Learner"]["tau"],
@@ -217,7 +219,9 @@ def main():
                 critic_learning_rate=config["Model"]["Critic"]["learning_rate"],
                 alpha_learning_rate=config["Model"]["Alpha"]["learning_rate"],
                 n_quantiles=config["Model"]["Critic"]["n_quantiles"],
-                top_quantiles_to_drop=config["Model"]["Critic"]["top_quantiles_to_drop"],
+                top_quantiles_to_drop=config["Model"]["Critic"][
+                    "top_quantiles_to_drop"
+                ],
                 n_critics=config["Model"]["Critic"]["count"],
                 clip_mean_min=config["Model"]["Actor"]["clip_mean_min"],
                 clip_mean_max=config["Model"]["Actor"]["clip_mean_max"],
@@ -303,4 +307,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

rl_toolkit/agents/dueling_dqn/agent.py

@@ -3,8 +3,8 @@
 import numpy as np
 import reverb
 import tensorflow as tf
-import wandb
 
+import wandb
 from rl_toolkit.networks.models import DuelingDQN
 from rl_toolkit.utils import VariableContainer
 
@@ -200,22 +200,24 @@ def collect(self, writer, policy):
             # Block until all the items have been sent to the server
             writer.end_episode()
 
-            # save the checkpoint
-            if self._total_episodes > 0:
-                if self._episode_reward > self._best_episode_reward:
-                    self._best_episode_reward = self._episode_reward
-                    self.save()
-                    print(
-                        f"Model is saved at {self._total_episodes} episode with score {self._best_episode_reward}"
-                    )
-                    wandb.log({"best_score": self._best_episode_reward}, commit=False)
-            else:
+            # Store best weights
+            if self._episode_reward > self._best_episode_reward:
                 self._best_episode_reward = self._episode_reward
+                self._best_episode = self._total_episodes
+                if self._save_path:
+                    os.makedirs(self._save_path, exist_ok=True)
+                    # Save model
+                    self.model.save_weights(
+                        os.path.join(self._save_path, "best_actor.h5")
+                    )
 
             # Logging
             print("=============================================")
             print(f"Epoch: {self._total_episodes}")
             print(f"Score: {self._episode_reward}")
+            print(
+                f"Best score: {self._best_episode_reward} (at epoch {self._best_episode})"
+            )
             print(f"Steps: {self._episode_steps}")
             print(f"TotalInteractions: {self._total_steps}")
             print(f"Train step: {self._train_step.numpy()}")
@@ -250,6 +252,8 @@ def collect(self, writer, policy):
     def run(self):
         # Init environment
         self._episode_reward = 0.0
+        self._best_episode_reward = float("-inf")
+        self._best_episode = 0
         self._episode_steps = 0
         self._total_episodes = 0
         self._total_steps = 0
@@ -267,18 +271,3 @@ def run(self):
             # Main loop
             while not self._stop_agents:
                 self.collect(writer, self.collect_policy)
-
-    def save(self, path=""):
-        if self._save_path:
-            try:
-                os.makedirs(os.path.join(os.path.join(self._save_path, path)))
-            except OSError:
-                print("The path already exist ❗❗❗")
-            finally:
-                # Save model
-                self.model.save_weights(
-                    os.path.join(
-                        os.path.join(self._save_path, path),
-                        f"dqn_{self._total_episodes}.h5",
-                    )
-                )

rl_toolkit/agents/dueling_dqn/learner.py

@@ -1,9 +1,9 @@
 import reverb
 import tensorflow as tf
-import wandb
 from tensorflow.keras.callbacks import LearningRateScheduler
 from wandb.integration.keras import WandbMetricsLogger
 
+import wandb
 from rl_toolkit.networks.callbacks import DQNAgentCallback, PrintLR, cosine_schedule
 from rl_toolkit.networks.models import DuelingDQN
 from rl_toolkit.utils import make_reverb_dataset

rl_toolkit/agents/dueling_dqn/tester.py

@@ -1,7 +1,7 @@
 import numpy as np
 import tensorflow as tf
-import wandb
 
+import wandb
 from rl_toolkit.networks.models import DuelingDQN
 
 from ...core.process import Process

rl_toolkit/agents/sac/agent.py

@@ -3,8 +3,8 @@
 import numpy as np
 import reverb
 import tensorflow as tf
-import wandb
 
+import wandb
 from rl_toolkit.networks.models import Actor
 from rl_toolkit.utils import VariableContainer
 
@@ -199,7 +199,9 @@ def collect(self, writer, max_steps, policy):
                 print("=============================================")
                 print(f"Epoch: {self._total_episodes}")
                 print(f"Score: {self._episode_reward}")
-                print(f"Best score: {self._best_episode_reward} (at epoch {self._best_episode})")
+                print(
+                    f"Best score: {self._best_episode_reward} (at epoch {self._best_episode})"
+                )
                 print(f"Steps: {self._episode_steps}")
                 print(f"TotalInteractions: {self._total_steps}")
                 print(f"Train step: {self._train_step.numpy()}")
@@ -233,7 +235,7 @@ def collect(self, writer, max_steps, policy):
     def run(self):
         # Init environment
         self._episode_reward = 0.0
-        self._best_episode_reward = float('-inf')
+        self._best_episode_reward = float("-inf")
         self._best_episode = None
         self._episode_steps = 0
         self._total_episodes = 0

rl_toolkit/agents/sac/learner.py

@@ -3,10 +3,10 @@
 import numpy as np
 import reverb
 import tensorflow as tf
-import wandb
 from tensorflow.keras.optimizers import Adam
 from wandb.integration.keras import WandbMetricsLogger
 
+import wandb
 from rl_toolkit.networks.callbacks import SACAgentCallback
 from rl_toolkit.networks.models import ActorCritic
 from rl_toolkit.utils import make_reverb_dataset
@@ -159,9 +159,7 @@ def save(self):
         if self._save_path:
             os.makedirs(self._save_path, exist_ok=True)
             # Save model
-            self.model.save_weights(
-                os.path.join(self._save_path, "actor_critic.h5")
-            )
+            self.model.save_weights(os.path.join(self._save_path, "actor_critic.h5"))
 
     def close(self):
         super(Learner, self).close()

rl_toolkit/agents/sac/tester.py

@@ -1,8 +1,8 @@
 import numpy as np
 import tensorflow as tf
-import wandb
 from dm_control import viewer
 
+import wandb
 from rl_toolkit.networks.models import Actor
 
 from ...core.process import Process

rl_toolkit/networks/callbacks/lr.py

@@ -1,8 +1,9 @@
 import numpy as np
 import tensorflow as tf
-import wandb
 from tensorflow.keras.callbacks import Callback
 
+import wandb
+
 
 def cosine_schedule(base_lr, total_steps, warmup_steps):
     def step_fn(epoch):

rl_toolkit/networks/models/actor.py

@@ -44,9 +44,7 @@ def __init__(
         self.fc_layers = []
 
         for m in units:
-            self.fc_layers.append(
-                Dense(units=m, activation="elu")
-            )
+            self.fc_layers.append(Dense(units=m, activation="elu"))
 
         # Deterministicke akcie
         self.mean = Dense(

rl_toolkit/networks/models/critic.py

@@ -29,9 +29,7 @@ def __init__(self, units: list, n_quantiles: int, merge_index: int, **kwargs):
 
         for i, m in enumerate(units):
             if i != self.merge_index:
-                self.fc_layers.append(
-                    Dense(units=m, activation="elu")
-                )
+                self.fc_layers.append(Dense(units=m, activation="elu"))
             else:
                 self.fc_layers.append(None)  # add empty layer instead of merge layer
 

rl_toolkit/networks/models/dueling.py

@@ -1,16 +1,14 @@
 import tensorflow as tf
 from tensorflow.keras import Model
 from tensorflow.keras.initializers import Orthogonal, TruncatedNormal
-from tensorflow.keras.layers import (
+from tensorflow.keras.layers import (  # GlobalMaxPooling1D,; Lambda,
     Add,
     Dense,
     Dropout,
+    GlobalAveragePooling1D,
     Layer,
     LayerNormalization,
     MultiHeadAttention,
-    GlobalAveragePooling1D,
-    GlobalMaxPooling1D,
-    Lambda,
 )
 
 
@@ -142,11 +140,11 @@ def __init__(
             for _ in range(num_layers)
         ]
 
-	# Reduce
+        # Reduce
         # self.flatten = Lambda(lambda x: x[:, -1])
         # self.flatten = GlobalMaxPooling1D()
         self.flatten = GlobalAveragePooling1D()
-        
+
         # Output
         self.V = Dense(
             1,
@@ -165,7 +163,6 @@ def call(self, inputs, training=None):
         for layer in self.e_layers:
             x = layer(x, training=training)
 
-
         # Reduce block
         x = self.flatten(x, training=training)
         # x = self.drop_out(x, training=training)