解决TensorFlow TF-Agents DQN collect_policy中的InvalidArgumentError：批量大小与张量形状匹配问题

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

# 假设的Q网络模型（为完整示例提供）
class SimpleQNetwork(tf.keras.Model):
    def __init__(self, observation_spec, action_spec):
        super().__init__()
        self._action_spec = action_spec
        num_actions = action_spec.maximum - action_spec.minimum + 1
        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
        self.dense2 = tf.keras.layers.Dense(num_actions)

    def call(self, observation, step_type=None, network_state=()):
        # 确保Q网络能够处理输入的observation形状
        # 如果observation_spec是 (1, amountMachines)，实际输入可能是 (batch_size, 1, amountMachines)
        if observation.shape.rank > len(self.input_spec.observation.shape):
            # 移除多余的维度，例如 (batch_size, 1, obs_dim) -> (batch_size, obs_dim)
            observation = tf.squeeze(observation, axis=1)

        x = self.dense1(tf.cast(observation, tf.float32))
        q_values = self.dense2(x)
        return q_values, network_state

# 定义环境规格 (TimeStepSpec)
discount = 0.95
reward_val = 0.0
learning_rate = 1e-3
amountMachines = 6 # 示例观测维度

time_step_spec = ts.TimeStep(
    # step_type, reward, discount 都是每个时间步的标量，因此 shape=()
    step_type=tensor_spec.BoundedTensorSpec(shape=(), dtype=tf.int32, minimum=0, maximum=2),
    reward=tensor_spec.TensorSpec(shape=(), dtype=tf.float32),
    discount=tensor_spec.TensorSpec(shape=(), dtype=tf.float32),
    # observation 的 shape 是单个观测值的形状，不包含批量维度
    observation=tensor_spec.TensorSpec(shape=(1, amountMachines), dtype=tf.int32)
)

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

# 初始化 DQN Agent
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
train_step_counter = tf.Variable(0)

model = SimpleQNetwork(time_step_spec.observation, action_spec)

agent = dqn_agent.DqnAgent(
    time_step_spec,
    action_spec,
    q_network=model,
    optimizer=optimizer,
    epsilon_greedy=1.0,
    td_errors_loss_fn=common.element_wise_squared_loss,
    train_step_counter=train_step_counter)
agent.initialize()

# 模拟环境状态获取函数
def get_states():
    # 假设返回一个形状为 (amountMachines,) 的 NumPy 数组或 Tensor
    return tf.constant([4, 4, 4, 4, 4, 6], dtype=tf.int32)

# 获取当前状态并添加批量维度
current_state = get_states() # 示例输出: tf.Tensor([4 4 4 4 4 6], shape=(6,), dtype=int32)
current_state_batch = tf.expand_dims(current_state, axis=0) # 形状变为 (1, 6)

# 示例标量值
step_type_val = 0
reward_val = 0.0
discount_val = 0.95

# 创建 TimeStep 实例
# 注意：对于标量值，即使批量大小为1，也需要包装成 [value] 以创建形状为 (1,) 的张量
time_step = ts.TimeStep(
    step_type=tf.convert_to_tensor([step_type_val], dtype=tf.int32), # 形状 (1,)
    reward=tf.convert_to_tensor([reward_val], dtype=tf.float32),     # 形状 (1,)
    discount=tf.convert_to_tensor([discount_val], dtype=tf.float32), # 形状 (1,)
    observation=current_state_batch # 形状 (1, 6)
)

print(f"TimeStep created with shapes:")
print(f"  step_type: {time_step.step_type.shape}")
print(f"  reward: {time_step.reward.shape}")
print(f"  discount: {time_step.discount.shape}")
print(f"  observation: {time_step.observation.shape}")

# 调用 collect_policy
action_step = agent.collect_policy.action(time_step)
print(f"Action Step: {action_step.action}")