Export Tensorflow SavedModel?

[AidanNelson]

Hi There,

I am exploring the possibility of training a model in Unity using the ML-Agents package, then exporting that model for use with Tensorflow.js in the browser. Eventually, I'd like to incorporate that model into ml5.js, an approachable wrapper library around tf.js. I think this workflow would be an exciting use of the training / curriculum learning structure of Unity, and allow people to create models which could be used in online demonstrations of reinforcement learning (...or maybe agent-based AI for online games?)

Currently, it is possible to convert Tensorflow models for use with Tensorflow.js as long as they are in the SavedModel, Keras model, or TensorFlow Hub module format.

Describe the solution you'd like
I would like to be able to export a trained model from Unity in one of the above Tensorflow formats (SavedModel, Keras model, or TensorFlow Hub module), for later conversion to Tensorflow.js using existing solutions.

I see that in the current release of ML-Agents, Barracuda is used to store models, and that a Tensorflow to Barracuda conversion script exists, but am unsure how possible it would be to convert a model from Barracuda back to Tensorflow and the extent of this work.

Any information about this would be most appreciated!

Aidan

A (vaguely) related issue:
#1802

UPDATE (2019.12.2)
I looked further and realized that Unity Ml-Agents still uses Tensorflow for training, and only converts the models to Barracuda '.nn' files on export for inference. Now I am trying to figure out how to export a SavedModel format in this export_model function, for later conversion to TFJS-capable model. Does anyone have experience with using the saved model builder and a sense for how it could be implemented here? I've altered the export_model function to use tf.compat.v1.saved_model.Builder in the following code:

def export_model(self):
        """
        Exports latest saved model to .nn format for Unity embedding.
        """
        with self.graph.as_default():
  
            # MY ADDITIONAL CODE:
            builder = tf.compat.v1.saved_model.Builder(self.model_path + "/SavedModel/")
            builder.add_meta_graph_and_variables(self.sess,
                                                [tf.saved_model.tag_constants.TRAINING],
                                                strip_default_attrs=True)
            builder.add_meta_graph([tf.saved_model.tag_constants.SERVING], strip_default_attrs=True)
            builder.save()
            # END OF ADDITIONAL CODE
                                
           target_nodes = ",".join(self._process_graph())
           graph_def = self.graph.as_graph_def()
            
            output_graph_def = graph_util.convert_variables_to_constants(
                self.sess, graph_def, target_nodes.replace(" ", "").split(",")
            )
            frozen_graph_def_path = self.model_path + "/frozen_graph_def.pb"
            with gfile.GFile(frozen_graph_def_path, "wb") as f:
                f.write(output_graph_def.SerializeToString())
            tf2bc.convert(frozen_graph_def_path, self.model_path + ".nn")
            logger.info("Exported " + self.model_path + ".nn file")

This runs without error, but produces a SavedModel without any Signature Definitions, as seen when I run the python saved_model_cli.py show --dir ./ --all as per this:

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

MetaGraphDef with tag-set: 'train' contains the following SignatureDefs:

This then throws an error when converting to TFJS using the tfjs-converter:

ValueError: Signature 'serving_default' does not exist. The following signatures are available: KeysView(_SignatureMap({}))

along with many warnings like this:

WARNING:tensorflow:Unable to create a python object for variable <tf.Variable 'beta2_power:0' shape=() dtype=float32_ref> because it is a reference variable. It may not be visible to training APIs. If this is a problem, consider rebuilding the SavedModel after running tf.compat.v1.enable_resource_variables().

If anyone can shed light on how to include SignatureDefs in the exported model, and export a valid SavedModel from Unity ML-Agents, that would be awesome!

Thanks again.

[shihzy]

Thank you for your comments. We’ve documented your feedback and will prioritize when appropriate.

[shihzy]

also - opening this up as a more general discussion to see if folks in the community have attempted to do what you are describing above.

[AidanNelson]

For anyone trying to use models from Unity in TensorFlow.js, here is the approach I took.

• Install Unity ML-Agents for development by following this guide here
• In ml-agents/mlagents/trainers/tf_policy.py script, add the following lines to the export_model function to export a TensorFlow SavedModel. Note that the required graph nodes are different for continuous action space (i.e. the agent takes action on float values) and discrete action space (i.e. the agent takes action on integer values), so you will need to uncomment the respective lines accordingly:

