솔직히 고백하겠습니다.
지난 주말, 저는 밥 먹는 시간도 아껴가며 **틱택토(Tic-Tac-Toe)** 게임을 구현하고 있었습니다.
"아니, CTO님이 최신 LLM도 아니고 고작 틱택토를요?"
라고 물으실 수 있습니다.
저도 처음엔 그렇게 생각했습니다.
하지만 강화학습(RL) 파이프라인의 효율성을 점검하다가, 문득 **기초 체력**의 중요성을 뼈저리게 느꼈습니다.
우리가 흔히 쓰는 PyTorch 기반의 루프는 직관적이지만, 대규모 시뮬레이션에서는 병목이 생길 수밖에 없습니다.
그래서 꺼내 든 카드가 바로 **Jax**입니다.
오늘은 Jax와 **PGX** 라이브러리를 활용해, 어떻게 강화학습 환경을 극적으로 최적화할 수 있는지 이야기해보려 합니다.
단순한 게임 구현기가 아닙니다.
이건 **속도와 리소스 관리**에 대한 이야기입니다.
먼저, 우리가 사용할 도구는 **PGX**입니다.
Jax로 작성된 게임 환경 라이브러리인데, 아주 흥미로운 구조를 가지고 있습니다.
핵심은 `State`라는 dataclass입니다.
여기서 가장 눈여겨볼 점은 **Observation(관측)** 데이터의 형태입니다.
보통 틱택토라 하면 3x3 배열을 떠올리시겠죠?
PGX는 `(3, 3, 2)` 형태의 boolean 배열을 사용합니다.
첫 번째 채널은 '내 말'의 위치, 두 번째 채널은 '상대 말'의 위치입니다.
왜 이렇게 복잡하게 할까요?
신경망이 현재 플레이어의 시점에서 판을 이해하기 훨씬 유리하기 때문입니다.
그리고 `current_player` 로직도 재미있습니다.
Player 0은 매 게임마다 랜덤으로 X(선공)나 O(후공)가 됩니다.
덕분에 신경망을 항상 Player 0에 맞춰 학습시켜도, 선공과 후공 상황을 모두 커버할 수 있습니다.
이런 디테일이 데이터 편향을 막아줍니다.
이제 모델을 만들어봅시다.
가볍게 **Flax**나 **NNX**를 사용해 **DQN(Deep Q-Network)**을 구축합니다.
틱택토 정도의 복잡도라면 거창한 트랜스포머는 필요 없습니다.
`Linear` 레이어 세 개면 충분합니다.
입력은 3x3x2 배열을 평탄화(Flatten)해서 넣습니다.
이때, 내 말은 1, 상대 말은 -1로 변환해주는 센스가 필요합니다.
출력은 9개의 칸에 대한 가치(Value)입니다.
여기서 **Masking** 기술이 들어갑니다.
이미 돌이 놓인 자리에 또 둘 수는 없으니까요.
`legal_action_mask`를 사용해 불가능한 수에는 `-inf` 값을 주어 선택되지 않게 막습니다.
자, 여기까지는 평범해 보입니다.
하지만 진짜 마법은 **실행** 단계에서 일어납니다.
보통의 파이썬 루프라면 게임을 하나씩 순차적으로 돌렸겠지만, Jax는 다릅니다.
`jax.vmap`
이 함수 하나가 게임의 판도를 바꿉니다.
단일 게임 로직을 작성해두고 `vmap`을 씌우면, 수천 개의 게임을 **병렬**로 돌려버립니다.
GPU 하나에서 수천 개의 틱택토 게임이 동시에 진행되는 겁니다.
이게 왜 중요할까요?
데이터 수집 속도가 비교가 안 되게 빨라집니다.
소스 코드 작성자가 테스트한 결과, 노트북 환경에서도 **단 15초** 만에 완벽한 플레이를 학습했다고 합니다.
기존 방식이었다면 커피 한 잔 마시고 와야 끝났을 작업입니다.
이 속도 차이는 비즈니스 관점에서 엄청난 **ROI(투자 대비 효과)** 상승을 의미합니다.
실험 주기가 짧아지면, 그만큼 더 많은 가설을 검증할 수 있으니까요.
물론, Jax의 함수형 프로그래밍 스타일이 처음엔 낯설 수 있습니다.
상태(State)를 직접 변경하는 게 아니라, 항상 새로운 상태를 반환받아야 하니까요.
하지만 이 **불변성(Immutability)** 덕분에 부작용 없는 깔끔한 병렬 처리가 가능해집니다.
이번 주말, 여러분도 거대한 모델만 바라보지 말고 가벼운 환경에서 **속도의 미학**을 체험해보시는 건 어떨까요?
작은 틱택토 게임 하나에서, 우리 시스템 전체를 관통하는 최적화의 힌트를 얻으실 수 있을 겁니다.
