如何从3D NumPy数组的最后一维中随机选取一个元素？

假设你的self.transition[s_dyna,a_dyna,:]是形状为(100,)的一维数组（当s_dyna和a_dyna都是单个整数索引时），直接用numpy.random.choice就能从最后一维（也就是这个一维数组）里随机选一个元素赋值给s_prime：

import numpy as np

# 先获取目标切片
target_slice = self.transition[s_dyna, a_dyna, :]
# 从切片中随机选取一个元素
s_prime = np.random.choice(target_slice)

如果你的切片结果是多维数组（比如s_dyna/a_dyna是多个索引，得到形状为(N, M, 100)的数组），需要对每个最后一维的100个元素随机选一个，生成形状为(N, M)的s_prime数组，可以用高效的索引方式避免循环：

# 获取切片的形状信息
slice_shape = self.transition[s_dyna, a_dyna, :].shape
# 生成对应维度的随机索引（每个位置选最后一维的一个随机下标）
rand_indices = np.random.randint(0, slice_shape[-1], size=slice_shape[:-1])
# 用NumPy高级索引取出对应元素
s_prime = self.transition[s_dyna, a_dyna, :][np.indices(slice_shape[:-1]) + (rand_indices,)]

如果追求代码直观性，也可以用np.apply_along_axis（效率略低于索引方式）：

def pick_random_element(arr):
    return np.random.choice(arr)

s_prime = np.apply_along_axis(pick_random_element, axis=-1, arr=self.transition[s_dyna, a_dyna, :])

补充说明 #

• np.random.choice默认是均匀随机选取，若需要带权重的随机选择，可传入p参数指定每个元素的概率。
• 若使用NumPy 1.17+版本，推荐用np.random.Generator的choice方法，性能更稳定：

  rng = np.random.default_rng()
  s_prime = rng.choice(target_slice)
  

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