def export_model(self):
        """
        Exports latest saved model to .nn format for Unity embedding.
        """

        with self.graph.as_default():
            graph_def = self.graph.as_graph_def()

            # BEGINNING OF ADDED CODE:
            # To learn more about nodes of the graph, uncomment the following lines:
            # for node in graph_def.node:
            #     print("-------")
            #     print(node.name)
            #     print(node.input)
            #     print(node.attr)
 
            # Uncomment for discrete vector action space:
            # vectorInputNode = self.graph.get_tensor_by_name("vector_observation:0")
            # actionMaskInput = self.graph.get_tensor_by_name("action_masks:0")
            # actionOutputNode = self.graph.get_tensor_by_name("action:0")
            # sigs = {}
            # sigs[tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
                tf.saved_model.signature_def_utils.predict_signature_def( \
                    {"in": vectorInputNode, "actionMask": actionMaskInput}, {"out": actionOutputNode})
            
            # Uncomment for continuous vector action space:
            # vectorInputNode = self.graph.get_tensor_by_name("vector_observation:0")
            # epsilonInputNode = self.graph.get_tensor_by_name("epsilon:0")
            # actionOutputNode = self.graph.get_tensor_by_name("action:0")
            # sigs = {}
            # sigs[tf.compat.v1.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
            #     tf.saved_model.signature_def_utils.predict_signature_def( \
            #         {"in": vectorInputNode,"epsilon": epsilonInputNode}, {"out": actionOutputNode})



            
            builder = tf.compat.v1.saved_model.Builder(self.model_path + "/SavedModel/")
            builder.add_meta_graph_and_variables( \
                self.sess, \
                [tf.saved_model.tag_constants.SERVING], \
                signature_def_map=sigs, \
                strip_default_attrs=True)
            builder.save()                                
            # END OF ADDED CODE
           
            target_nodes = ",".join(self._process_graph())
            
            output_graph_def = graph_util.convert_variables_to_constants(
                self.sess, graph_def, target_nodes.replace(" ", "").split(",")
            )
            frozen_graph_def_path = self.model_path + "/frozen_graph_def.pb"
            with gfile.GFile(frozen_graph_def_path, "wb") as f:
                f.write(output_graph_def.SerializeToString())
            tf2bc.convert(frozen_graph_def_path, self.model_path + ".nn")
            logger.info("Exported " + self.model_path + ".nn file")
     

• the exported SavedModel should be able to be converted to a web model by the latest tensorflowjs-converter (tfjs version 1.4.0, currently)

If anyone can clarify the different output nodes: action vs. action_probs/action_probs, please let me know!

[stale]

This issue has been automatically marked as stale because it has not had activity in the last 14 days. It will be closed in the next 14 days if no further activity occurs. Thank you for your contributions.

[stale]

This issue has been automatically closed because it has not had activity in the last 28 days. If this issue is still valid, please ping a maintainer. Thank you for your contributions.

[AidanNelson]

@unityjeffrey I would love to open this thread up again now that MLAgents is in an official release.

The ability to export models as Tensorflow SavedModel format would allow for interoperability with Tensorflow.js, a very well supported web-based inference engine with a large community around it. The ability to use models trained in Unity on the web (even outside of the Unity ecosystem on the web) would open up a lot of possibilities for Unity ML Agents as a training platform. The advantages of web-based inference are numerous. My particular interest is in creating web based educational tools around machine learning.

Here is a model trained in Unity, then exported as a Tensorflow SavedModel, converted for use with Tensorflow.js, and loaded in a javascript sketch: https://www.aidanjnelson.com/balance-ball/

Thoughts?

[shihzy]

hi @AidanNelson - quick question for you. Can you describe a use case where you would want to run an ML-Agents model via Tensorflow.js? Barracuda allows for ML-Agents models to be run in WebGL games.

[AidanNelson]

Sure!

Part of what is great about the Unity ML-Agents Framework is the ability to quickly spin up a training environment and get started experimenting with RL without a deep knowledge of the difference between SAC and PPO or the ability to tune the various training hyperparameters. ML-Agents provides an easier (and more powerful) access point to training RL than I've seen elsewhere, in an environment which many people are already familiar with. That said, there are a number of contexts in which someone might want the final version of their project to exist online outside of a Unity WebGL build. The first one that comes to mind is a classroom environment.

Teachers teaching Javascript (in a middle or high school environment, for instance) with the ability to incorporate a simple RL model into their curriculum would be able to:

• live-code on an environment in javascript and see results without needing to first compile,
• have students build applications using these models without needing to know how to use Unity,
• have students build applications using these models without needing to have Unity installed on their machines (in computer lab environment, for example),
• interact with other Javascript libraries / frameworks / browser APIs (webcam, etc.) in a familiar manner,
• host javascript sketches on various platforms (p5.js web editor, Glitch, etc.)

Personally, I want to work with incorporating models into projects using p5.js and ml5.js, libraries which have large and welcoming communities around them and are specifically geared towards accessibility and education.

Any thoughts?