[AI] Fix tensorboard

Author: Herklos

ai-example.py

@@ -73,9 +73,12 @@ def init_argparse() -> argparse.ArgumentParser:
     parser.add_argument("-p", "--plot", action=argparse.BooleanOptionalAction)
     parser.add_argument('-w', '--weights', type=str, help='a trained model weights')
     parser.add_argument("-d", "--days", type=int, default=365)
+    parser.add_argument("-ev", "--evaluate", action=argparse.BooleanOptionalAction)
+    parser.add_argument("-ep", "--epochs", type=int, default=20)
     return parser
 
 
+
 def main():
     parser = init_argparse()
     args = parser.parse_args()
@@ -90,7 +93,7 @@ def main():
     agent = obs.DQNAgent(action_size)
 
     logdir = "tensorboard_logs/scalars/" + datetime.now().strftime("%Y%m%d-%H%M%S")
-    tensorboard_callback = TensorBoard(log_dir=logdir)
+    tensorboard_callback = TensorBoard(log_dir=logdir, histogram_freq=1, write_images=False, batch_size=args.batch_size)
 
     if args.weights:
         print(f"Loading model {args.weights}...")
@@ -107,7 +110,9 @@ def main():
 
         if args.train and len(agent.memory) > args.batch_size:
             print("Starting replay...")
-            agent.replay(args.batch_size, tensorboard_callback)
+            score = agent.replay(args.batch_size, args.epochs, args.evaluate, tensorboard_callback)
+            if args.evaluate:
+                print(f"Score = {score}")
 
         if args.train and (episode + 1) % 10 == 0:  # checkpoint weights
             print("Saving...")

octobot_script/ai/agents.py

@@ -23,7 +23,7 @@ def act(self, state):
         act_values = self.model.predict(state)
         return np.argmax(act_values[0])  # returns action
 
-    def replay(self, batch_size=32, tensorboard_callback=None):
+    def replay(self, batch_size=32, epochs=1, evaluate=False, tensorboard_callback=None):
         # pylint: disable=unsubscriptable-object
         """ vectorized implementation; 30x speed up compared with for loop """
         minibatch = random.sample(self.memory, batch_size)
@@ -44,10 +44,14 @@ def replay(self, batch_size=32, tensorboard_callback=None):
         # make the agent to approximately map the current state to future discounted reward
         target_f[range(batch_size), actions] = target
 
-        self.model.fit(states, target_f, epochs=1, verbose=0)
+        self.model.fit(states, target_f, batch_size=batch_size, epochs=epochs, verbose=0, callbacks=[tensorboard_callback])
 
         if self.epsilon > self.epsilon_min:
             self.epsilon *= self.epsilon_decay
+        
+        if evaluate:
+            return self.model.evaluate(states, target_f, batch_size=32)
+        return 0
 
     def load(self, name):
         self.model.load_weights(name)