[ray] Ray RLlib의 비동기 학습 성능 최적화: PULL 기반 EnvRunnerStateServer 도입

PR 링크: ray-project/ray#63849 상태: Merged | 변경: +523 / -82

들어가며

Ray RLlib의 비동기 알고리즘(APPO, IMPALA)은 학습 중인 가중치를 EnvRunner로 전달하기 위해 기존에 PUSH 모델을 사용해 왔습니다. 하지만 이 방식은 EnvRunner가 sample() 루프를 도는 동안 새로운 가중치 업데이트가 발생하면, 해당 업데이트를 무시하거나 강제로 덮어쓰는 과정에서 오프폴리시(off-policy) 지연이 발생하는 문제가 있었습니다. 본 PR은 이를 해결하기 위해 중앙 집중식 EnvRunnerStateServer를 도입하여, EnvRunner가 필요할 때 최신 가중치를 직접 PULL 해가는 방식으로 아키텍처를 개선했습니다.

코드 분석

1. AlgorithmConfig 설정 확장

rllib/algorithms/algorithm_config.py에서는 새로운 PULL 모델을 제어하기 위한 설정값이 추가되었습니다.

# Before
self.broadcast_env_runner_states = True

# After
self.use_env_runner_state_server = False
self.env_runner_state_server_max_concurrency = 16

2. EnvRunnerStateServer 도입 및 공유

rllib/algorithms/impala/impala.py에서는 EnvRunnerStateServer 액터를 생성하고, 이를 모든 EnvRunner와 공유하도록 설정합니다.

# After: 새로운 서버 액터 생성 및 공유
server_cls = ray.remote(
    num_cpus=0,
    max_restarts=-1,
    max_concurrency=self.config.env_runner_state_server_max_concurrency,
)(EnvRunnerStateServer)
self._env_runner_state_server = server_cls.remote()

def _share_state_server(env_runner, server=self._env_runner_state_server):
    env_runner._env_runner_state_server = server

self.env_runner_group.foreach_env_runner(func=_share_state_server, local_env_runner=True)

3. PUSH에서 PULL로의 로직 전환

기존에는 Algorithm이 모든 EnvRunner에 강제로 상태를 밀어넣었으나, 이제는 서버에 상태를 업데이트하고 EnvRunner가 sample() 호출 시점에 이를 가져갑니다.

# After: PULL 모델 적용
if self._env_runner_state_server is not None:
    env_runner_state = self.env_runner_group.get_merged_env_runner_state(...)
    self._env_runner_state_server.push.remote(env_runner_state)
else:
    self.env_runner_group.sync_env_runner_states(...)

왜 이게 좋은가

1. 오프폴리시 지연 감소: 실험 결과, diff_num_grad_updates_vs_sampler_policy 지표가 평균 20% 개선되었습니다. 이는 EnvRunner가 항상 가장 최신의 가중치를 사용하여 샘플링하기 때문입니다.
2. 처리량(Throughput) 유지: 가중치 전송을 ObjectRef 기반으로 최적화하여 네트워크 오버헤드를 최소화했습니다.
3. 회복 탄력성(Robustness): EnvRunnerStateServer가 실패하더라도 max_restarts=-1 설정을 통해 자동으로 복구되며, EnvRunner는 재시도 로직을 통해 학습 흐름을 끊지 않습니다.

교훈: 분산 시스템에서 상태 동기화가 병목이 될 때, PUSH(서버 주도) 방식보다는 PULL(클라이언트 주도) 방식이 시스템의 유연성과 최신성(freshness)을 보장하는 데 더 유리할 수 있습니다.

리뷰어 피드백 반영

리뷰 과정에서 env_to_module 파이프라인이 누락되었을 때의 예외 처리 문제가 제기되었으며, 이는 기존 sync_env_runner_states의 동작을 계승하되 안전성을 확보하는 방향으로 논의되었습니다. 또한, 테스트 코드에서 use_server 플래그에 따른 액터 생성 여부를 명확히 검증하도록 개선되었습니다.

참고 자료

• https://docs.ray.io/en/latest/ray-core/actors.html
• https://docs.ray.io/en/latest/rllib/rllib-training.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.