From e9550bd4167212fb818486173d7d77b243a6968d Mon Sep 17 00:00:00 2001
From: GreenWizard <sizeof.2011@gmail.com>
Date: Sat, 2 Jan 2021 16:08:49 +0200
Subject: [PATCH] misc refactoring

---
 Agent/DQNEnsembleAgent.py | 32 ++++++++++----------------------
 Agent/MaskedSoftmax.py    | 14 ++++++++++++++
 view_maze.py              |  2 +-
 3 files changed, 25 insertions(+), 23 deletions(-)
 create mode 100644 Agent/MaskedSoftmax.py

diff --git a/Agent/DQNEnsembleAgent.py b/Agent/DQNEnsembleAgent.py
index b7fb0ae..0783341 100644
--- a/Agent/DQNEnsembleAgent.py
+++ b/Agent/DQNEnsembleAgent.py
@@ -3,35 +3,23 @@
 import tensorflow.keras as keras
 import tensorflow.keras.layers as layers
 import tensorflow as tf
+from Agent.MaskedSoftmax import MaskedSoftmax
 
 def combineModels(models, combiner):
   shape = models[0].layers[0].input_shape[0][1:]
   inputs = layers.Input(shape=shape)
   actionsMask = layers.Input(shape=(4, ))
-  res = layers.Lambda(combiner)([actionsMask] + [ x(inputs) for x in models ])
-  return keras.Model(inputs=[inputs, actionsMask], outputs=res)
-
-def maskedSoftmax(mask, inputs):
-  mask = tf.where(tf.equal(mask, 1))
-  return [
-    tf.sparse.to_dense(
-      tf.sparse.softmax(
-        tf.sparse.SparseTensor(
-          indices=mask,
-          values=tf.gather_nd(x, mask),
-          dense_shape=tf.shape(x, out_type=tf.int64)
-        )
-      )
-    ) for x in inputs
-  ]
   
-def multiplyOutputs(inputs):
-  outputs = maskedSoftmax(inputs[0], inputs[1:])
+  predictions = [ layers.Reshape((1, -1))(
+    MaskedSoftmax()( x(inputs), actionsMask )
+  ) for x in models ]
   
-  res = 1 + outputs[0]
-  for x in outputs[1:]:
-    res = tf.math.multiply(res, 1 + x)
-  return res
+  res = layers.Lambda(combiner)( layers.Concatenate(axis=1)(predictions)  )
+  return keras.Model(inputs=[inputs, actionsMask], outputs=res)
+
+@tf.function
+def multiplyOutputs(outputs):
+  return tf.math.reduce_prod(1 + outputs, axis=1)
 
 ENSEMBLE_MODE = {
   'multiply': multiplyOutputs
diff --git a/Agent/MaskedSoftmax.py b/Agent/MaskedSoftmax.py
new file mode 100644
index 0000000..1cc2858
--- /dev/null
+++ b/Agent/MaskedSoftmax.py
@@ -0,0 +1,14 @@
+import tensorflow as tf
+
+class MaskedSoftmax(tf.keras.layers.Layer):
+  def call(self, inputLayer, mask):
+    mask = tf.where(tf.equal(mask, 1))
+    return tf.sparse.to_dense(
+      tf.sparse.softmax(
+        tf.sparse.SparseTensor(
+          indices=mask,
+          values=tf.gather_nd(inputLayer, mask),
+          dense_shape=tf.shape(inputLayer, out_type=tf.int64)
+        )
+      )
+    )
diff --git a/view_maze.py b/view_maze.py
index f590cbc..777b216 100644
--- a/view_maze.py
+++ b/view_maze.py
@@ -2,7 +2,6 @@
 # -*- coding: utf-8 -*-
 import tensorflow as tf
 import os
-from Agent.DQNEnsembleAgent import DQNEnsembleAgent
 # limit GPU usage
 gpus = tf.config.experimental.list_physical_devices('GPU')
 tf.config.experimental.set_virtual_device_configuration(
@@ -15,6 +14,7 @@
 import pygame.locals as G
 import random
 from Agent.DQNAgent import DQNAgent
+from Agent.DQNEnsembleAgent import DQNEnsembleAgent
 import glob
 from collections import namedtuple
 from model import createModel