refactor(model): remove episode long RNN state tracking

R2D2 and friends use this to great effect, but either I don't train long enough or I messed up the implementation. Remove it for now to simplify model usage, and let LSTM only operate on the current batch of observations (which are a rolling window of timesteps)

BREAKING CHANGE: this removes long-term RNN state tracking across episodes. Tracking the state was a significant amount of code and it wasn't clear that it made the model substantially better at any given task.

The overhead associated with keeping lots of hidden states in memory and calculating state histories was not insignificant on CPU training setups as well.

Author: justindujardin

libraries/mathy_python/mathy/agents/a3c/agent.py

@@ -37,7 +37,7 @@ def __init__(self, args: A3CConfig, env_extra: dict = None):
         self.action_size = env.action_space.n
         self.log_dir = os.path.join(self.args.model_dir, "tensorboard")
         self.writer = tf.summary.create_file_writer(self.log_dir)
-        init_window = env.initial_window(self.args.lstm_units)
+        init_window = env.initial_window()
         self.global_model = get_or_create_policy_model(
             args=args, predictions=self.action_size, is_main=True, env=env.mathy
         )

libraries/mathy_python/mathy/agents/a3c/worker.py

@@ -221,53 +221,16 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
         ep_steps = 1
         time_count = 0
         done = False
-        last_observation: MathyObservation = env.reset(rnn_size=self.args.lstm_units)
+        last_observation: MathyObservation = env.reset()
         last_text = env.state.agent.problem
         last_action = -1
         last_reward = -1
 
         selector = self.build_episode_selector(env)
 
-        # Set RNN to 0 state for start of episode
-        selector.model.embedding.reset_rnn_state()
-
-        # Start with the "init" sequence [n] times
-        for i in range(self.args.num_thinking_steps_begin):
-            rnn_state_h = tf.squeeze(selector.model.embedding.state_h.numpy())
-            rnn_state_c = tf.squeeze(selector.model.embedding.state_c.numpy())
-            seq_start = env.state.to_start_observation(rnn_state_h, rnn_state_c)
-            try:
-                window = observations_to_window([seq_start, last_observation])
-                selector.model.call(window.to_inputs())
-            except BaseException as err:
-                print_error(
-                    err, f"Episode begin thinking steps prediction failed.",
-                )
-                continue
-
         while not done and A3CWorker.request_quit is False:
             if self.args.print_training and self.worker_idx == 0:
                 env.render(self.args.print_mode, None)
-            # store rnn state for replay training
-            rnn_state_h = tf.squeeze(selector.model.embedding.state_h.numpy())
-            rnn_state_c = tf.squeeze(selector.model.embedding.state_c.numpy())
-            rnn_history_h = episode_memory.rnn_weighted_history(self.args.lstm_units)[0]
-            last_rnn_state = [rnn_state_h, rnn_state_c]
-
-            # named tuples are read-only, so add rnn state to a new copy
-            last_observation = MathyObservation(
-                nodes=last_observation.nodes,
-                mask=last_observation.mask,
-                values=last_observation.values,
-                type=last_observation.type,
-                time=last_observation.time,
-                rnn_state_h=tf.squeeze(rnn_state_h),
-                rnn_state_c=tf.squeeze(rnn_state_c),
-                rnn_history_h=rnn_history_h,
-            )
-            # before_rnn_state_h = selector.model.embedding.state_h.numpy()
-            # before_rnn_state_c = selector.model.embedding.state_c.numpy()
-
             window = episode_memory.to_window_observation(
                 last_observation, window_size=self.args.prediction_window_size
             )
@@ -277,33 +240,20 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
                     last_window=window,
                     last_action=last_action,
                     last_reward=last_reward,
-                    last_rnn_state=last_rnn_state,
                 )
             except BaseException as err:
                 print_error(err, "failed to select an action during an episode step")
                 continue
 
             # Take an env step
             observation, reward, done, _ = env.step(action)
-            rnn_state_h = tf.squeeze(selector.model.embedding.state_h.numpy())
-            rnn_state_c = tf.squeeze(selector.model.embedding.state_c.numpy())
-
-            # TODO: make this a unit test, check that EpisodeMemory states are not equal
-            #       across time steps.
-            # compare_states_h = tf.math.equal(before_rnn_state_h,rnn_state_h)
-            # compare_states_c = tf.math.equal(before_rnn_state_h,rnn_state_h)
-            # assert before_rnn_state_h != rnn_state_h
-            # assert before_rnn_state_c != rnn_state_c
 
             observation = MathyObservation(
                 nodes=observation.nodes,
                 mask=observation.mask,
                 values=observation.values,
                 type=observation.type,
                 time=observation.time,
-                rnn_state_h=rnn_state_h,
-                rnn_state_c=rnn_state_c,
-                rnn_history_h=rnn_history_h,
             )
 
             new_text = env.state.agent.problem
