基于TensorFlow的NES俄罗斯方块RL Agent Reshape维度错误求助

核心问题 #

你的错误源于两个关键问题：

1. 动作分布选择错误：你的动作是5个二元开关（每个动作0/1），应该使用伯努利分布而非分类分布。原代码中错误地为每个动作分配2个logits并尝试用分类分布处理，导致维度不匹配。
2. 训练数据维度未处理：训练时从回放缓冲区取出的数据带有时间步维度（num_steps=2），observations形状为(64,2,204)，但你的网络call方法只提取了batch_size=64，忽略了时间步维度，导致logits总数量与reshape目标形状不匹配。

修正后的自定义Actor网络 #

替换原CustomActorDistributionNetwork为以下代码，解决维度和分布问题：

class CustomActorDistributionNetwork(network.Network):
    def __init__(self, input_tensor_spec, output_tensor_spec, fc_layer_params=(100,), name='CustomActorDistributionNetwork'):
        super(CustomActorDistributionNetwork, self).__init__(input_tensor_spec, state_spec=(), name=name)

        self._output_tensor_spec = output_tensor_spec
        self._fc_layers = [tf.keras.layers.Dense(num_units, activation='relu') for num_units in fc_layer_params]
        # 每个二元动作只需要1个logits（伯努利分布的logits）
        num_actions = np.prod(output_tensor_spec.shape)
        self._logits_layer = tf.keras.layers.Dense(num_actions, activation=None)

    def call(self, observations, step_type=None, network_state=(), training=False):
        del step_type  # Unused.
        observations = tf.convert_to_tensor(observations)
        
        # 处理带时间步的输入：将batch和time维度合并，处理后再拆分
        # 输入形状可能是 [batch_size, time_steps, obs_dim]（训练时）或 [batch_size, obs_dim]（推理时）
        original_shape = tf.shape(observations)
        if len(original_shape) == 3:
            # 训练模式：合并batch和time维度
            batch_time_size = original_shape[0] * original_shape[1]
            obs_reshaped = tf.reshape(observations, [batch_time_size, original_shape[2]])
        else:
            # 推理模式：直接使用batch维度
            obs_reshaped = observations
            batch_time_size = original_shape[0]

        # 全连接层处理
        for layer in self._fc_layers:
            obs_reshaped = layer(obs_reshaped, training=training)

        # 生成logits：每个动作对应1个伯努利logits
        logits = self._logits_layer(obs_reshaped, training=training)
        num_actions = np.prod(self._output_tensor_spec.shape)
        
        # 恢复原始维度结构
        if len(original_shape) == 3:
            logits = tf.reshape(logits, [original_shape[0], original_shape[1], num_actions])
            # 定义独立伯努利分布：每个动作是独立的二元选择
            action_distribution = tfd.Independent(
                tfd.Bernoulli(logits=logits, dtype=tf.int32),
                reinterpreted_batch_ndims=1)
        else:
            logits = tf.reshape(logits, [original_shape[0], num_actions])
            action_distribution = tfd.Independent(
                tfd.Bernoulli(logits=logits, dtype=tf.int32),
                reinterpreted_batch_ndims=1)

        return action_distribution, network_state

关键修正点说明 #

1. 输出层单元数调整：从num_actions*2改为num_actions，因为每个二元动作只需要1个logits来控制伯努利分布的概率。
2. 多维度输入处理：自动识别训练时的[batch, time, obs]形状和推理时的[batch, obs]形状，合并/拆分维度保证网络在两种模式下都能正常工作。
3. 动作分布修正：使用tfd.Bernoulli替代tfd.Categorical，并通过tfd.Independent将5个独立的伯努利分布组合成一个联合分布，匹配你的5维二元动作空间。

其他注意事项 #

• 原环境的_step方法中奖励计算reward = log_time ** self._score可能会导致奖励值爆炸或梯度消失，建议调整为更合理的奖励函数（比如消除行数的即时奖励+生存奖励）。
• 训练时回放缓冲区的num_steps=2会生成带时间步的轨迹，修正后的网络已经适配这种输入，无需修改数据集配置。

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