ProRL Agent: 다중 턴 LLM 에이전트 강화학습을 위한 서비스형 롤아웃 인프라

복잡한 다중 턴 LLM 에이전트 훈련 시 발생하는 롤아웃 병목 현상을 해결하기 위해 훈련과 실행 환경을 분리한 새로운 아키텍처를 제시합니다. 이를 통해 GPU 자원 낭비를 줄이고, 루트 권한이 없는 고성능 컴퓨팅(HPC) 환경에서도 대규모 에이전트 학습이 가능해집니다.

강화학습에서 에이전트가 환경과 상호작용하며 데이터를 쌓는 과정을 '롤아웃'이라고 한다. 기존 시스템은 이 롤아웃 로직이 훈련 코드 내부에 섞여 있어, 파일 입출력이 많은 환경 상호작용과 연산량이 많은 GPU 학습이 자원을 서로 점유하며 전체 속도를 늦추는 문제가 있었다.

ProRL Agent는 이를 '서비스형 롤아웃'으로 전환하여 해결한다. 훈련 서버는 단순히 작업 요청만 보내고, 실제 환경 구축과 도구 실행은 별도의 서버가 전담한다. 이는 마치 식당에서 요리사(훈련)와 서빙(롤아웃)의 역할을 분리하여 회전율을 높이는 것과 같은 원리다.

특히 다중 턴 에이전트는 한 번의 데이터 수집에 수십 번의 도구 호출이 필요해 시간이 매우 오래 걸린다. ProRL Agent는 이 과정을 비동기 파이프라인으로 처리하고 여러 추론 서버를 동적으로 관리하여, 훈련 서버가 데이터가 도착할 때까지 기다리는 유휴 시간을 획기적으로 줄인다.

시스템은 Sandbox Environment, ProRL Agent Server, RL Trainer의 세 가지 계층으로 구성된다. Trainer는 HTTP API를 통해 작업을 제출하며, Server는 이를 수신하여 내부 큐에 할당한다.

롤아웃 과정은 INIT(샌드박스 생성), RUN(에이전트 루프), EVAL(보상 평가)의 3단계 비동기 파이프라인으로 처리된다. 각 단계는 독립된 워커 풀을 가지며, [이전 단계 완료 데이터 입력 → 현재 단계 워커 처리 → 다음 단계 큐 출력] 순으로 연산이 흐르도록 설계되어 병렬성을 극대화한다.

LLM 추론 백엔드 관리를 위해 Min-Heap 기반 로드 밸런싱을 적용한다. [등록된 모든 백엔드의 현재 작업 수 입력 → 최소 작업 수 백엔드 선택 → 작업 할당 및 카운터 갱신] 과정을 통해 추론 부하를 균등하게 분산하고 프리픽스 캐시 효율을 높인다.

SWE-Bench Verified 벤치마크에서 Qwen3 모델군(4B, 8B, 14B)을 대상으로 실험한 결과, 모든 규모에서 성능 향상을 확인했다. 특히 8B 모델은 기존 SkyRL-Agent 대비 약 2배의 성능 향상을 기록하며 인프라의 효과를 입증했다.

STEM, MATH, Coding 도메인에서도 안정적인 학습 성능을 보였다. STEM 에이전트의 경우 60단계 학습 후 평균 보상이 0.2에서 0.65로 상승했으며, 수학 및 코드 생성 작업에서도 학습이 진행됨에 따라 Pass@1 지표가 꾸준히 개선되었다.

시스템 분석 결과, 제안된 비동기 파이프라인과 도구 실행 최적화를 통해 GPU 활용률을 78%까지 달성했다. 또한 노드 수가 증가함에 따라 처리량이 거의 선형적으로 증가하는 우수한 확장성을 확인했다.

ProRL Agent Server는 FastAPI 기반으로 구축되었으며, Singularity 컨테이너 런타임을 사용하여 루트 권한 없이도 격리된 환경을 제공한다. 각 컨테이너 인스턴스에는 고유한 루프백 IP가 할당되어 다수의 에이전트가 동일 노드에서 포트 충돌 없이 실행될 수 있다.

도구 실행 효율화를 위해 ptyprocess 기반의 직접 의사 터미널을 도입하여 Bash 명령 실행 지연을 줄였다. 또한 IPython 커널에 인프로세스 API로 직접 연결하여 기존 Jupyter 게이트웨이 방식에서 발생하던 네트워크 오버헤드를 제거했다.

데이터 일관성을 위해 Token-in/Token-out 통신 방식을 채택했다. [Token ID 시퀀스 입력 → LLM 추론 및 로그 확률 계산 → 결과 Token ID 출력] 과정을 통해 텍스트 변환 시 발생하는 토큰 불일치 문제를 방지하고 훈련 안정성을 높였다.

Dynamic Sampling Policy Optimization (DAPO) 알고리즘을 지원하며, 정보가 없는 샘플을 필터링하는 과정에서 발생하는 유휴 시간을 줄이기 위해 비동기 보충(Replenishment) 메커니즘을 구현하여 최대 처리량을 유지한다.

더 다양한 환경 지원과 클러스터 규모에서의 견고성 향상이 향후 과제로 명시됨.