M1 Mac下使用TF-Agents训练DQN时出现malloc错误求解决

版本适配修复 #

• 你当前使用的TF 2.6.2、tf-agents 0.10.0版本对ARM架构M1芯片的原生支持存在已知内存泄漏bug，直接替换为适配M1的版本组合即可解决90%以上的同类问题：
  • 使用Apple官方conda源安装tensorflow-deps=2.10
  • 安装tensorflow-macos==2.10、tensorflow-metal==0.6
  • 安装匹配的tf-agents版本：tf-agents==0.15.0
• 不要使用Rosetta2转译的x86版本Python和TF依赖，所有依赖优先使用arm64架构版本。

数据集并行读取冲突修复 #

你构造回放缓冲区数据集时开启的多线程读取和tf.function包装的训练步骤存在线程竞争，触发野指针访问导致malloc校验错误，按以下顺序修改：

1. 把memory.as_dataset的参数修改为：

sample = memory.as_dataset(
    sample_batch_size=params.batch_size.value,
    single_deterministic_pass=False,
    num_parallel_calls=1,
    num_steps=2
)

直接移除prefetch(3)配置，关闭并行读取和预取，测试错误是否复现。
2. 若修改后仍报错，移除agent_tf.train = common.function(agent_tf.train)这行代码，关闭训练步骤的图编译，用eager模式执行验证。
3. 若eager模式无报错，可在代码开头添加配置关闭M1下存在问题的XLA编译：

import tensorflow as tf
tf.config.optimizer.set_jit(False)

变量冲突修复 #

你同时定义了train_step_counter和global_step两个独立的全局步数变量，epsilon衰减和agent训练分别使用不同的步数变量，会导致tf.function追踪时出现变量泄漏，直接统一变量即可：

• 删掉train_step_counter = tf.Variable(0)这行代码
• 初始化DQNAgent时直接传入train_step_counter=global_step

自定义环境校验 #

检查你封装的自定义TF环境的所有输出：

• 观测值、奖励、终止标记等返回值必须是tf.Tensor类型，禁止返回numpy数组或者Python原生数值
• 所有返回值的dtype、shape必须和你定义的observation_spec、time_step_spec完全匹配，避免类型不匹配导致的异常内存拷贝。

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