feat(MathyWindowObservation): add option to return inputs using numpy instead of tf.Tensor

Author: justindujardin

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

@@ -39,13 +39,18 @@ def __init__(self, args: A3CConfig, env_extra: dict = None):
         self.writer = tf.summary.create_file_writer(self.log_dir)
         init_window = env.initial_window(self.args.lstm_units)
         self.global_model = get_or_create_policy_model(
-            args=args, env_actions=self.action_size, is_main=True, env=env.mathy
+            args=args, predictions=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()
-            tf_model.call_graph(inputs)
+
+            @tf.function
+            def trace_fn():
+                return tf_model.call(inputs)
+
+            trace_fn()
             tf.summary.trace_export(
                 name="PolicyValueModel", step=0, profiler_outdir=self.log_dir
             )

libraries/mathy_python/mathy/agents/embedding.py

@@ -135,16 +135,11 @@ def reset_rnn_state(self, force: bool = False):
         if self.episode_reset_state_c or force is True:
             self.state_c.assign(tf.zeros([1, self.config.lstm_units]))
 
-    @tf.function
-    def call_graph(self, inputs: MathyInputsType) -> tf.Tensor:
-        """Autograph optimized function"""
-        return self.call(inputs)
-
-    def call(self, features: MathyInputsType) -> tf.Tensor:
-        nodes = tf.convert_to_tensor(features[ObservationFeatureIndices.nodes])
-        values = tf.convert_to_tensor(features[ObservationFeatureIndices.values])
-        type = tf.cast(features[ObservationFeatureIndices.type], dtype=tf.float32)
-        time = tf.cast(features[ObservationFeatureIndices.time], dtype=tf.float32)
+    def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
+        nodes = features[ObservationFeatureIndices.nodes]
+        values = features[ObservationFeatureIndices.values]
+        type = features[ObservationFeatureIndices.type]
+        time = features[ObservationFeatureIndices.time]
         nodes_shape = tf.shape(features[ObservationFeatureIndices.nodes])
         batch_size = nodes_shape[0]  # noqa
         sequence_length = nodes_shape[1]

libraries/mathy_python/mathy/agents/mcts.py

@@ -129,7 +129,7 @@ def search(self, env_state: MathyEnvState, rnn_state: List[Any], isRootNode=Fals
                 rnn_history_h=rnn_state[0],
             )
             observations = observations_to_window([obs]).to_inputs()
-            out_policy, state_v = self.model.predict_next(observations, use_graph=False)
+            out_policy, state_v = self.model.predict_next(observations)
             out_rnn_state = [
                 tf.squeeze(self.model.unwrapped.embedding.state_h).numpy(),
                 tf.squeeze(self.model.unwrapped.embedding.state_c).numpy(),

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -7,7 +7,7 @@
 import numpy as np
 import srsly
 import tensorflow as tf
-from tensorflow.python.keras import backend as K
+from tensorflow.keras import backend as K
 from wasabi import msg
 
 from ..env import MathyEnv
@@ -124,13 +124,6 @@ def apply_pi_mask(
         )
         return negative_mask_logits
 
-    @tf.function
-    def call_graph(
-        self, inputs: MathyInputsType
-    ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
-        """Autograph optimized function"""
-        return self.call(inputs)
-
 
 class ThincPolicyValueModel(thinc.model.Model[ArrayNd, Tuple[Array1d, Array2d]]):
     @property
@@ -139,15 +132,10 @@ def unwrapped(self) -> TFPVModel:
         assert isinstance(tf_shim, TensorFlowShim), "only tensorflow shim is supported"
         return tf_shim._model
 
-    def predict_next(
-        self, inputs: MathyInputsType, use_graph: bool = False
-    ) -> Tuple[tf.Tensor, tf.Tensor]:
+    def predict_next(self, inputs: MathyInputsType) -> Tuple[tf.Tensor, tf.Tensor]:
         """Predict one probability distribution and value for the
         given sequence of inputs """
-        if use_graph:
-            logits, values, masked = self.unwrapped.call_graph(inputs)
-        else:
-            logits, values, masked = self.unwrapped.call(inputs)
+        logits, values, masked = self.unwrapped.call(inputs)
         # take the last timestep
         masked = masked[-1][:]
         flat_logits = tf.reshape(tf.squeeze(masked), [-1])
@@ -198,7 +186,7 @@ def _load_model(
 
 def get_or_create_policy_model(
     args: BaseConfig,
-    env_actions: int,
+    predictions: int,
     is_main=False,
     required=False,
     env: MathyEnv = None,
@@ -235,7 +223,7 @@ def get_or_create_policy_model(
                 print_error=False,
             )
 
-    model = PolicyValueModel(args=args, predictions=env_actions, name="agent")
+    model = PolicyValueModel(args=args, predictions=predictions, name="agent")
     init_inputs = initial_state.to_inputs()
 
     def handshake_keras(m: ThincPolicyValueModel):

libraries/mathy_python/mathy/agents/zero/self_play_runner.py

@@ -39,7 +39,7 @@ def get_model(self, game):
             )
 
             model: PolicyValueModel = get_or_create_policy_model(
-                args=config, env_actions=game.action_space.n, is_main=True
+                args=config, predictions=game.action_space.n, is_main=True
             )
             return model
 

libraries/mathy_python/mathy/state.py

@@ -77,18 +77,23 @@ class MathyWindowObservation(NamedTuple):
     rnn_state_c: WindowRNNStateFloatList
     rnn_history_h: WindowRNNStateFloatList
 
-    def to_inputs(self) -> MathyInputsType:
+    def to_inputs(self, as_tf_tensor: bool = True) -> MathyInputsType:
         import tensorflow as tf
 
+        def to_res(in_value):
+            if as_tf_tensor is True:
+                return tf.convert_to_tensor(in_value, dtype=tf.float32)
+            return np.asarray(in_value, dtype="float32")
+
         result = [
-            tf.convert_to_tensor(self.nodes),
-            tf.convert_to_tensor(self.mask),
-            tf.convert_to_tensor(self.values),
-            tf.convert_to_tensor(self.type),
-            tf.convert_to_tensor(self.time),
-            tf.convert_to_tensor(self.rnn_state_h),
-            tf.convert_to_tensor(self.rnn_state_c),
-            tf.convert_to_tensor(self.rnn_history_h),
+            to_res(self.nodes),
+            to_res(self.mask),
+            to_res(self.values),
+            to_res(self.type),
+            to_res(self.time),
+            to_res(self.rnn_state_h),
+            to_res(self.rnn_state_c),
+            to_res(self.rnn_history_h),
         ]
         for r in result:
             for s in r.shape:

libraries/website/docs/snippets/ml/policy_value_serialization.ipynb

@@ -38,7 +38,7 @@
     "instance.train()\n",
     "# Load the model back in\n",
     "model_two = get_or_create_policy_model(\n",
-    "    args=args, env_actions=PolySimplify().action_size, is_main=True\n",
+    "    args=args, predictions=PolySimplify().action_size, is_main=True\n",
     ")\n",
     "# Comment this out to keep your model\n",
     "shutil.rmtree(model_folder)"

libraries/website/docs/snippets/ml/policy_value_serialization.py

@@ -26,7 +26,7 @@
 instance.train()
 # Load the model back in
 model_two = get_or_create_policy_model(
-    args=args, env_actions=PolySimplify().action_size, is_main=True
+    args=args, predictions=PolySimplify().action_size, is_main=True
 )
 # Comment this out to keep your model
 shutil.rmtree(model_folder)

libraries/website/docs/snippets/ml/thinc_basic.ipynb

@@ -0,0 +1,86 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "# This file is generated from a Mathy (https://mathy.ai) code example.\n",
+    "!pip install mathy --upgrade\n",
+    "import os\n",
+    "import shutil\n",
+    "import tempfile\n",
+    "from typing import List, Tuple\n",
+    "\n",
+    "import numpy as np\n",
+    "\n",
+    "from mathy import envs\n",
+    "from mathy.agents.base_config import BaseConfig\n",
+    "from mathy.agents.embedding import MathyEmbedding\n",
+    "from mathy.env import MathyEnv\n",
+    "from mathy.state import MathyEnvState, MathyObservation, observations_to_window\n",
+    "from thinc.api import TensorFlowWrapper\n",
+    "from thinc.layers import (\n",
+    "    Embed,\n",
+    "    Linear,\n",
+    "    MeanPool,\n",
+    "    ReLu,\n",
+    "    Softmax,\n",
+    "    chain,\n",
+    "    list2ragged,\n",
+    "    with_array,\n",
+    "    with_list,\n",
+    ")\n",
+    "from thinc.model import Model\n",
+    "from thinc.shims.tensorflow import TensorFlowShim\n",
+    "from thinc.types import Array, Array1d, Array2d, ArrayNd\n",
+    "\n",
+    "# Mathy env setup and initial observations\n",
+    "args = BaseConfig()\n",
+    "env: MathyEnv = envs.PolySimplify()\n",
+    "state: MathyEnvState = env.get_initial_state()[0]\n",
+    "observation: MathyObservation = env.state_to_observation(\n",
+    "    state, rnn_size=args.lstm_units\n",
+    ")\n",
+    "window = observations_to_window([observation, observation])\n",
+    "inputs = window.to_inputs()\n",
+    "\n",
+    "X = [inputs]  # TODO: why do I need to wrap inputs for the tf wrapper?\n",
+    "input_shape = window.to_input_shapes()\n",
+    "embeddings = TensorFlowWrapper(MathyEmbedding(args), input_shape=input_shape)\n",
+    "embeddings.initialize([inputs])\n",
+    "\n",
+    "embed_Y = embeddings.predict([inputs])\n",
+    "# Shape = (2, 23, 128) = (num_observations, padded_sequence_len, vector_width)\n",
+    "\n",
+    "# The policy head is a softmax(actions) for each node in the sequence.\n",
+    "policy_head = chain(embeddings, Softmax(6))\n",
+    "policy_head.initialize([inputs])\n",
+    "policy_Y = policy_head.predict([inputs])\n",
+    "# Shape (desired) = (2, 23, 6)\n",
+    "\n",
+    "# The value head is normally a linear transformation from the\n",
+    "# output embedding layer's RNN state. I haven't tried mixing\n",
+    "# that tensor in here. TODO: try that\n",
+    "value_head = chain(embeddings, MeanPool(), Linear(1))\n",
+    "value_head.initialize([inputs])\n",
+    "value_Y = value_head.predict([inputs])\n",
+    "# Shape (desired) = (2, 1)\n",
+    "\n",
+    "# Combined [policy_head, value_head] outputs without invoking embeddings twice?\n",
+    "model: Model[ArrayNd, Tuple[Array2d, Array1d]] = ...\n",
+    "\n",
+    "model.initialize([inputs])\n",
+    "\n",
+    "Y = model.predict([inputs])\n",
+    "model.to_disk(f\"training/model\")"
+   ]
+  }
+ ],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

libraries/website/docs/snippets/ml/thinc_basic.py

@@ -0,0 +1,67 @@
+import os
+import shutil
+import tempfile
+from typing import List, Tuple
+
+import numpy as np
+
+from mathy import envs
+from mathy.agents.base_config import BaseConfig
+from mathy.agents.embedding import MathyEmbedding
+from mathy.env import MathyEnv
+from mathy.state import MathyEnvState, MathyObservation, observations_to_window
+from thinc.api import TensorFlowWrapper
+from thinc.layers import (
+    Embed,
+    Linear,
+    MeanPool,
+    ReLu,
+    Softmax,
+    chain,
+    list2ragged,
+    with_array,
+    with_list,
+)
+from thinc.model import Model
+from thinc.shims.tensorflow import TensorFlowShim
+from thinc.types import Array, Array1d, Array2d, ArrayNd
+
+# Mathy env setup and initial observations
+args = BaseConfig()
+env: MathyEnv = envs.PolySimplify()
+state: MathyEnvState = env.get_initial_state()[0]
+observation: MathyObservation = env.state_to_observation(
+    state, rnn_size=args.lstm_units
+)
+window = observations_to_window([observation, observation])
+inputs = window.to_inputs()
+
+X = [inputs]  # TODO: why do I need to wrap inputs for the tf wrapper?
+input_shape = window.to_input_shapes()
+embeddings = TensorFlowWrapper(MathyEmbedding(args), input_shape=input_shape)
+embeddings.initialize([inputs])
+
+embed_Y = embeddings.predict([inputs])
+# Shape = (2, 23, 128) = (num_observations, padded_sequence_len, vector_width)
+
+# The policy head is a softmax(actions) for each node in the sequence.
+policy_head = chain(embeddings, Softmax(6))
+policy_head.initialize([inputs])
+policy_Y = policy_head.predict([inputs])
+# Shape (desired) = (2, 23, 6)
+
+# The value head is normally a linear transformation from the
+# output embedding layer's RNN state. I haven't tried mixing
+# that tensor in here. TODO: try that
+value_head = chain(embeddings, MeanPool(), Linear(1))
+value_head.initialize([inputs])
+value_Y = value_head.predict([inputs])
+# Shape (desired) = (2, 1)
+
+# Combined [policy_head, value_head] outputs without invoking embeddings twice?
+model: Model[ArrayNd, Tuple[Array2d, Array1d]] = ...
+
+model.initialize([inputs])
+
+Y = model.predict([inputs])
+model.to_disk(f"training/model")