시리즈: Learning RL and JAX in Public - from zero to DeepMind.

Day 6은 REINFORCE가 작동하는 것으로 끝났습니다. 에이전트는 Q-value 대신 확률을 사용하여 그리드월드 (gridworld)를 탐색하는 법을 배웠습니다. 깔끔한 아이디어였고, 깔끔한 코드였습니다.

하지만 노이즈가 심했습니다. 학습이 느렸습니다. 어떤 에피소드에서는 에이전트가 운이 좋았고, 어떤 에피소드에서는 운이 나빴으며, 네트워크는 그 차이로 인해 계속 혼란을 겪었습니다. 신호는 존재했지만 너무 많은 분산 (variance) 속에 파묻혀 있었습니다.

Day 7은 그 해결책입니다.

이해를 도운 농구 비유

농구 선수가 한 경기에서 15점을 득점했다고 상상해 보세요. 그것이 좋은 걸까요, 나쁜 걸까요?

그것은 전적으로 그 선수에게 무엇이 정상인지에 달려 있습니다. 평소에 10점을 득점한다면 훌륭한 경기입니다. 만약 평소에 25점을 득점한다면 나쁜 경기입니다. 숫자 그 자체만으로는 아무것도 알려주지 않습니다. 맥락 (context)이 필요합니다.

REINFORCE는 맥락이 없습니다.

이 단 한 번의 뺄셈이 엄청난 양의 노이즈 (noise)를 제거합니다. "기대보다 더 나은" 것이 이제 단순히 큰 숫자가 아니라 의미 있는 신호 (signal)가 되기 때문에, 운 좋게 얻은 에피소드 (lucky episode) 문제도 사라집니다.

업데이트 방식

두 네트워크 모두 매 에피소드마다 업데이트됩니다. 먼저 크리틱 (critic)이 업데이트되어 (추정 능력을 향상시키기 위해), 그다음 액터 (actor)가 업데이트됩니다 (최신 추정치를 사용하여):

# critic: 리턴 (returns)을 정확하게 예측하도록 학습
critic_loss = (actual_return - critic(state)) ** 2

...

크리틱은 단순한 회귀 (regression)처럼 학습합니다. 액터는 REINFORCE와 유사하게 학습하되 더 깨끗한 신호를 사용합니다. 이들은 시간이 지남에 따라 함께 발전합니다.

훈련 출력에서 확인한 내용

코드는 마지막에 크리틱의 가치 추정치 (value estimates)를 그리드 형태로 출력합니다. 훈련 후에는 대략 다음과 같이 보였습니다:

목표 지점 근처 상태: 높은 가치 (~0.7 ~ 0.9)
구멍 근처 상태: 낮은 가치 (음수)
안전한 중간 상태: 중간 가치

크리틱은 액터가 플레이하는 것을 지켜보는 것만으로 그리드월드 (gridworld)의 레이아웃을 학습했습니다. 아무도 구멍이 어디에 있는지 알려주지 않았습니다. 경험을 통해 스스로 알아낸 것입니다.

그 과정을 지켜보는 것은 정말 멋진 경험이었습니다.

이것이 모든 것의 기반이 되는 아키텍처 (architecture)인 이유

제가 계속해서 PPO와 GRPO를 언급하고 있습니다. 그 직접적인 연결 고리는 다음과 같습니다:

• PPO (ChatGPT 학습에 사용됨): 액터-크리틱 (actor-critic) 방식입니다. 업데이트가 과도하게 발생하는 것을 방지하기 위해 클립 (clip)을 추가합니다.

• A3C (DeepMind, 2016): 액터-크리틱 방식입니다. 더 다양한 경험을 수집하기 위해 여러 에이전트 (agents)를 병렬로 실행합니다.

• GRPO (DeepSeek-R1): 액터-크리틱의 변형입니다. 단일 크리틱을 여러 에피소드에 걸친 그룹 비교 (group comparisons)로 대체합니다. "내 추정치보다 더 나은가?" 대신 "이 배치 (batch) 내의 다른 에피소드들보다 더 나은가?"라고 묻습니다.

이 시리즈의 앞부분에서 제가 만든 forge 프로젝트는 GRPO를 구현합니다. 프로젝트를 만들 당시에는 보상 정규화 (reward normalization)가 왜 그런 방식으로 작동하는지 완전히 이해하지 못했습니다. 이제는 압니다. 그룹 내에서 어드밴티지 (advantage)를 계산하는 것입니다.

액터-크리틱을 이해하는 것이 바로 빠져 있던 마지막 조각이었습니다.

한눈에 보는 계보:

Actor-Critic
    |
    ├── A3C  (DeepMind, 2016)   - 병렬 액터 (parallel actors)
...

모든 모델의 공통점은 다음과 같습니다: 행동을 선택하는 액터 (actor) + 품질을 추정하는 무언가. 차이점은 어드밴티지 (advantage)를 계산하고 학습을 안정화하는 방식에 있습니다.

Day 8: DeepMind의 신경망 라이브러리인 Haiku로 넘어갑니다. 알고리즘은 동일하지만, 코드는 훨씬 더 깔끔합니다. 바로 이 지점부터 코드베이스가 실제 연구용 코드처럼 보이기 시작합니다.

이 시리즈의 모든 코드는 일자별로 정리되어 제 GitHub에 있습니다: https://github.com/MadhumithaKolkar/jax-rl-lab

모두 즐거운 학습 되시길 바랍니다!

~ Madhumitha Kolkar (index_0)