训练60维离散动作DQN时出现mtrand.pyx概率含NaN报错的解决方案

根因澄清 #

这个报错的核心触发逻辑不是numpy本身的计算缺陷，是训练过程中DQN网络前向输出的Q值出现NaN/Inf，经过Softmax策略层计算后得到全无效的概率分布，最终传给np.random.choice时触发参数校验报错。

• 低维度动作空间下不会稳定复现的原因：动作维度低时，Q值计算的数值溢出概率小，即使个别Q值出现偏差，也不会导致整个概率数组全量失效；当动作维度到60维时，Q值输出层的数值波动范围更大，只要单个Q值出现Inf，整个Softmax计算就会得到全NaN的概率数组，直接触发报错。
• 本地找不到mtrand.pyx是正常现象：1.19.5版本的numpy是通过pip安装的预编译wheel包，Cython源码不会直接释放到本地路径，需要手动重新编译才能修改，改造成本极高，完全没必要为了这个场景折腾。
• 注意 就算把对应官方numpy修复打上去，只要根源的Q值NaN问题不解决，还是会在numpy侧抛出“概率包含NaN”的报错，修复逻辑只是把报错位置往上传了一层，没有解决根本问题。

分层可落地方案（从低改造成本到根源解决） #

1. 最小改动兜底拦截（10分钟可完成） #

直接修改本地环境中keras-rl2的策略层代码，在传入numpy随机选择逻辑前加异常校验，不用动numpy本身：

• 定位文件路径：你的venv路径/Lib/site-packages/rl/policy.py
• 找到SoftmaxPolicy类的select_action方法，定位到原逻辑中计算完probs、即将调用np.random.choice的代码位置，替换为带兜底校验的逻辑：

  # 替换原probs计算完成后的执行逻辑
  # 先校验概率数组的有效性
  if not np.all(np.isfinite(probs)):
      # 出现NaN/Inf时直接退化为均匀随机选动作，跳过无效参数传给numpy
      action = np.random.randint(0, nb_actions)
  else:
      # 加一层概率归一化兜底，避免浮点累计误差导致概率总和偏离1
      probs = probs / probs.sum()
      action = np.random.choice(range(nb_actions), p=probs)
  

• 这个改动完全覆盖报错触发的场景，不会改变原有训练逻辑，只是在异常值出现时做容错，避免程序直接崩溃中断训练。

2. 根源解决训练数值不稳定问题（从根上减少NaN出现） #

兜底只是避免程序崩溃，还是要从训练环节调整减少Q值异常：

• 降低DQN的初始学习率：60维动作空间下Q值输出层的数值范围更大，过高的学习率很容易导致权重震荡出NaN，建议将原来的学习率从1e-3调整到1e-4再重新训练
• 对奖励值做标准化：把单步奖励缩放到[-1, 1]区间，避免奖励幅值过大导致Q值计算溢出为Inf/NaN
• 给网络输出层加数值截断：在自定义的DQN主网络的最后输出层，增加tf.clip_by_value(q_output, -10, 10)的逻辑，限制Q值的输出范围，避免Softmax计算时因为输入值过大/过小溢出
• 统一设置numpy、tensorflow的全局随机种子，减少权重初始化的随机偏差，避免训练初期就出现权重发散的问题

3. 如果你确实要应用numpy的官方修复逻辑 #

不要手动找源码编译，直接安装兼容版本的numpy：

• 你当前的TensorFlow 2.4.1仅兼容numpy 1.19.x版本，直接将现有numpy从1.19.5升级到1.19.6即可，这个版本已经把对应分支的修复逻辑合入，完全满足依赖兼容要求，不需要重新编译Cython源码，直接替换安装就能生效。
• 不要升级numpy到1.20以上版本，会直接触发TensorFlow 2.4.1的二进制依赖兼容报错。

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