Convert an TF Agents ActorDistributionNetwork into a Tensorflow lite model

I would like to convert the ActorDistributionModel from a trained PPOClipAgent into a Tensorflow Lite model for deployment. How should I accomplish this?

I have tried following this tutorial (see section at bottom converting policy to TFLite), but the network outputs a single action (the policy) rather than the density function over actions that I desire.

I think perhaps something like this could work:

 tf.compat.v2.saved_model.save(actor_net, saved_model_path, signature=?)

... if I knew how to set the signature parameter. That line of code executes without error when I omit the signature parameter, but I get the following error on load (I assume because the signature is not set up correctly):

converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)
File "/home/ais/salesmentor.ai/MDPSolver/src/solver/ppo_budget.py", line 336, in train_eval
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path)  


  File "/home/ais/.local/lib/python3.9/site-packages/tensorflow/lite/python/lite.py", line 1275, in from_saved_model
    raise ValueError("Only support a single signature key.")
ValueError: Only support a single signature key.

Solution

This appears to work. I won't accept the answer until I have completed an end-to-end test, though.

def export_model(actor_net, observation_spec, saved_model_path):
  predict_signature = {
      'action_pred':
          tf.function(func=lambda x: actor_net(x, None, None)[0].logits,
                      input_signature=(tf.TensorSpec(shape=observation_spec.shape),)
                     )
  }
  tf.saved_model.save(actor_net, saved_model_path, signatures=predict_signature)

  # Convert to TensorFlow Lite model.
  converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_path,
                                                       signature_keys=["action_pred"])
  converter.target_spec.supported_ops = [
      tf.lite.OpsSet.TFLITE_BUILTINS,  # enable TensorFlow Lite ops.
      tf.lite.OpsSet.SELECT_TF_OPS  # enable TensorFlow ops.
  ]
  tflite_policy = converter.convert()
  with open(os.path.join(saved_model_path, 'policy.tflite'), 'wb') as f:
    f.write(tflite_policy)

The solution wraps the actor_net in a lambda because I was unable to figure out how to specify the signature with all three expected arguments. Through the lambda, I convert the function into using a single argument (a tensor). I expect to pass None to the other two arguments in my use case, so there is nothing lost in this approach.