如何用np.random.choice高效批量选择对应概率分布的动作？

给定动作数组与多组动作概率分布，需要无循环高效生成每组概率对应的动作，替代Python循环以提升大量episode生成时的速度。

实现代码 #

import numpy as np

actions = np.array([0, 1, 2, 3])
probs = np.array([[0.5, 0.4, 0.05, 0.05], [0.05, 0.05, 0.1, 0.8]])

# 计算每组概率的累积分布
cumulative_probs = np.cumsum(probs, axis=1)
# 生成与概率组数匹配的随机数
random_values = np.random.rand(probs.shape[0], 1)
# 找到每个随机数对应的动作索引
action_indices = np.argmax(cumulative_probs > random_values, axis=1)
# 映射到动作数组得到最终结果
output = actions[action_indices].tolist()

print(output)  # 示例输出：[0, 3]

效率说明 #

• 全程使用numpy矢量化操作，所有计算在底层C层面完成，彻底避免Python循环的性能开销；
• 处理十万级甚至百万级别的概率组时，速度比循环实现快100~1000倍，完全适配大量episode生成的场景。

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