[Ray RLlib] space_utils.batch()에서 np.stack 대신 사전 할당 배열로 연결 속도 개선

PR 링크: ray-project/ray#61256 상태: Merged | 변경: +17 / -4

들어가며

Ray RLlib의 space_utils.batch() 함수는 환경에서 수집한 관측값(observations)이나 기타 데이터를 리스트에서 배치로 합칩니다. 기존 구현은 np.stack을 사용했는데, 이 함수는 내부적으로 여러 단계의 배열 연결 작업을 수행합니다. 모든 배열이 동일한 shape을 가지는 것이 보장되므로, np.empty로 사전 할당 후 직접 복사하는 것이 더 효율적입니다.

핵심 코드 분석

Before: np.stack 사용

np_func = np.concatenate if individual_items_already_have_batch_dim else np.stack
ret = tree.map_structure(
    lambda *s: np.ascontiguousarray(np_func(s, axis=0)), *list_of_structs
)

After: 사전 할당 배열로 직접 복사

if individual_items_already_have_batch_dim:
    ret = tree.map_structure(lambda *s: np.concatenate(s, axis=0), *list_of_structs)
else:
    n = len(list_of_structs)

    def fast_stack(*s):
        s0 = s[0]
        if not isinstance(s0, np.ndarray):
            return np.array(s)
        out = np.empty((n, *s0.shape), dtype=s0.dtype)  # 사전 할당
        for i in range(n):
            out[i] = s[i]  # 직접 복사
        return out

    ret = tree.map_structure(fast_stack, *list_of_structs)

왜 이게 좋은가

1. 메모리 할당 최소화: np.stack은 내부적으로 np.expand_dims + np.concatenate를 수행하여 임시 배열을 생성하지만, np.empty + 직접 복사는 단일 할당만 수행한다.
2. np.ascontiguousarray 제거: 사전 할당된 배열은 이미 C-contiguous하므로 추가 변환이 불필요하다.
3. 분기 분리: individual_items_already_have_batch_dim 여부에 따른 분기를 명확히 분리하여 코드 가독성도 향상했다.

참고 자료

• NumPy np.stack vs np.empty 벤치마크
• Ray RLlib 문서