fix(policy_value_model): value head was not learning from hidden state

 - Maybe it was a dumb idea to make the value head a single-layer projection of the hidden state.
 - Reduce the output sequence and use a dense layer activated with ReLu.

Author: justindujardin

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -63,7 +63,12 @@ def __init__(
         self.args = args
         self.predictions = predictions
         self.embedding = MathyEmbedding(self.args)
-        # self.embedding.compile(optimizer=self.optimizer, run_eagerly=True)
+        self.value_head = tf.keras.layers.Dense(
+            self.args.units,
+            name="value_head",
+            kernel_initializer="he_normal",
+            activation="relu",
+        )
         self.value_logits = tf.keras.layers.Dense(
             1, name="value_logits", kernel_initializer="he_normal", activation=None,
         )
@@ -114,7 +119,9 @@ def _call(
         inputs = features_window
         # Extract features into contextual inputs, sequence inputs.
         sequence_inputs = self.embedding(inputs)
-        values = self.normalize_v(self.value_logits(self.embedding.state_h))
+        values = self.value_head(tf.reduce_mean(sequence_inputs, axis=1))
+        values = self.normalize_v(values)
+        values = self.value_logits(values)
         logits = self.normalize_pi(self.policy_logits(sequence_inputs))
         mask_logits = self.apply_pi_mask(logits, features_window)
         if call_print is True: