feat(tensorflow): update to 2.1.0

Author: justindujardin

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

@@ -42,15 +42,16 @@ def __init__(self, args: A3CConfig, env_extra: dict = None):
             args=args, env_actions=self.action_size, is_main=True, env=env.mathy
         )
         with self.writer.as_default():
+            tf_model: tf.keras.Model = self.global_model.unwrapped
             tf.summary.trace_on(graph=True)
             inputs = init_window.to_inputs()
-            self.global_model.call_graph(inputs)
+            tf_model.call_graph(inputs)
             tf.summary.trace_export(
                 name="PolicyValueModel", step=0, profiler_outdir=self.log_dir
             )
             tf.summary.trace_off()
             if self.args.verbose:
-                print(self.global_model.summary())
+                print(tf_model.summary())
 
     def train(self):
         res_queue = Queue()
@@ -68,7 +69,7 @@ def train(self):
                 args=self.args,
                 teacher=self.teacher,
                 worker_idx=i,
-                optimizer=self.global_model.optimizer,
+                optimizer=self.global_model.unwrapped.optimizer,
                 result_queue=res_queue,
                 writer=self.writer,
             )
@@ -93,3 +94,4 @@ def train(self):
         # Do a final save after joining to get the very latest model
         self.global_model.save()
         print("Done. Bye!")
+

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

@@ -26,7 +26,7 @@
 from .. import action_selectors
 from ..episode_memory import EpisodeMemory
 from ..mcts import MCTS
-from ..policy_value_model import PolicyValueModel, get_or_create_policy_model
+from ..policy_value_model import ThincPolicyValueModel, get_or_create_policy_model
 from ..trfl import discrete_policy_entropy_loss, td_lambda
 from .config import A3CConfig
 from .util import record, truncate
