keras agents fails in DQNAgent using PQC during clonation for target - keras

Question

I have some issues using keras-rl2 with tensorflow_quantum and VQC (using identical architecture as https://www.tensorflow.org/quantum/tutorials/quantum_reinforcement_learning)
After the creation of the model and DqnAgent, in dqn.compile:
############################################################
def generate_model_Qlearning(qubits, n_layers, n_actions, observables, target):
qubits = cirq.GridQubit.rect(1, n_qubits)
ops = [cirq.Z(q) for q in qubits]
observables = [ops[0]*ops[1], ops[2]*ops[3]] # Z_0*Z_1 for
action 0 and Z_2*Z_3 for action 1
input_tensor = tf.keras.Input(shape=(len(qubits), ),
dtype=tf.dtypes.float32, name='input')
re_uploading_pqc = ReUploadingPQC(qubits, n_layers,
observables, activation='tanh')([input_tensor])
process = tf.keras.Sequential(
[Rescaling(len(observables))],
name=target*"Target"+"Q-values"
)
Q_values = process(re_uploading_pqc)
model = tf.keras.Model(inputs=[input_tensor],
outputs=Q_values)
return model
############################################################
model = generate_model_Qlearning(qubits, n_layers, n_actions,
observables, False)
model_target = generate_model_Qlearning(qubits, n_layers,
n_actions, observables, True)
model_target.set_weights(model.get_weights())
dqn = DQNAgent(model=model, enable_double_dqn = True,
nb_actions=num_actions,
dqn.compile(Adam(learning_rate=1e-3), metrics=['mae'])
history = dqn.fit(env, nb_steps=50000, visualize=False,
verbose=2)
The following exception appears:
---------------------------------------------------------------------------
NotImplementedError Traceback (most recent call last)
Input In [119], in <module>
----> 1 dqn.compile(Adam(learning_rate=1e-3), metrics=['mae'])
3 history = dqn.fit(env, nb_steps=50000,
4 visualize=False,
5 verbose=2)
File ~.local/lib/python3.8/site-packages/rl/agents/dqn.py:167, in DQNAgent.compile(self, optimizer, metrics)
164 metrics += [mean_q] # register default metrics
166 # We never train the target model, hence we can set the optimizer and loss arbitrarily.
--> 167 **self.target_model = clone_model(self.model, self.custom_model_objects)**
168 self.target_model.compile(optimizer='sgd', loss='mse')
169 self.model.compile(optimizer='sgd', loss='mse')
File ~.local/lib/python3.8/site-packages/rl/util.py:13, in clone_model(model, custom_objects)
9 def clone_model(model, custom_objects={}):
10 # Requires Keras 1.0.7 since get_config has breaking changes.
11 config = {
12 'class_name': model.__class__.__name__,
---> 13 **'config': model.get_config(),**
14 }
15 clone = model_from_config(config, custom_objects=custom_objects)
16 clone.set_weights(model.get_weights())
File ~.local/lib/python3.8/site-packages/keras/engine/functional.py:685, in Functional.get_config(self)
684 def get_config(self):
--> 685 return copy.deepcopy(get_network_config(self))
File ~.local/lib/python3.8/site-packages/keras/engine/functional.py:1410, in get_network_config(network, serialize_layer_fn)
1407 node_data = node.serialize(_make_node_key, node_conversion_map)
1408 filtered_inbound_nodes.append(node_data)
-> 1410 layer_config = serialize_layer_fn(layer)
1411 layer_config['name'] = layer.name
1412 layer_config['inbound_nodes'] = filtered_inbound_nodes
File ~.local/lib/python3.8/site-packages/keras/utils/generic_utils.py:510, in serialize_keras_object(instance)
507 if _SKIP_FAILED_SERIALIZATION:
508 return serialize_keras_class_and_config(
509 name, {_LAYER_UNDEFINED_CONFIG_KEY: True})
--> 510 raise e
511 serialization_config = {}
512 for key, item in config.items():
File ~.local/lib/python3.8/site-packages/keras/utils/generic_utils.py:505, in serialize_keras_object(instance)
503 name = get_registered_name(instance.__class__)
504 try:
--> 505 config = instance.get_config()
506 except NotImplementedError as e:
507 if _SKIP_FAILED_SERIALIZATION:
File ~.local/lib/python3.8/site-packages/keras/engine/base_layer_v1.py:497, in Layer.get_config(self)
494 # Check that either the only argument in the `__init__` is `self`,
495 # or that `get_config` has been overridden:
496 if len(extra_args) > 1 and hasattr(self.get_config, '_is_default'):
--> 497 raise NotImplementedError('Layers with arguments in `__init__` must '
498 'override `get_config`.')
499 return config
NotImplementedError: Layers with arguments in `__init__` must override `get_config`.
the topology:
It could be great if this library let us specify the dqn_target instead of doing Clone.
Because working with hybrid neural networks with a cirquit with parameters in a layer, it's difficult to serialize it. So, when it runs the line: model.get_config(), it fails.
Any idea to solve it?
Thanks!