按照 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() 函数,以合并多个优化器,但我失败了。