@@ -50,7 +50,7 @@ def __init__(
         self,
         args: A3CConfig,
         action_size: int,
-        global_model: PolicyValueModel,
+        global_model: ThincPolicyValueModel,
         optimizer,
         greedy_epsilon: Union[float, List[float]],
         result_queue: Queue,
@@ -132,7 +132,7 @@ def run(self):
                 if win_pct is not None:
                     with self.writer.as_default():
                         student = self.teacher.students[self.worker_idx]
-                        step = self.global_model.optimizer.iterations
+                        step = self.global_model.unwrapped.optimizer.iterations
                         if self.worker_idx == 0:
                             tf.summary.scalar(
                                 f"win_rate/{student.topic}", data=win_pct, step=step
@@ -216,7 +216,7 @@ def build_episode_selector(
             )
         return selector
 
-    def run_episode(self, episode_memory: EpisodeMemory):
+    def run_episode(self, episode_memory: EpisodeMemory) -> float:
         env_name = self.teacher.get_env(self.worker_idx, self.iteration)
         env = gym.make(env_name, **self.env_extra)
         episode_memory.clear()
@@ -233,16 +233,16 @@ def run_episode(self, episode_memory: EpisodeMemory):
         selector = self.build_episode_selector(env)
 
         # Set RNN to 0 state for start of episode
-        selector.model.embedding.reset_rnn_state()
+        selector.model.unwrapped.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())
+            rnn_state_h = tf.squeeze(selector.model.unwrapped.embedding.state_h.numpy())
+            rnn_state_c = tf.squeeze(selector.model.unwrapped.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())
+                selector.model([window.to_inputs()], is_train=True)
             except BaseException as err:
                 print_error(
                     err, f"Episode begin thinking steps prediction failed.",
@@ -253,8 +253,8 @@ def run_episode(self, episode_memory: EpisodeMemory):
             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_state_h = tf.squeeze(selector.model.unwrapped.embedding.state_h.numpy())
+            rnn_state_c = tf.squeeze(selector.model.unwrapped.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]
 
@@ -269,8 +269,8 @@ def run_episode(self, episode_memory: EpisodeMemory):
                 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()
+            # before_rnn_state_h = selector.model.unwrapped.embedding.state_h.numpy()
+            # before_rnn_state_c = selector.model.unwrapped.embedding.state_c.numpy()
 
             window = episode_memory.to_window_observation(last_observation)
             try:
@@ -287,8 +287,8 @@ def run_episode(self, episode_memory: EpisodeMemory):
 
             # 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())
+            rnn_state_h = tf.squeeze(selector.model.unwrapped.embedding.state_h.numpy())
+            rnn_state_c = tf.squeeze(selector.model.unwrapped.embedding.state_c.numpy())
 
             # TODO: make this a unit test, check that EpisodeMemory states are not equal
             #       across time steps.
@@ -379,7 +379,7 @@ def maybe_write_episode_summaries(
         assert self.worker_idx == 0, "only write summaries for greedy worker"
         # Track metrics for all workers
         name = self.teacher.get_env(self.worker_idx, self.iteration)
-        step = self.global_model.optimizer.iterations
+        step = self.global_model.unwrapped.optimizer.iterations
         with self.writer.as_default():
             agent_state = last_state.agent
             steps = int(last_state.max_moves - agent_state.moves_remaining)
@@ -396,28 +396,30 @@ def maybe_write_episode_summaries(
                 step=step,
             )
 
-    def maybe_write_histograms(self):
+    def maybe_write_histograms(self) -> None:
         if self.worker_idx != 0:
             return
-        step = self.global_model.optimizer.iterations.numpy()
+        step = self.global_model.unwrapped.optimizer.iterations.numpy()
         next_write = self.last_histogram_write + self.args.summary_interval
         if step >= next_write or self.last_histogram_write == -1:
             with self.writer.as_default():
                 self.last_histogram_write = step
-                for var in self.local_model.trainable_variables:
+                for var in self.local_model.unwrapped.trainable_variables:
                     tf.summary.histogram(
-                        var.name, var, step=self.global_model.optimizer.iterations
+                        var.name,
+                        var,
+                        step=self.global_model.unwrapped.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,
+                    self.local_model.unwrapped.embedding.state_c,
+                    step=self.global_model.unwrapped.optimizer.iterations,
                 )
                 tf.summary.histogram(
                     "memory/lstm_h",
-                    self.local_model.embedding.state_h,
-                    step=self.global_model.optimizer.iterations,
+                    self.local_model.unwrapped.embedding.state_h,
+                    step=self.global_model.unwrapped.optimizer.iterations,
                 )
 
     def update_global_network(
@@ -442,10 +444,10 @@ def update_global_network(
                 self.ep_aux_loss[k] = 0.0
             self.ep_aux_loss[k] += aux_losses[k].numpy()
         # Calculate local gradients
-        grads = tape.gradient(total_loss, self.local_model.trainable_weights)
+        grads = tape.gradient(total_loss, self.local_model.unwrapped.trainable_weights)
         # Push local gradients to global model
 
-        zipped_gradients = zip(grads, self.global_model.trainable_weights)
+        zipped_gradients = zip(grads, self.global_model.unwrapped.trainable_weights)
         # Assert that we always have some gradient flow in each trainable var
 
         # TODO: Make this a unit test. It degrades performance at train time
@@ -460,7 +462,7 @@ def update_global_network(
 
         self.optimizer.apply_gradients(zipped_gradients)
         # Update local model with new weights
-        self.local_model.set_weights(self.global_model.get_weights())
+        self.local_model.unwrapped.set_weights(self.global_model.unwrapped.get_weights())
         episode_memory.clear()
 
     def finish_episode(self, episode_reward, episode_steps, last_state: MathyEnvState):
@@ -487,7 +489,7 @@ def finish_episode(self, episode_reward, episode_steps, last_state: MathyEnvStat
                 episode_reward, episode_steps, last_state
             )
 
-            step = self.global_model.optimizer.iterations.numpy()
+            step = self.global_model.unwrapped.optimizer.iterations.numpy()
             next_write = self.last_model_write + A3CWorker.save_every_n_episodes
             if step >= next_write or self.last_model_write == -1:
                 self.last_model_write = step
@@ -512,12 +514,12 @@ def compute_policy_value_loss(
         episode_memory: EpisodeMemory,
         gamma=0.99,
     ):
-        step = self.global_model.optimizer.iterations
+        step = self.global_model.unwrapped.optimizer.iterations
         if done:
             bootstrap_value = 0.0  # terminal
         else:
             # Predict the reward using the local network
-            _, values, _ = self.local_model.call(
+            _, values, _ = self.local_model.unwrapped.call(
                 observations_to_window([observation]).to_inputs()
             )
             # Select the last timestep
@@ -536,7 +538,8 @@ def compute_policy_value_loss(
         batch_size = len(episode_memory.actions)
         sequence_length = len(episode_memory.observations[0].nodes)
         inputs = episode_memory.to_episode_window().to_inputs()
-        logits, values, trimmed_logits = self.local_model(inputs, apply_mask=False)
+        logits, values, trimmed_logits = self.local_model.unwrapped(inputs, apply_mask=False)
+        # TODO: don't call unwrapped here
 
         logits = tf.reshape(logits, [batch_size, -1])
         policy_logits = tf.reshape(trimmed_logits, [batch_size, -1])
@@ -615,7 +618,7 @@ def compute_loss(
         gamma=0.99,
     ):
         with self.writer.as_default():
-            step = self.global_model.optimizer.iterations
+            step = self.global_model.unwrapped.optimizer.iterations
             loss_tuple = self.compute_policy_value_loss(
                 done, observation, episode_memory
             )

libraries/mathy_python/mathy/agents/action_selectors.py

@@ -3,14 +3,15 @@
 import numpy as np
 
 from ..state import MathyEnvState, MathyWindowObservation
-from .policy_value_model import PolicyValueModel
+from .policy_value_model import ThincPolicyValueModel
 from .mcts import MCTS
+import thinc
 
 
 class ActionSelector:
     """An episode-specific selector of actions"""
 
-    def __init__(self, *, model: PolicyValueModel, episode: int, worker_id: int):
+    def __init__(self, *, model: ThincPolicyValueModel, episode: int, worker_id: int):
         self.model = model
         self.worker_id = worker_id
         self.episode = episode

libraries/mathy_python/mathy/agents/mcts.py

@@ -4,6 +4,7 @@
 import numpy
 
 from ..env import MathyEnv
+from .policy_value_model import ThincPolicyValueModel
 from ..state import MathyEnvState, observations_to_window
 from ..util import is_terminal_transition
 
@@ -28,7 +29,7 @@ class MCTS:
     def __init__(
         self,
         env: MathyEnv,
-        model: Any,
+        model: ThincPolicyValueModel,
         cpuct: float = 1.0,
         num_mcts_sims: int = 15,
         epsilon: float = 0.25,
@@ -130,8 +131,8 @@ def search(self, env_state: MathyEnvState, rnn_state: List[Any], isRootNode=Fals
             observations = observations_to_window([obs]).to_inputs()
             out_policy, state_v = self.model.predict_next(observations, use_graph=False)
             out_rnn_state = [
-                tf.squeeze(self.model.embedding.state_h).numpy(),
-                tf.squeeze(self.model.embedding.state_c).numpy(),
+                tf.squeeze(self.model.unwrapped.embedding.state_h).numpy(),
+                tf.squeeze(self.model.unwrapped.embedding.state_c).numpy(),
             ]
             self.Rs[s] = out_rnn_state
             self.Ps[s] = out_policy