본문으로 건너뛰기

[Ray RLlib] SingleAgentEnvRunner의 validate 호출 위치 최적화로 3.1배 속도 향상

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

들어가며

Ray RLlib의 SingleAgentEnvRunner에서 에피소드 데이터를 수집하는 핵심 루프에서 validate() 메서드가 전체 샘플링 시간의 13.86%를 차지하고 있었습니다. 원인을 추적해보니, validate()__init__, add_reset_data, 그리고 **매 스텝마다 호출되는 add_step_data**에서 모두 호출되고 있었습니다. 환경의 매 타임스텝마다 전체 에피소드를 검증하는 것은 과도한 방어적 프로그래밍입니다.

핵심 코드 분석

validate 호출 위치 변경

Before (add_env_step 내부):

def add_env_step(self, ...):
    # ... 관찰, 행동, 보상 등을 에피소드에 추가 ...

    # Validate our data.
    self.validate()

    # Step time stats.
    self._last_step_time = time.perf_counter()

After (to_numpy와 concat_episode로 이동):

def add_env_step(self, ...):
    # ... 관찰, 행동, 보상 등을 에피소드에 추가 ...
    # validate 호출 제거

    # Step time stats.
    self._last_step_time = time.perf_counter()

def to_numpy(self) -> "SingleAgentEpisode":
    # Check that the episode data is correct
    self.validate()
    # ... 나머지 numpy 변환 로직 ...

def concat_episode(self, other: "SingleAgentEpisode") -> None:
    # Validate both this and the other episode
    self.validate()
    other.validate()
    # ... 나머지 연결 로직 ...

왜 이게 좋은가

  1. 3.1배 속도 향상: add_step_data의 누적 실행 시간이 16.7초(13.86%)에서 5.43초(4.52%)로 감소했습니다.
  2. 검증 수준 유지: to_numpy()concat_episode()는 에피소드가 완료된 후 호출되므로, 최종 데이터의 무결성은 여전히 보장됩니다.
  3. 핫 패스 최적화: 환경 스텝은 가장 자주 실행되는 코드 경로입니다. 여기서의 오버헤드 감소는 전체 학습 시간에 직접적으로 영향을 줍니다.
  4. 코드 변경 최소화: 4줄 추가, 4줄 삭제만으로 의미 있는 성능 개선을 달성했습니다.

참고 자료

댓글

관련 포스트

PR Analysis 의 다른글