使用 TensorFlow Agents 的自定义 Keras 层中不同变量的不同学习率

Question

按照 TensorFlow Agents 的 Train-a-deep-Q-network 教程，我正在实现一个顺序模型，其中包括一个自定义层作为 q_net，以作为参数传递给 DqnAgent(..., q_network=q_net, . ..) 班级：

自定义图层的简化草图：

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, n_layers, activation="linear", name="MyLayer"):
        super(CustomLayer, self).__init__(name=name)
        self.n_layers = n_layers

    self.alphas = tf.Variable(1., trainable=True)
    self.betas = tf.Variable(1., trainable=True)

    ...

使用部分代码：

layer1 = tf.keras.layers.Dense(64)
layer2 = CustomLayer(16)
layer3 = tf.keras.layers.Dense(4)

q_net = sequential.Sequential([layer1, layer2, layer3])

agent = dqn_agent.DqnAgent(
    train_env.time_step_spec(),
    train_env.action_spec(),
    q_network=q_net,
    ...)

此自定义层有两组不同的可训练 tf.Variables，self.alphas 和 self.betas（类似于自定义密集层中的“self.kernel”和“self.bias”）。现在，我希望对两组不同的参数有不同的学习率，这可以使用梯度带、多个优化器和 model.trainable_variables 属性来实现，就像在本教程中对 irreforme_update() 函数所做的那样：

optim_A = tf.keras.optimizers.SGD(learning_rate=0.1)
optim_B = tf.keras.optimizers.SGD(learning_rate=0.001)

def reinforce_update(states, actions, returns, model):
    ...

    with tf.GradientTape() as tape:
        ...

    grads = tape.gradient(loss, model.trainable_variables)
    for optimizer, w in zip([optim_A, optim_B], [0, 1]):
        optimizer.apply_gradients([(grads[w], model.trainable_variables[w])])

但是，我很难将其应用到 TensorFlow Agents 设置中，其中训练是通过 DqnAgent.train() 方法进行的。

使用 DqnAgent.train() 函数时，有没有办法让自定义层的不同参数集具有不同的学习率？

我尝试以某种方式修改 DqnAgent.train() 函数，以合并多个优化器，但我失败了。