TensorFlow Agents DQN InvalidArgumentError: 解决 collect_policy 中的形状不匹配问题

理解 tf_agents 中的数据结构：TimeStepSpec 与 TimeStep

在 tf_agents 中，timestepspec 和 timestep 是定义环境交互数据的核心结构。

• TimeStepSpec：这是一个规范（specification），它定义了单个时间步中每个组件（如 step_type、reward、discount、observation）的预期数据类型（dtype）和形状（shape）。这里的 shape 指的是单个时间步中该组件的形状，不包含批处理维度。
• TimeStep：这是实际的数据容器，它包含了当前时间步的真实张量值。在 tf_agents 中，TimeStep 通常是批处理的，这意味着每个组件的张量都会有一个额外的最外层维度，表示批处理大小（batch_size）。

理解 spec.shape（单个元素形状）与 tensor.shape（批处理张量形状）之间的区别至关重要。如果 spec.shape 是 S，且批处理大小为 B，那么对应的实际张量形状应该是 (B,) + S。

错误根源：标量形状的误解

InvalidArgumentError: 'then' and 'else' must have the same size. but received: [1] vs. [] 这个错误通常发生在策略内部，特别是 epsilon_greedy_policy 等探索性策略中，它们会使用 tf.where 等条件语句来选择贪婪动作或随机动作。当 tf.where(condition, x, y) 被调用时，x 和 y 必须能够广播到 condition 的形状。如果 x 和 y 的形状不一致，就会触发此错误。

导致此问题的常见误区在于对 TimeStepSpec 中标量组件的形状定义。考虑以下原始的 TimeStepSpec 定义片段：

time_step_spec = TimeStep(
    step_type = tensor_spec.BoundedTensorSpec(shape=(1,), dtype=tf.int32, minimum=0, maximum=2),
    reward = tensor_spec.TensorSpec(shape=(1,), dtype=tf.float32),
    discount = tensor_spec.TensorSpec(shape=(1,), dtype=tf.float32),
    observation = tensor_spec.TensorSpec(shape=(1,amountMachines), dtype=tf.int32)
)

以及对应的 TimeStep 数据创建片段（假设 batch_size = 1）：

time_step = TimeStep(
    step_type=tf.convert_to_tensor([step_type_value], dtype=tf.int32), # 结果形状 (1,)
    reward=tf.convert_to_tensor([reward_value], dtype=tf.float32),     # 结果形状 (1,)
    discount=tf.convert_to_tensor([discount_value], dtype=tf.float32), # 结果形状 (1,)
    observation=current_state_batch # 结果形状 (1, amountMachines)
)

乍一看，time_step_spec 中 shape=(1,) 与 time_step 中 tf.convert_to_tensor([value]) 产生的 (1,) 形状似乎是匹配的。然而，这里的关键在于 tf_agents 对 TimeStepSpec 中 shape 的解释：

• 如果一个组件（如 reward）在单个时间步中是一个标量值，那么其 TimeStepSpec 中的 shape 应该定义为 ()（空元组），表示一个零维张量。
• 当实际创建 TimeStep 张量时，如果 batch_size 为 1，则一个标量值需要被包装成 [value]，这会创建一个形状为 (1,) 的一维张量。这个 (1,) 形状是 (batch_size,) + spec.shape 的结果，其中 batch_size=1 且 spec.shape=()。

因此，当 TimeStepSpec 中将标量定义为 shape=(1,) 时，tf_agents 内部会期望每个批处理元素都是一个形状为 (1,) 的张量。但实际传入的 tf.convert_to_tensor([value]) 产生的 (1,) 形状，在某些内部操作（如 tf.where）中，可能会与预期产生混淆，导致一个分支得到 (1,) 而另一个得到 () 的动作张量，从而引发 InvalidArgumentError。

解决方案：修正 TimeStepSpec 中的标量形状

解决此问题的关键是确保 TimeStepSpec 正确反映每个组件在单个时间步内的实际维度。对于像 step_type、reward 和 discount 这样的标量值，它们的 shape 应该定义为 ()。

正确的 TimeStepSpec 定义：

import tensorflow as tf
from tf_agents.trajectories import time_step as ts
from tf_agents.specs import tensor_spec
from tf_agents.agents.dqn import dqn_agent
from tf_agents.utils import common

# 假设 amountMachines 已定义，例如 amountMachines = 6
amountMachines = 6
learning_rate = 1e-3
train_step_counter = tf.Variable(0) # 示例

# 修正后的 time_step_spec
time_step_spec = ts.TimeStep(
    step_type = tensor_spec.BoundedTensorSpec(shape=(), dtype=tf.int32, minimum=0, maximum=2), # 修正：shape=()
    reward = tensor_spec.TensorSpec(shape=(), dtype=tf.float32),                                 # 修正：shape=()
    discount = tensor_spec.TensorSpec(shape=(), dtype=tf.float32),                               # 修正：shape=()
    observation = tensor_spec.TensorSpec(shape=(amountMachines,), dtype=tf.int32)                # 修正：shape=(amountMachines,)
)

num_possible_actions = 729
action_spec = tensor_spec.BoundedTensorSpec(
    shape=(), dtype=tf.int32, minimum=0, maximum=num_possible_actions - 1)

# 假设

以上就是TensorFlow Agents DQN InvalidArgumentError: 解决 collect_policy 中的形状不匹配问题的详细内容，更多请关注php中文网其它相关文章！