TF Agents自定义环境训练报uniform_table张量不兼容错误

错误根因 #

该InvalidArgumentError本质是Reverb的uniform_table表初始化时定义的张量规格，和实际采样取出的张量shape/dtype不匹配：报错指向扁平化索引4位置的字段，要求是一维变长int32张量（shape [?]），实际写入的是二维[2,1]张量，触发点在数据集迭代取数的IteratorGetNext算子，属于典型的自定义环境/数据链路维度对齐问题。

排查步骤 #

• 第一步先定位错位字段：在调用回放缓冲区写入方法add_batch前，加日志打印待写入轨迹（Trajectory）扁平化后的所有字段信息，直接找到索引4对应的字段：

# 插入到replay_buffer.add_batch(traj) 之前
for idx, field in enumerate(tf.nest.flatten(traj)):
    print(f"索引{idx} | dtype: {field.dtype} | shape: {field.shape}")

运行1步数据收集即可定位到shape为[2,1]的异常字段，90%以上概率是离散动作、步类型（step_type）、奖励（reward）、折扣因子（discount）这类本应为标量/一维的字段，被多余添加了长度为1的尾部维度。

• 核对自定义环境的spec定义：检查cGame类中action_spec()、time_step_spec()、observation_spec()返回的张量规格，是否和_reset()、_step()方法实际返回的张量shape完全匹配。常见错误是给离散动作、奖励这类标量字段定义了shape为(1,)的spec，而非标准的0维()。
• 核对Reverb表初始化逻辑：检查回放缓冲区初始化时传入的dataset_spec是否直接复用agent的collect_data_spec，是否存在手动修改spec给某个字段额外加维度的操作。

修复方案 #

1. 清理多余维度：定位到异常字段后，对带多余长度1维度的字段调用tf.squeeze(field, axis=-1)压掉尾部冗余维度。
   • 错误写法参考：动作采样后用双层括号包裹为np.array([[action]])，单样本下shape为(1,)，拼接batch后就会变成[batch_size, 1]的二维形状；
   • 正确写法：标量字段直接返回0维张量/数组，batch维度拼接后自然为要求的一维[batch_size]形状，和Reverb表的[?]规格匹配。
2. 修正环境spec：所有离散动作、step_type、reward、discount字段的spec shape必须设为()，observation的spec按实际21×21网格的观测形状设置即可，不要给标量字段加任何额外维度。
3. 并行环境场景额外校验：如果使用了2个并行环境收集数据（对应异常shape里的第一维长度2），要保证所有时间步标量字段输出shape为[num_envs]，观测shape为[num_envs, 21, 21, ...]，不要在维度拼接时把时间步维度、batch维度顺序搞反。

验证方式 #

修复后先不启动训练，单独测试数据链路写入+采样流程：

# 采样测试代码
dataset = replay_buffer.as_dataset(
    sample_batch_size=32,
    num_steps=2 # DQN默认用2步相邻样本计算TD误差
)
test_iter = iter(dataset)
test_batch = next(test_iter) # 该行不报错即说明维度完全对齐

测试通过后再启动正式训练循环，就不会再触发该取数报错。

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