自定义TensorFlow训练函数时梯度无法计算的问题求助

参考TensorFlow官方示例修改自定义训练函数，加入强化学习的ε-greedy动作选择逻辑后，梯度计算功能失效，TensorFlow返回错误：

ValueError: No gradients provided for any variable: ['dense/kernel:0', 'dense/bias:0'].

问题代码 #

class CustomModel(keras.Model):
    def train_step(self, data):
        x, y = data

        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)  # Forward pass

            # Compute our own loss
            metric = tf.math.argmin(y_pred, axis=1)

            loss = keras.losses.mean_squared_error(y, metric)

        # Compute gradients
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Compute our own metrics
        loss_tracker.update_state(loss)
        mae_metric.update_state(y, y_pred)
        return {"loss": loss_tracker.result(), "mae": mae_metric.result()}

核心问题分析 #

tf.math.argmin是不可导的离散操作，它直接返回张量中最小值的索引，这个过程切断了y_pred（模型输出）到loss之间的梯度传播链路。因为离散索引没有梯度信息，TensorFlow无法计算loss对模型可训练变量的梯度，最终导致tape.gradient()返回空梯度。

解决方案 #

强化学习中，动作选择（ε-greedy）是离散决策，但loss计算必须基于模型输出的连续价值（如Q值）来构建可导路径：

1. 分离动作选择与梯度计算：ε-greedy的动作选择逻辑放在tf.GradientTape上下文之外，避免离散操作干扰梯度图。
2. 基于连续输出计算loss：不要用动作索引计算loss，而是用模型输出的对应动作的价值（如Q值）。

修改示例（以DQN为例）

class CustomModel(keras.Model):
    def __init__(self, num_actions, epsilon=0.1, gamma=0.99, **kwargs):
        super().__init__(**kwargs)
        self.num_actions = num_actions
        self.epsilon = epsilon
        self.gamma = gamma
        # 初始化指标
        self.loss_tracker = keras.metrics.Mean(name="loss")
        self.mae_metric = keras.metrics.MeanAbsoluteError(name="mae")

    def train_step(self, data):
        # data包含状态、动作、奖励、下一个状态等RL数据（根据你的场景调整）
        x, actions, rewards, next_x, done = data

        with tf.GradientTape() as tape:
            # 前向传播得到当前状态的Q值
            current_q = self(x, training=True)
            # 计算下一个状态的最大Q值
            next_q = self(next_x, training=False)
            max_next_q = tf.reduce_max(next_q, axis=1)
            # 计算目标Q值
            target_q = rewards + (1 - done) * self.gamma * max_next_q
            # 取出当前模型对所选动作的Q值
            selected_current_q = tf.reduce_sum(current_q * tf.one_hot(actions, self.num_actions), axis=1)
            # 计算可导的MSE loss
            loss = keras.losses.mean_squared_error(target_q, selected_current_q)

        # 计算并应用梯度
        trainable_vars = self.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # 更新指标
        self.loss_tracker.update_state(loss)
        self.mae_metric.update_state(target_q, selected_current_q)
        return {"loss": self.loss_tracker.result(), "mae": self.mae_metric.result()}

    # ε-greedy动作选择放在tape外，推理时使用
    def select_action(self, state):
        if tf.random.uniform(()) < self.epsilon:
            return tf.random.uniform((1,), 0, self.num_actions, dtype=tf.int32)
        else:
            q_values = self(state, training=False)
            return tf.math.argmax(q_values, axis=1)

关键要点 #

• 离散动作选择（如argmin/argmax）不能参与梯度计算图，必须放在GradientTape之外。
• loss必须基于模型输出的连续张量（如Q值、策略概率）构建，保证梯度能从loss反向传播到模型参数。

内容的提问来源于stack exchange，提问作者mindstorm84