Actor-Critic强化学习中神经网络输出NaN的原因及解决方法

构建了如下Actor网络用于Actor-Critic强化学习，输入为取值范围0-1的缩放图像：

obs_shape = (288, 512, 3)
num_actions = 2

input = layers.Input(shape=obs_shape)
x = layers.Conv2D(16, (3, 3), activation='relu')(input)
x = layers.MaxPool2D((2, 2))(x)
x = layers.Conv2D(32, (3, 3), activation='relu')(x)
x = layers.MaxPool2D((2, 2))(x)
x = layers.Conv2D(64, (3, 3), activation='relu')(x)
x = layers.MaxPool2D((2, 2))(x)
x = layers.Flatten()(x)
x = layers.Dense(64, activation='relu')(x)
x = layers.Dense(32, activation='relu')(x)
actions = layers.Dense(num_actions, activation='softmax')(x)
actor = keras.Model(input, actions, name="actor")

训练仅两个回合后，执行actor.predict(observation)得到的结果数组中包含NaN值，需排查原因并解决。

• 学习率过高：RMSprop默认学习率（0.001）在Actor-Critic场景下可能偏大，策略梯度的方差本就较高，大学习率会引发梯度爆炸，直接导致权重更新为NaN。
• Softmax数值溢出：如果最后一层Dense输出的logits值过大，exp(logits)会超出浮点数范围变为inf，归一化后就会产生NaN。
• 自定义梯度计算漏洞：若手动实现策略梯度（未使用Keras内置损失），可能存在未处理的除以零、log(0)等操作，导致梯度出现inf/NaN并传播到参数更新中。
• 奖励信号异常：未做归一化的奖励若出现极大值，会导致损失计算爆炸，进而让参数更新后变为NaN。
• 梯度未裁剪：Actor-Critic训练中梯度波动大，未对梯度做裁剪会导致参数更新幅度过大，触发数值异常。

• 调低学习率并尝试衰减：将RMSprop的学习率降至0.0001或更小，例如：

  optimizer = keras.optimizers.RMSprop(learning_rate=1e-4)
  

  也可添加学习率衰减，比如每回合按比例降低学习率。
• 稳定Softmax计算：避免直接在Dense层加softmax，先计算logits并做数值裁剪，再手动计算softmax：

  logits = layers.Dense(num_actions)(x)
  logits = tf.clip_by_value(logits, -10.0, 10.0)  # 限制logits范围，防止exp溢出
  actions = tf.nn.softmax(logits)
  actor = keras.Model(input, actions, name="actor")
  

• 检查并修复自定义梯度：若手动计算梯度，需给log操作添加小epsilon避免log(0)，同时做梯度裁剪：

  probs = actor(observation)
  log_probs = tf.log(probs + 1e-8)  # 添加epsilon防止log(0)
  grads = tf.gradients(loss, actor.trainable_variables)
  grads, _ = tf.clip_by_global_norm(grads, 1.0)  # 全局梯度裁剪
  optimizer.apply_gradients(zip(grads, actor.trainable_variables))
  

• 归一化奖励信号：对每回合的奖励做标准化处理，例如：

  rewards = np.array(episode_rewards)
  rewards = (rewards - np.mean(rewards)) / (np.std(rewards) + 1e-8)
  

  或直接将奖励裁剪到[-1, 1]范围内。
• 添加梯度裁剪与保守初始化：在优化器中启用梯度裁剪，同时给网络层使用更稳定的初始化方式：

  # 优化器添加梯度裁剪
  optimizer = keras.optimizers.RMSprop(learning_rate=1e-4, clipnorm=1.0)
  # 卷积和全连接层使用he_normal初始化（适配relu激活）
  x = layers.Conv2D(16, (3, 3), activation='relu', kernel_initializer='he_normal')(input)
  # ...其余层同理
  

• 验证输入数据：确认训练过程中输入的observation没有NaN值，可在数据加载后添加np.all(np.isfinite(observation))检查。

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