如何将元素为等长1D numpy.ndarray的1D数组转为2D numpy.ndarray？

我有一个形状为(256,)的1D numpy.ndarray，其元素均为长度相同的1D ndarray，无法直接转换为常规的2D numpy.ndarray。该数组示例如下：

[array([ 1.24836708, -0.00366597, -0.00617445,  0.00593824, -0.00499844,
        -0.008723  ,  0.0017313 ,  0.0032415 , -0.00921341, -0.08843152,
        -0.7156603 , -1.36149851,  1.03884429, -1.07907692, -1.08093478,
         0.40311199,  1.69189478])
 array([ 1.25213268e+00,  5.91834461e-06, -3.67174629e-03,  4.31587243e-03,
         2.05513245e-03,  4.03030160e-03, -3.70626458e-03, -3.19801204e-03,
        -4.03959450e-03, -3.16571953e-04, -3.36570170e-03,  4.23967454e-03,
        -4.80364918e-03, -5.40294083e-05, -7.57197110e-04,  8.55342504e-04,
         8.52148544e-06])
 array([  1.22506234,  -0.6000973 ,  -0.65958838,   0.01343306,
         -0.28506279,  -0.64566119,  -0.02085904,  -0.18191512,
         -2.45525694,  -1.91504442, -10.        , -10.        ,
          0.28079188,   2.70344181, -10.        ,  -4.4974326 ,
          4.82187824])
 ...

相关代码片段如下：

def sample(self, batch_size, **kwargs):
    """
    Sample a batch of experiences.
    Args:
        batch_size: int
            How many experience tuples to sample as a batch.
    Returns:
        state_batch: np.ndarray
            batch of current observed states
        action_batch: np.ndarray
            batch of actions executed given current states
        reward_batch: np.ndarray
            rewards received as results of executing action_batch
        next_state_batch: np.ndarray
            next state observed after executing action_batch
        done_mask: np.ndarray
            done_mask[i] = 1 if executing act_batch[i] resulted in the end of an episode
        discount: np.ndarray
            product of gammas for N-step returns
    """
    assert len(self._buffer) >= batch_size
    
    idx = np.random.randint(0, len(self._buffer), size=batch_size)
    batch = np.array(self._buffer)[idx]

    weights = np.zeros(batch_size)
    idxes = np.zeros(batch_size)

    return [np.array(batch[:, i]) for i in range(6)] + [weights, idxes]

其中self._buffer是一个列表，每个元素格式为[长度17的列表, 长度6的列表, int, 长度17的列表, bool, int]，问题中的数组即为np.array(batch[:, 0])。我尝试过np.array、array.reshape方法，但均无效，目前能想到的方法是先扁平化数组再reshape，请问是否有更简便的方法将其转换为如下所示的常规2D数组？

[ 1.28988438e+00 -3.28926461e-01 -3.16421640e-01  8.41842633e-03
 -6.04719630e-01 -3.21522510e-01  2.62200694e-02 -6.86004568e-01
 -2.75207792e+00  7.86307393e-01 -1.00000000e+01 -1.00000000e+01
  3.95934592e-01 -6.10089078e+00 -1.00000000e+01  5.53911043e-02
 -9.57694240e+00]
[ 1.20355496 -0.08770949  0.02464436 -0.11603001  1.19935722 -0.01665168
 -0.13457423 -0.03690821 -1.60970334  0.07964383 -5.48408598  0.10511611
 -8.75844077 -6.11991489 -2.52924686 -4.49214368  1.07501073]
[ 1.23668679e+00  1.71734326e-02  1.66840531e-02  5.26031800e-03
 -2.82087376e-02 -3.94339947e-03  2.50997621e-02  8.48180561e-02
  1.36369562e-01 -5.52168069e-01  6.43659328e-01  5.59083718e-01
  9.81923803e-03  1.00000000e+01  2.63192025e-01  8.68684414e-01
 -4.33114962e+00]
 ...

方法1：使用np.stack()或np.vstack()直接转换 #

针对已有的np.array(batch[:, 0])，直接用np.stack()或者np.vstack()就能把元素为一维数组的一维ndarray转换成二维数组：

# 两种方法任选其一
state_batch = np.stack(batch[:, 0])
# 或者
state_batch = np.vstack(batch[:, 0])

这两个方法会将形状为(N,)、每个元素是(M,)的ndarray转换为(N, M)的标准2D数组，完全符合需求。

方法2：在采样阶段直接生成二维数组 #

可以避免先将self._buffer转成np.array再切片，直接从列表中按索引提取元素后转换，这样生成的数组直接就是二维的：

def sample(self, batch_size, **kwargs):
    assert len(self._buffer) >= batch_size
    
    idx = np.random.randint(0, len(self._buffer), size=batch_size)
    # 直接从buffer中取元素，避免生成元素为数组的一维ndarray
    batch = [self._buffer[i] for i in idx]
    
    # 直接生成二维数组
    state_batch = np.array([item[0] for item in batch])
    action_batch = np.array([item[1] for item in batch])
    reward_batch = np.array([item[2] for item in batch])
    next_state_batch = np.array([item[3] for item in batch])
    done_mask = np.array([item[4] for item in batch])
    discount = np.array([item[5] for item in batch])
    
    weights = np.zeros(batch_size)
    idxes = np.zeros(batch_size)
    
    return [state_batch, action_batch, reward_batch, next_state_batch, done_mask, discount] + [weights, idxes]

这种方式从根源上避免了生成问题中的一维数组，代码更直接，也减少了不必要的中间转换。

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