How to save/load tensorflow 2.3.1 serialized model (DoubleDuelingDQN Agent) with custom arguments?

Question

Following a couple of tutorials and Stackoverflow questions, I managed to get my Double Dueling DQN Agent learning in my custom Environment.

Now I am unable to save and load my Agent's network model with its advantage() argument. How do I save & load such custom serialized networks?

Python 3.8.5 Tensorflow 2.3.1 Keras 2.4.3

My approach so far was simply:

class DuelingDeepQNAgent():
    ....
    def load_model(self, model_name):
        self.q_network = load_model(model_name)
        self.target_network = load_model(model_name)

    def save_model(self, model_name):
        self.q_network.save(model_name, save_format='tf')

Loading the saved model and continuing on learning results in that error:

agent = MyAgent.DuelingDeepQNAgent()
agent.load_model(model_name)
.... 
# within learning process
q_values = self.q_network.advantage(state)

AttributeError: 'DuelingDeepQNetwork' object has no attribute 'advantage'

Network:

class DuelingDeepQNetwork(keras.Model):
    def __init__(self, input_shape, output_shape, n_filter=64, kernel_size=3, n_neurons1=256, n_neurons2=128):
        super(DuelingDeepQNetwork, self).__init__()

        self.input_layer = keras.layers.InputLayer(input_shape=(None, *input_shape))
        # self.lstm1 = keras.layers.LSTM(units=n_neurons1, activation='relu', stateful=True, return_state=True)
        self.conv1d1 = keras.layers.Conv1D(filters=n_filter, kernel_size=kernel_size, padding='same', activation='relu')
        self.pool1 = keras.layers.AvgPool1D(pool_size=2)
        self.conv1d2 = keras.layers.Conv1D(filters=n_filter, kernel_size=kernel_size, padding='same', activation='relu')
        self.pool2 = keras.layers.AvgPool1D(pool_size=2)
        self.flat1 = keras.layers.Flatten()
        self.dens_1 = keras.layers.Dense(n_neurons1, activation='relu')
        self.dens_2 = keras.layers.Dense(n_neurons2, activation='relu')

        self.V = keras.layers.Dense(1, activation=None)  # Value layer
        self.A = keras.layers.Dense(output_shape, activation=None)  # Advantage layer

    def call(self, state):
        x = self.input_layer(state)
        x = self.conv1d1(x)
        x = self.pool1(x)
        x = self.conv1d2(x)
        x = self.pool2(x)
        x = self.flat1(x)
        x = self.dens_1(x)
        x = self.dens_2(x)

        A = self.A(x)
        V = self.V(x)
        Q = V + (A - tf.math.reduce_mean(A, axis=1, keepdims=True))

        return Q

    def advantage(self, state):
        x = self.input_layer(state)
        x = self.conv1d1(x)
        x = self.pool1(x)
        x = self.conv1d2(x)
        x = self.pool2(x)
        x = self.flat1(x)
        x = self.dens_1(x)
        x = self.dens_2(x)

        A = self.A(x)
        return A

Answer 1

EDIT: My previous answer was making an hypothesis on the missing part of the code but I believe the below is actually the right answer.

I actually worked a bit more on a similar issue, the problem you probably have is that when you load the saved H5 model, it doesn't load the associated functions of the custom class. What you need to do is to save/load only the weights to ensure you keep the class methods.

def save_model(self):
     self.q_network.save_weights("model_name") #no *.h5 extension
def load_model(self):
    self.q_network.load_weights("model_name")