I am using the TF-Agntes Turotial 7 to set up and train a SAC on a custom environment. The training is complete and working extremely well. I now whish to export the trained policy to a TF-Lite model but this is deemed unsuccessful as the policy setup in TF-Agents’s Tutorial 7 are eager policies and not TF Policies. Any help as to convert the eager policy to a TF-Lite model will be appreciated.
I apologize for the delayed response, I far I know If you don’t want to use TF policy, then you can also use the saved_model directly with the Python env through the use of py_tf_eager_policy.SavedModelPyTFEagerPolicy . something like below
Note that this only works when eager mode is enabled.
eager_py_policy = py_tf_eager_policy.SavedModelPyTFEagerPolicy(
policy_dir, eval_py_env.time_step_spec(), eval_py_env.action_spec())
# Note that we're passing eval_py_env not eval_env.
run_episodes_and_create_video(eager_py_policy, eval_py_env, eval_py_env)
Convert policy to TFLite
See TensorFlow Lite converter for more details.
converter = tf.lite.TFLiteConverter.from_saved_model(policy_dir, signature_keys=["action"])
tflite_policy = converter.convert()
with open(os.path.join(tempdir, 'policy.tflite'), 'wb') as f:
f.write(tflite_policy)
Please refer this tutorial for details, If I have missed something here please let me know.
Thank you for your cooperation and patience.