Error when saving model with tensorflow-agents

Question

I am trying to save a model with tensorflow-agents. First I define the following:

collect_policy = tf_agent.collect_policy
saver = PolicySaver(collect_policy, batch_size=None)

and then save the model like this:

saver.save('my_directory/')

This works ok in google colab but I am getting the following error in my local PC.

AttributeError: module 'tensorflow.python.saved_model.nested_structure_coder' has no attribute 'StructureCoder'

These are the library versions I am using:

tensorflow 2.9.1

tf-agents 0.11.0

Answer 1

Tl;dr

Make sure you have the right tensorflow-probability version that is compatible for 2.9.x and tf-agents 0.11.0

pip uninstall tensorflow-probability
pip install tensorflow-probability==0.17.0

(0.19.0 for TF 2.11, 0.18.0 for TF 2.10 or look at the release notes)

Also make sure to restart your kernel from the notebook.

What the problem was

StructureCoder has been moved to tensorflow API. Therefore, other dependent libraries have made changes like this in tf-agent and like this in tensorflow-probability. Your machine is somehow picking up an older version that depends on the previous version of nested_structure_coder.

For me, I was using tensorflow 2.9.0 tf-agents 0.13.0 tensorflow-probabilities 0.17.0

Try making an explicit import in your notebook:

import tensorflow_probability
print(tensorflow_probability.__version__) # was 0.17.0 for me