@@ -315,38 +265,12 @@ def run_episode(self, episode_memory: EpisodeMemory) -> float:
                 observation=last_observation,
                 action=action,
                 reward=reward,
-                grouping_change=grouping_change,
                 value=value,
             )
             if time_count == self.args.update_gradients_every or done:
                 if done and self.args.print_training and self.worker_idx == 0:
                     env.render(self.args.print_mode, None)
 
-                # TODO: Make this a unit test?
-                # Check that the LSTM h/c states changed over time in the episode.
-                #
-                # NOTE: in practice it seems every once in a while the state doesn't
-                # change, and I suppose this makes sense if the LSTM thought the
-                # existing state was... fine?
-                #
-                # check_rnn = None
-                # for obs in episode_memory.observations:
-                #     if check_rnn is not None:
-                #         h_equal_indices = (
-                #             tf.squeeze(tf.math.equal(obs.rnn_state_h, check_rnn[0]))
-                #             .numpy()
-                #             .tolist()
-                #         )
-                #         c_equal_indices = (
-                #             tf.squeeze(tf.math.equal(obs.rnn_state_c, check_rnn[1]))
-                #             .numpy()
-                #             .tolist()
-                #         )
-                #         assert False in h_equal_indices
-                #         assert False in c_equal_indices
-
-                #     check_rnn = [obs.rnn_state_h, obs.rnn_state_c]
-
                 self.update_global_network(done, observation, episode_memory)
                 self.maybe_write_histograms()
                 time_count = 0
@@ -406,17 +330,6 @@ def maybe_write_histograms(self) -> None:
                     tf.summary.histogram(
                         var.name, var, step=self.global_model.optimizer.iterations,
                     )
