Error in loading policy from saved model.

[n3011]

I was trying to load a saved model policy using the following code, but it ends with errors pasted below

saved_policy = tf.saved_model.load('policy_0')

Software versions:

• Tensorflow =2.3.0
• tf-agents = 0.6.0

Traceback (most recent call last): File "", line 1, in File "/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 603, in load return load_internal(export_dir, tags, options) File "/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 633, in load_internal ckpt_options) File "/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 130, in init self._load_all() File "/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 147, in _load_all self._setup_functions_structures() File "/lib/python3.6/site-packages/tensorflow/python/saved_model/load.py", line 191, in _setup_functions_structures original_outputs = coder.decode_proto(proto.output_signature) File "/lib/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 127, in decode_proto return self._map_structure(proto, self._get_decoders()) File "/lib/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 82, in _map_structure return do(pyobj, recursion_fn) File "/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 245, in do_decode items = [(pair.key, decode_fn(pair.value)) for pair in key_value_pairs] File "/lib/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 245, in items = [(pair.key, decode_fn(pair.value)) for pair in key_value_pairs] File "/lib/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 82, in _map_structure return do(pyobj, recursion_fn) File "/lib/python3.6/site-packages/tensorflow/python/saved_model/nested_structure_coder.py", line 552, in do_decode "version of TensorFlow.)" % type_spec_proto.type_spec_class_name) ValueError: The type '_DistributionTypeSpec' is not supported by this version of TensorFlow. (The object you are loading must have been created with a newer version of TensorFlow.) `

[kbanoop]

Can you try these: https://github.com/tensorflow/agents/blob/master/tf_agents/policies/policy_loader.py https://github.com/tensorflow/agents/blob/bdf6694dec482f8f894ee61dea750e1c0d48ad3e/tf_agents/policies/py_tf_eager_policy.py#L102

[n3011]

Using tf_agents.policies.policy_loader.load I can load policy. Thanks.

In this case, my concern is how I can use TF Serving if the policy can not be loaded with standard saved_model.load?

[windmaple]

+1

I'm trying to convert a policy to TFLite and hit the same error. Not being fully 'Saved Model' is a real barrier.

[sguada]

Does it work on TF 2.4?

[windmaple]

Just tried nightly. Didn't work.

[ebrevdo]

_DistributionTypeSpec is a Tensorflow Probability object that we use in our policies because we serialize the _distribution() method as well. Tensorflow Probability renamed this object to _TFPTypeSpec about 4 months ago: https://github.com/tensorflow/probability/commit/2ffe65af9fbc1f1ce0a66a33073a243f3f089ed3

So in this case, the issue may be due to a policy saved with an older version of TFP but loaded with a recent version of TFP. I'll include the TFP team to see if they can make this backwards compatible.

[ebrevdo]

Another possible issue here is that the type spec is only registered if you import tfp.experimental; unfortunately may be lazy-loaded (https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/__init__.py#L118). If you could run two experiments:

1. Import tensorflow_probability then try to do saved_model.load
2. Import tensorflow_probabitlity.experimental then try to do saved_model.load

Let us know if doing either of these fixes the issue.

[david942j]

Facing the same issue with tensorflow==2.4.1, tf_agents==0.7.1

And adding "import tensorflow_probability" fixes the issue (tf.saved_model.load works)

[cserpell]

1. Import tensorflow_probability then try to do saved_model.load

I was facing a similar issue, not being able to load a policy. I added the import tensorflow_probability line at the beginning of the module, and it worked.