feat(embedding): use LSTMs for batch and time axes

 - refactor inputs to have a dense transform for the node values coming in
 - the time LSTM contextualizes the episode timesteps in the batch
 - the nodes LSTM contextualizes each node of the timesteps
 - the combination works well for multitask training, and exceeds previous models on acc of the poly task

Author: justindujardin

libraries/mathy_python/mathy/agents/embedding.py

@@ -19,27 +19,16 @@ def __init__(self, config: BaseConfig, **kwargs):
         self.token_embedding = tf.keras.layers.Embedding(
             input_dim=MathTypeKeysMax,
             output_dim=self.config.embedding_units,
-            name="nodes_embedding",
+            name="nodes_input",
             mask_zero=True,
         )
-
         # +1 for the value
         # +1 for the time
         # +2 for the problem type hashes
-        self.concat_size = (
-            4
-            if self.config.use_env_features
-            else 1
-            if self.config.use_node_values
-            else 0
-        )
-        if self.config.use_env_features:
-            self.time_dense = tf.keras.layers.Dense(
-                self.config.units, name="time_input"
-            )
-            self.type_dense = tf.keras.layers.Dense(
-                self.config.units, name="type_input"
-            )
+        self.concat_size = 4
+        self.values_dense = tf.keras.layers.Dense(self.config.units, name="values_input")
+        self.time_dense = tf.keras.layers.Dense(self.config.units, name="time_input")
+        self.type_dense = tf.keras.layers.Dense(self.config.units, name="type_input")
         self.in_dense = tf.keras.layers.Dense(
             self.config.units,
             # In transform gets the embeddings concatenated with the
@@ -61,28 +50,39 @@ def __init__(self, config: BaseConfig, **kwargs):
             NormalizeClass = tf.keras.layers.BatchNormalization
         self.out_dense_norm = NormalizeClass(name="out_dense_norm")
         self.nodes_lstm_norm = NormalizeClass(name="lstm_nodes_norm")
+        self.time_lstm_norm = NormalizeClass(name="time_lstm_norm")
         self.lstm_nodes = tf.keras.layers.LSTM(
             self.config.lstm_units,
             name="nodes_lstm",
             time_major=False,
             return_sequences=True,
             return_state=True,
         )
+        self.lstm_time = tf.keras.layers.LSTM(
+            self.config.lstm_units,
+            name="time_lstm",
+            time_major=True,
+            return_sequences=True,
+        )
         self.lstm_attention = tf.keras.layers.Attention()
 
     def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         output = tf.concat(
             [
                 self.token_embedding(features[ObservationFeatureIndices.nodes]),
-                tf.expand_dims(features[ObservationFeatureIndices.values], axis=-1),
+                self.values_dense(
+                    tf.expand_dims(features[ObservationFeatureIndices.values], axis=-1)
+                ),
                 self.type_dense(features[ObservationFeatureIndices.type]),
                 self.time_dense(features[ObservationFeatureIndices.time]),
             ],
             axis=-1,
             name="input_vectors",
         )
         output = self.in_dense(output)
+        output = self.lstm_time(output)
+        output = self.time_lstm_norm(output)
         output, state_h, state_c = self.lstm_nodes(output)
-        output = self.lstm_attention([output, state_h])
         output = self.nodes_lstm_norm(output)
+        output = self.lstm_attention([output, state_h])
         return self.out_dense_norm(self.output_dense(output))

libraries/mathy_python/mathy/agents/episode_memory.py

@@ -29,6 +29,13 @@ def clear(self):
         self.rewards = []
         self.values = []
 
+    def clear_except_window(self, window_size: int):
+        """Clear all data except for a window's worth of elements"""
+        self.observations = self.observations[-window_size:]
+        self.actions = self.actions[-window_size:]
+        self.rewards = self.rewards[-window_size:]
+        self.values = self.values[-window_size:]
+
     def to_window_observation(
         self, observation: MathyObservation, window_size: int = 3
     ) -> MathyWindowObservation:

libraries/mathy_python/mathy/envs/complex_simplify.py

@@ -26,7 +26,7 @@ def get_env_namespace(self) -> str:
     def max_moves_fn(
         self, problem: MathyEnvProblem, config: MathyEnvProblemArgs
     ) -> int:
-        return problem.complexity * 5
+        return problem.complexity * 8
 
     def problem_fn(self, params: MathyEnvProblemArgs) -> MathyEnvProblem:
         """Given a set of parameters to control term generation, produce

libraries/mathy_python/mathy/envs/poly_simplify.py

@@ -31,17 +31,7 @@ def get_env_namespace(self) -> str:
     def max_moves_fn(
         self, problem: MathyEnvProblem, config: MathyEnvProblemArgs
     ) -> int:
-        if config.difficulty == MathyEnvDifficulty.easy:
-            multiplier = 4
-        elif problem.complexity < 5:
-            multiplier = 2
-        elif problem.complexity < 7:
-            multiplier = 3
-        elif problem.complexity < 12:
-            multiplier = 4
-        else:
-            multiplier = 3
-        return problem.complexity * multiplier
+        return problem.complexity * 6
 
     def transition_fn(
         self,