I am trying to save a batch-constrained Q-learning model written by AurelianTactics (code repo here) with the end goal of converting to TFLite. The BCQ class is written with TensorFlow functions but is not itself a TensorFlow class object. The existing code saves the model using a function that calls tf.compat.v1.train.Saver(), which saves the into checkpoint file and ckpt files with .data-00000-of-00001, .index, .meta extensions. To my knowledge these only save model weights, etc. and not the actual model itself, so the model class file can read and continue training it but without the base code you cannot select an action. In addition, you cannot convert the files to a TensorFlow Lite model. Attempting to use tf.saved_model.save() gives the error

ValueError: Checkpoint was expecting a trackable object (an object derived from TrackableBase), got <BCQ.BCQ object at 0x00000143F797F2B0>. If you believe this object should be trackable (i.e. it is part of the TensorFlow Python API and manages state), please open an issue.

How should I go about rewriting the code so that I am able to save the files to a TensorFlow savedmodel and convert to TFLite?

I also don’t fully understand the difference between writing a model in the way it was done here using TensorFlow functions and writing a model as a tf.Module class.

My full code is here, with only small modifications from the original. The only files of interest should be main.py and BCQ.py. DDPG trains a policy to collect data with, generate_buffer generates a replay buffer for BCQ to train on, and main trains BCQ policy.

Using Python 3.7.1, trfl 1.1.0, numpy 1.16.5, scipy 1.7.3, tensorflow-gpu 1.14.0, tensorflow-probability 0.7, gast 0.2.2, gym 0.12.5, and pyglet 1.4.0 to avoid package conflicts

I have no idea how I should go about changing the code. I know I should try to make the class inherit tf.Module but I don’t know if the easiest way to do this would be through Keras or something else. I would appreciate any pointers. Thank you!