强化学习Actor-Critic模型np.random.choice报错：'p'必须为一维数组

问题背景 #

作为强化学习新手，参考Keras官方CartPole Actor-Critic示例修改代码，希望基于概率生成0到max_x-1之间的随机整数x。但当前x_probs的形状为(1,10,10)，运行代码时触发ValueError，提示'p' must be 1-dimensional。

报错详情 #

ValueError
'p' must be 1-dimensional
  File "C:\Users\cglab\Desktop\Match3\Train.py", line 48, in main
    x = np.random.choice(range(max_x), p=np.squeeze(x_probs))
  File "C:\Users\cglab\Desktop\Match3\Train.py", line 124, in <module>
    main(ACTION_TYPE, MAP_PLAYER, MAX_PLAY_PER_STEP, MAX_STEP, N_TILE_TYPE, MAX_X, MAX_Y, MAX_MOVE_PER_PLAY, N_TARGET)
ValueError: 'p' must be 1-dimensional

问题代码 #

class CustomModel(tf.keras.Model):
    def __init__(self, num_hidden, max_x, max_y, n_tile_type):
        super(CustomModel, self).__init__()
        self.common = layers.Dense(num_hidden, activation="relu")
        self.x_probs = layers.Dense(max_x, activation="softmax")
        self.y_probs = layers.Dense(max_y, activation="softmax")
        self.tile_prob = layers.Dense(n_tile_type, activation="softmax")
        self.critic = layers.Dense(1)

    def call(self, inputs):
        common = self.common(inputs)
        x_probs = self.x_probs(common)
        y_probs = self.y_probs(common)
        tile_prob = self.tile_prob(common)
        critic = self.critic(common)

        return [x_probs, y_probs, tile_prob, critic]

max_x = 10
max_y = 10
max_tile_type = 5
model = CustomModel(512, max_x, max_y, max_tile_type)

state = np.full((max_y, max_x), -1)
state = tf.convert_to_tensor(state, dtype=tf.float32)
state = tf.expand_dims(state, 0)

x_probs, y_probs, tile_probs, critic_value = model(state)
x = np.random.choice(max_x, p=np.squeeze(x_probs)) #error

问题原因 #

输入的state形状是(1,10,10)（三维），Keras的Dense层默认会对输入的最后一维进行变换，保留前面的维度。因此：

1. 输入经过self.common（Dense(512)）后，输出形状变为(1,10,512)
2. 再经过self.x_probs（Dense(10, softmax)）后，输出形状变为(1,10,10)
3. 即使使用np.squeeze(x_probs)，也只能去掉第一个维度，得到(10,10)的二维数组，不符合np.random.choice对p参数必须是一维数组的要求。

修改方案 #

在模型中加入Flatten层，将三维输入展平为二维（batch_size, feature_num），确保后续Dense层输出为正确的二维（batch_size, action_num）。

修改后的模型代码：

class CustomModel(tf.keras.Model):
    def __init__(self, num_hidden, max_x, max_y, n_tile_type):
        super(CustomModel, self).__init__()
        self.flatten = layers.Flatten()  # 新增Flatten层
        self.common = layers.Dense(num_hidden, activation="relu")
        self.x_probs = layers.Dense(max_x, activation="softmax")
        self.y_probs = layers.Dense(max_y, activation="softmax")
        self.tile_prob = layers.Dense(n_tile_type, activation="softmax")
        self.critic = layers.Dense(1)

    def call(self, inputs):
        # 先展平输入
        flattened_inputs = self.flatten(inputs)
        common = self.common(flattened_inputs)
        x_probs = self.x_probs(common)
        y_probs = self.y_probs(common)
        tile_prob = self.tile_prob(common)
        critic = self.critic(common)

        return [x_probs, y_probs, tile_prob, critic]

修改后，x_probs的形状会变成(1,10)，使用np.squeeze(x_probs)即可得到(10,)的一维数组，满足np.random.choice的参数要求，原报错代码行可正常运行。

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