강화학습 라이브러리 LeanRL 소개
LeanRL은 PyTorch 기반의 강화 학습(Reinforcement Learning, RL) 알고리즘을 위한 경량 라이브러리로, CleanRL에서 파생(forked)되었습니다. LeanRL의 주된 목적은 RL 사용자들에게 학습 시간을 절반 이상 단축할 수 있는 최적화 기법을 제공하는 것입니다. 이 라이브러리는 torch.compile과 cudagraphs 같은 PyTorch 2.x의 기능을 적극 활용하여 RL 스크립트의 성능을 극대화합니다. 이를 통해 사용자는 복잡한 설정 없이도 최적화된 성능을 경험할 수 있습니다.
LeanRL은 CleanRL의 원래 정신을 계승하여, 각 알고리즘의 모든 세부 사항을 단일 파일에 포함하는 방식으로 구현됩니다. 주요 기능은 다음과 같습니다:
- 단일 파일 구현: 각 강화 학습 알고리즘 변형에 대한 모든 세부 사항이 독립적인 단일 파일에 포함되어 있어 코드 관리가 간편합니다.
- 최적화된 성능: PyTorch 2의 새로운 기능을 활용한 경량화된 버전으로, 코드 실행을 최적화하고 데이터를 효율적으로 처리합니다.
torch.compile: 연산자 결합(fusion)과 같은 최적화 기법으로 오버헤드를 줄입니다.cudagraphs: CUDA에서 실행되는 작업을 캡처하고 나중에 동일한 그래프를 재실행하여 오버헤드를 제거합니다.tensordict: CUDA에서 데이터를 빠르게 복사하고 함수 호출을 효율화하며 목표 파라미터 업데이트를 빠르게 처리합니다.torch.vmap: Q-값 네트워크 실행을 벡터화하여 처리 속도를 높입니다.
LeanRL의 속도 향상 기법
LeanRL은 성능 향상을 위해 PyTorch 2.x에서 도입된 기능들을 활용합니다. 이 중에서 가장 주목할 만한 기법은 다음과 같습니다:
torch.compile()
torch.compile은 PyTorch 2.0에서 도입된 기능으로, 파이썬 코드를 저수준의 C++ 연산으로 변환하여 학습과 추론의 속도를 향상시킵니다. 이를 통해 파이썬 인터프리터와 C++ 런타임 간의 경계를 넘나드는 비용을 줄일 수 있습니다. PyTorch의 기본 실행 모드와 달리, torch.compile은 단일 C++ 실행 파일을 생성하여 이러한 경계를 최소화합니다. 또한, 그래프 중단(graph breaks)을 잘 처리하여 대부분의 파이썬 코드를 컴파일할 수 있도록 설계되었습니다.
cudagraphs
강화 학습은 일반적으로 얕은 네트워크를 사용하는 경우가 많아 CPU 오버헤드가 많이 발생합니다. torch.compile은 컴파일된 코드의 실행 시 CPU 오버헤드가 발생할 수 있는데, 이는 연산자 결합의 이점을 상쇄할 수 있습니다. 이를 해결하기 위해 PyTorch는 cudagraphs 기능을 제공합니다. cudagraphs는 CUDA에서 실행되는 작업을 캡처한 후, 이후 동일한 작업을 반복하여 실행하는 방식을 사용합니다. 이를 통해 CPU와의 상호작용을 줄이고 속도를 크게 향상시킬 수 있습니다.
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
z_buffer = func(x_buffer, y_buffer)
x_buffer.copy_(x)
y_buffer.copy_(y)
graph.replay()
z = z_buffer.clone()
라이선스 정보
LeanRL은 MIT 라이선스 하에 공개되었습니다. 자세한 내용은 GitHub 저장소의 LICENSE 파일에서 확인할 수 있습니다.
LeanRL GitHub 저장소
CleanRL GitHub 저장소
이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. 
파이토치 한국 사용자 모임
이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일
로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)
아래
쪽에 좋아요
를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ 