Small RL Models at the Speed of Light — LeanRL

RL 학습 — 특히 PPO / SAC 같은 on-policy 또는 off-policy 학습 — 의 한 step 안에서 모델 forward 가 1ms 이하로 끝난다. 그 짧은 시간에 GPU launch overhead, environment step, replay sampling, gradient update 가 모두 끼어든다. 결과 — GPU 사용률이 5%~30% 사이.

강의가 던지는 두 질문.

1. RL 학습의 병목이 정확히 어디인가 — 작은 모델이라 compute 가 아니다. 그러면 launch overhead, host-device 전송, environment step 중 어디?
2. 그 병목을 어떻게 풀 수 있나 — vectorized env, cudagraph, torch.compile 의 조합으로.

본 노트는 transcript 실패로 — LeanRL repo (github.com/pytorch-labs/LeanRL) 의 README, 일반적 RL 인프라 지식 (CleanRL, Stable Baselines, IsaacGym 등), 그리고 PyTorch torch.compile / cudagraph 의 동작 원리로 재구성. speaker 정보가 누락되어 본문에 “speaker 의 인용” 형식 안 씀.

supervised learning 의 throughput 단위는 tokens/sec. RL 은 environment steps/sec. 정책 모델이 작아도 환경이 느리면 학습이 느리고, 환경이 빨라도 정책이 작아서 GPU 가 비면 학습이 느리다.

그래서 RL 의 GPU 활용률은 통상적으로 — single-env CPU 환경 (gym/MuJoCo) 에서 1~5%, vectorized CPU env (gym vector) 에서 5~15%, GPU env (IsaacGym) 에서 20~50%. LeanRL 의 목표는 마지막 카테고리에서 60~80% 까지 끌어올리는 것.

단일 환경(env) 대신 수천 개 환경을 동시에 step. 각 환경의 state 가 batch dim 으로 묶여서 — policy forward 가 batch=4096 의 큰 op 이 됨. compute-bound 영역으로 진입.

# LeanRL 의 일반 패턴 (의사코드)
import torch
import isaacgym as ig

env = ig.make_env("Cartpole", num_envs=4096,
                  device="cuda")
policy = MLP(obs_dim, act_dim).cuda()
policy = torch.compile(policy, mode="reduce-overhead")

for step in range(N):
    obs = env.get_obs()                  # (4096, obs_dim) on cuda
    action = policy(obs)                 # (4096, act_dim) on cuda
    env.step(action)                     # env 시뮬레이션 GPU 위
    rewards = env.get_rewards()          # (4096,) on cuda
    # gradient update — 같은 GPU 위에서

환경 수가 ~4096 정도 되면 — 모델 forward 가 batch 작아서 빠르고 + env step 이 GPU 에서 빠르게 — wall-clock 학습 시간이 single-env 의 100배 이상 빨라지는 사례도 있음. 강의에서 정확한 숫자를 언급했을 가능성 높음 (원본 영상 확인 필요).

큰 RL 시스템 (IMPALA, SEED, A3C) 의 표준 구조 — actor 는 환경 위에서 trajectory 수집, learner 는 gradient update. 둘이 비동기로 돌아 — 한쪽이 다른쪽을 안 기다림. LeanRL 은 작은 모델 RL 에서도 같은 패턴을 쓴다.

2-layer MLP 의 forward 도 — PyTorch 안에서는 50~100 개의 작은 op 으로 분해됨. linear, gelu, dropout, residual 등. 각자 GPU kernel launch 가 일어나고, 각 launch overhead 가 ~10μs. 모델 자체는 50μs 인데 launch overhead 가 500μs.

LeanRL 의 디자인 철학은 단순함. JAX 나 Triton 으로 새 framework 를 짜는 게 아니라 — PyTorch 의 표준 도구 (compile, cudagraph) 를 “알맞게” 쓰는 것만으로 큰 효과. 사용자가 새 언어 안 배워도 됨.

정확한 벤치마크 숫자는 LeanRL repo 에 있지만 — 일반적 패턴은 “같은 알고리즘 (PPO 등) 을 같은 환경 (Cartpole, Humanoid) 에서 같은 hyperparam 으로 학습” 의 wall-clock 비교.

강의에서 “speed of light” 라는 표현은 — RL 학습이 GPU 가 native 로 줄 수 있는 throughput 의 최대치에 도달한다는 의미. 그 자리에 도달하면 다음 병목은 알고리즘의 sample efficiency 가 됨 — “환경을 100배 빠르게 굴려도 학습 step 수가 같으면 똑같이 학습이 오래 걸린다”.

Stable Baselines 3 는 RL 의 사실상 표준 — 알고리즘 구현이 검증되고, 사용자 base 가 큼. 하지만 GPU 활용 측면에서 — design 자체가 CPU env 시대에 맞춰져 있어서 새 도구 (compile, cudagraph) 가 잘 안 끼워짐.

LeanRL 의 위치는 — “SB3 를 대체” 가 아니라 “GPU env 시대의 best practice 를 합친 reference”. 기존 SB3 사용자는 그대로 유지, 새 시작은 LeanRL — 의 분기.

LeanRL 의 현재 sweet spot 은 single-GPU + small policy + GPU env. 이 자리를 넘어가면 — 더 큰 정책 (LLM-RL), multi-GPU (DDP), multi-task (population-based) 같은 새 도전.

1. LLM-RL (RLHF, RLHF-PPO, GRPO) — 정책이 7B+ LLM. forward 가 100ms+ 라 launch overhead 무의미. cudagraph 효과 작음. compile 만으로 충분. 다른 차원의 도전 (memory, kv cache).
2. multi-GPU PPO — actor 들을 각 GPU 에 분산. learner 가 한 곳에 모음. NCCL 통신 패턴 새로 디자인.
3. population-based training — 여러 hyperparam set 의 정책을 동시 학습. 각 정책이 독립이라 GPU 활용 100% 가능.
4. self-play — 정책이 자신을 상대로 학습. actor pool 의 latest weight vs old weight 의 동시 평가.
5. continuous learning — 한 정책이 task 를 갈아타며 학습. catastrophic forgetting 의 인프라 차원 답.

• speaker 정보 누락 — 강의 메타에 speaker 가 없음. LeanRL 은 PyTorch Labs 의 결과물이므로 — Vincent Moens, Albert Bou 등 contributor 중 한 명일 가능성. 원본 영상 확인 필요.
• 벤치마크 정확한 숫자 — “30~100배” 같은 추정값. 강의에서 정확한 비교 표가 등장했을 것이지만 본 노트는 일반적 패턴으로 재구성.
• 지원 알고리즘 목록 — PPO, SAC 외에 어디까지 LeanRL 이 구현했는지 — repo 에서 직접 확인 필요. DDPG, TD3, A2C 가능성.
• compile mode 의 정확한 선택 — “reduce-overhead” vs “max-autotune” 중 어느 쪽이 RL 워크로드에 권장인지.
• actor-learner 의 구현 디테일 — separate process 인지 same-process 별도 stream 인지. 두 패턴의 trade-off.
• replay buffer 의 GPU 구현 — ring buffer 의 정확한 패턴. lock 처리.
• algorithmic 보정 (V-trace 등) — actor-learner 의 staleness 를 어떻게 처리하는지 — algorithm 쪽의 변경 여부.
• JAX-RL 과의 직접 비교 — Brax + JAX 의 throughput 과 LeanRL 의 throughput — 비슷한 환경에서 어느 쪽이 빠른지.