-                # Write out current LSTM hidden/cell states
-                tf.summary.histogram(
-                    "memory/lstm_c",
-                    self.local_model.embedding.state_c,
-                    step=self.global_model.optimizer.iterations,
-                )
-                tf.summary.histogram(
-                    "memory/lstm_h",
-                    self.local_model.embedding.state_h,
-                    step=self.global_model.optimizer.iterations,
-                )
 
     def update_global_network(
         self, done: bool, observation: MathyObservation, episode_memory: EpisodeMemory,

libraries/mathy_python/mathy/agents/action_selectors.py

@@ -26,7 +26,6 @@ def select(
         last_window: MathyWindowObservation,
         last_action: int,
         last_reward: float,
-        last_rnn_state: List[float],
     ) -> Tuple[int, float]:
         raise NotImplementedError(self.select)
 
@@ -46,7 +45,6 @@ def select(
         last_window: MathyWindowObservation,
         last_action: int,
         last_reward: float,
-        last_rnn_state: List[float],
     ) -> Tuple[int, float]:
 
         probs, value = self.model.predict_next(last_window.to_inputs())
@@ -85,7 +83,6 @@ def select(
         last_window: MathyWindowObservation,
         last_action: int,
         last_reward: float,
-        last_rnn_state: List[float],
     ) -> Tuple[int, float]:
-        probs, value = self.mcts.estimate_policy(last_state, last_rnn_state)
+        probs, value = self.mcts.estimate_policy(last_state)
         return np.argmax(probs), value

libraries/mathy_python/mathy/agents/embedding.py

@@ -13,18 +13,9 @@
 
 
 class MathyEmbedding(tf.keras.Model):
-    def __init__(
-        self,
-        config: BaseConfig,
-        episode_reset_state_h: Optional[bool] = True,
-        episode_reset_state_c: Optional[bool] = True,
-        **kwargs,
-    ):
+    def __init__(self, config: BaseConfig, **kwargs):
         super(MathyEmbedding, self).__init__(**kwargs)
         self.config = config
-        self.episode_reset_state_h = episode_reset_state_h
-        self.episode_reset_state_c = episode_reset_state_c
-        self.init_rnn_state()
         self.token_embedding = tf.keras.layers.Embedding(
             input_dim=MathTypeKeysMax,
             output_dim=self.config.embedding_units,
@@ -69,24 +60,15 @@ def __init__(
         if self.config.normalization_style == "batch":
             NormalizeClass = tf.keras.layers.BatchNormalization
         self.out_dense_norm = NormalizeClass(name="out_dense_norm")
-        if self.config.use_lstm:
-            self.nodes_lstm_norm = NormalizeClass(name="lstm_nodes_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_attention = tf.keras.layers.Attention()
-        else:
-            self.densenet = DenseNetStack(
-                units=self.config.units,
-                num_layers=6,
-                output_transform=self.output_dense,
-                normalization_style=self.config.normalization_style,
-            )
-            self.dense_attention = tf.keras.layers.Attention()
+        self.nodes_lstm_norm = NormalizeClass(name="lstm_nodes_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_attention = tf.keras.layers.Attention()
 
     def compute_output_shape(self, input_shapes: List[tf.TensorShape]) -> Any:
         nodes_shape: tf.TensorShape = input_shapes[0]
@@ -98,27 +80,6 @@ def compute_output_shape(self, input_shapes: List[tf.TensorShape]) -> Any:
             )
         )
 
-    def init_rnn_state(self):
-        """Track RNN states with variables in the graph"""
-        self.state_c = tf.Variable(
-            tf.zeros([1, self.config.lstm_units]),
-            trainable=False,
-            name="embedding/rnn/agent_state_c",
-        )
-        self.state_h = tf.Variable(
-            tf.zeros([1, self.config.lstm_units]),
-            trainable=False,
-            name="embedding/rnn/agent_state_h",
-        )
-
-    def reset_rnn_state(self, force: bool = False):
-        """Zero out the RNN state for a new episode"""
-        if self.episode_reset_state_h or force is True:
-            self.state_h.assign(tf.zeros([1, self.config.lstm_units]))
-
-        if self.episode_reset_state_c or force is True:
-            self.state_c.assign(tf.zeros([1, self.config.lstm_units]))
-
     def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         nodes = features[ObservationFeatureIndices.nodes]
         values = features[ObservationFeatureIndices.values]
@@ -128,10 +89,6 @@ def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         batch_size = nodes_shape[0]  # noqa
         sequence_length = nodes_shape[1]
 
-        in_rnn_state_h = features[ObservationFeatureIndices.rnn_state_h]
-        in_rnn_state_c = features[ObservationFeatureIndices.rnn_state_c]
-        in_rnn_history_h = features[ObservationFeatureIndices.rnn_history_h]
-
         with tf.name_scope("prepare_inputs"):
             values = tf.expand_dims(values, axis=-1, name="values_input")
             query = self.token_embedding(nodes)
@@ -159,8 +116,6 @@ def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
             output, state_h, state_c = self.lstm_nodes(query)
             output = self.lstm_attention([output, state_h])
             output = self.nodes_lstm_norm(output)
-            self.state_h.assign(state_h[-1:])
-            self.state_c.assign(state_c[-1:])
         else:
             # Non-recurrent model
             output = self.densenet(query)

libraries/mathy_python/mathy/agents/episode_memory.py

@@ -7,34 +7,9 @@
     MathyObservation,
     MathyWindowObservation,
     observations_to_window,
-    rnn_placeholder_state,
 )
 
 
-def rnn_weighted_history(
-    observations: List[MathyObservation], rnn_size: int = 128,
-):
-    """Build a historical LSTM state: https://arxiv.org/pdf/1810.04437.pdf
-
-    Take the mean of the previous LSTM states for this episode. """
-
-    import tensorflow as tf
-
-    if len(observations) > 0:
-        in_h = []
-        in_c = []
-        for obs in observations:
-            in_h.append(obs.rnn_state_h)
-            in_c.append(obs.rnn_state_c)
-        # Take the mean of the historical states:
-        memory_context_h = tf.reduce_mean(in_h, axis=0)
-        memory_context_c = tf.reduce_mean(in_c, axis=0)
-    else:
-        memory_context_h = rnn_placeholder_state(rnn_size)
-        memory_context_c = memory_context_h
-    return [memory_context_h, memory_context_c]
-
-
 class EpisodeMemory(object):
     # Observation from the environment
     observations: List[MathyObservation]
@@ -44,8 +19,6 @@ class EpisodeMemory(object):
     rewards: List[float]
     # Estimated value from the model
     values: List[float]
-    # Grouping Control error from the environment
-    grouping_changes: List[float]
 
     def __init__(self):
         self.clear()
@@ -55,7 +28,6 @@ def clear(self):
         self.actions = []
         self.rewards = []
         self.values = []
-        self.grouping_changes = []
 
     def to_window_observation(
         self, observation: MathyObservation, window_size: int = 3
@@ -74,14 +46,9 @@ def store(
         observation: MathyObservation,
         action: int,
         reward: float,
-        grouping_change: float,
         value: float,
     ):
         self.observations.append(observation)
         self.actions.append(action)
         self.rewards.append(reward)
         self.values.append(value)
-        self.grouping_changes.append(grouping_change)
-
-    def rnn_weighted_history(self, rnn_size):
-        return rnn_weighted_history(self.observations, rnn_size)

libraries/mathy_python/mathy/agents/example.py

@@ -10,5 +10,5 @@ def example() -> MathyWindowObservation:
     passed forward through a Mathy model. """
     env = PolySimplify()
     state = env.get_initial_state()[0]
-    observation = env.state_to_observation(state, rnn_size=BaseConfig().lstm_units)
+    observation = env.state_to_observation(state)
     return observations_to_window([observation])