게임을 플레이하려면 보통 게임 엔진과 플레이어가 따로 있어야 한다. 그런데 'tinyworldformer'라는 오픈소스 프로젝트는 이 둘을 모두 신경망으로 대체하고, 그것도 5달러짜리 마이크로컨트롤러인 라즈베리 파이 피코 2 W 위에서 동시에 돌린다. 두 개의 작은 네트워크가 쓰이는데, 첫 번째는 '월드 모델'로 현재 화면과 입력된 행동을 받아 다음 화면과 보상을 예측한다. 두 번째는 '액터'로 그렇게 생성된 화면을 보고 다음 행동을 고른다. 칩 위에서 두 네트워크는 실제 게임을 전혀 보지 않는다. 액터는 월드 모델이 꿈처럼 그려낸 화면만 보고, 월드 모델은 액터가 고른 행동만 받는다. 말 그대로 신경망이 꾼 꿈 안에서 또 다른 신경망이 게임을 하는 셈이다.
가중치를 재사용하는 '루프 트랜스포머'
이 프로젝트의 출발점은 '루프 트랜스포머(looped transformer)'에 대한 호기심이었다. 보통 트랜스포머는 서로 다른 층을 여러 개 쌓지만, 루프 방식은 동일한 가중치를 가진 하나의 블록을 여러 번 반복 적용한다. 가중치를 재사용하므로 메모리 소모가 적고, 그래서 월드 모델을 마이크로컨트롤러에 올릴 만큼 작게 만들 수 있지 않을까 하는 발상이었다. 피코 2 W를 고른 이유는 값이 싸면서도 부동소수점 연산장치(FPU)를 갖췄기 때문이고, 퐁·브레이크아웃·팩맨·플래피버드·둠라이트 같은 작은 게임을 택한 이유는 실제 규칙을 numpy로 직접 짜서 생성된 프레임을 한 칸 한 칸 대조할 수 있었기 때문이다.
모델 구조는 단출하다. 화면을 2×2(브레이크아웃)에서 6×6(팩맨)까지의 정사각형 패치로 잘라 게임당 36~80개의 토큰으로 만든다. 그 위에 어텐션과 ReLU MLP로 이뤄진 트랜스포머 블록 하나를, 매번 입력을 더해주며 1~4회 반복 적용한다. 반복 횟수가 많을수록 정확해지지만 느려지는 교환 관계가 있어, 피코에서는 적당한 프레임 속도를 위해 2회 루프를 쓴다. 공개된 둠라이트 예시를 보면 한 번 통과했을 때는 있지도 않은 적을 그려 넣지만, 반복할수록 틀린 셀이 92→35→12→0개로 줄어든다.
정확도만으로는 부족하다 — '꿈을 검증하기'
흥미로운 지점은 모델의 출력 점수를 실제 프레임으로 바꾸는 디코딩 규칙이 모델 자체만큼 중요했다는 고백이다. 초기 브레이크아웃 모델은 공을 두 개 그리는 버그가 있었는데, 전역적으로 하나만 고르게 하는 규칙 하나를 넣자 '정확 프레임' 비율이 91.6%에서 99.7%로 뛰었다. 또한 녹화된 데이터에 대한 정확도만으로는 신뢰할 수 없다는 점도 짚는다. 에이전트는 데이터에 없던 상태로 꿈을 몰아가고, 그곳의 작은 오류를 파고들기 때문이다. 그래서 핵심 검증은 에이전트가 플레이하는 도중에 꿈 화면에서 게임 상태를 거꾸로 읽어내 실제 규칙을 적용한 결과와 비교하는 방식으로 이뤄진다.
정수 8비트까지 내려간 양자화
피코에 올리기 위한 핵심 기술은 int8 양자화다. 월드 모델의 모든 행렬은 출력 행마다 하나의 스케일을 갖는 int8이고, 활성값도 루프 인덱스마다 범위가 다르다는 점을 반영해 int8로 둔다. 곱셈-누산은 RP2350의 SMLAD 명령으로 한 번에 두 번씩 처리하고, 소프트맥스는 1024개짜리 지수 테이블로 근사한다. 반면 정규화·잔차·임베딩은 FPU를 활용해 float32로 남긴다. 학습 단계에서 동일한 반올림을 적용하는 양자화 인지 학습(QAT)을 쓰기 때문에, 배포되는 int8 모델이 곧 학습된 모델 그 자체다. 나아가 int8 가중치를 러스트 소스로 내보낸 뒤 cargo test로 파이썬 정수 명세와 비트 단위까지 일치하는지 확인한다. 펌웨어는 embassy 기반 no_std 러스트로, 실행 중 메모리 할당이 전혀 없이 가중치와 버퍼를 모두 정적 배열로 둔다.
온디바이스 학습도 눈여겨볼 만하다. 액터는 재생 버퍼도, 배치도, 타깃 네트워크도 없이 한 전이씩 학습하는 Stream AC(λ) 알고리즘(Elsayed, Vasan & Mahmood, 2024)을 쓴다. 보드는 처음부터 새 에이전트가 자기 꿈 속에서 스스로 학습하도록 부팅되고, 버튼을 누르면 사전학습된 액터로 전환하거나 사용자가 직접 플레이할 수도 있다. 보상 설계의 사례도 솔직하게 공유된다. 둠라이트에서 에이전트가 시작 모퉁이에 웅크려 사격만 하는 문제가 있었는데, 최근 자주 방문한 칸에 머무르면 벌점을 주는 idle-penalty를 넣자 순찰을 돌며 처치 수가 두 배가 됐다고 한다.
실무적 관점에서 이 프로젝트는 생성 모델과 강화학습 에이전트를 극단적으로 압축해 저가 MCU에 통합한 TinyML 엔지니어링의 교과서적 사례다. 양자화-러스트 내보내기 파이프라인은 numpy와 torch에만 의존하도록 분리돼 다른 프로젝트가 그대로 재사용할 수 있고, 동일한 엔진이 웹어셈블리로도 컴파일돼 같은 시드로 같은 에피소드를 프레임 단위까지 재현한다. 다만 한계도 분명하다. 대상은 수백 줄짜리 장난감 게임이고, 데스크톱에서 한 프레임을 0.7~1.5ms에 꿈꾸는 성능은 측정됐지만 정작 피코 자체의 프레임 속도는 아직 측정하지 않았다고 밝힌다. 가장 큰 메모리를 차지하는 임베딩 테이블도 여전히 float32로 남아 있어 추가 최적화 여지가 있다. 그럼에도 '엣지에서 돌아가는 월드 모델'이라는 발상이 실제 실리콘 위에서 어디까지 가능한지를 구체적 수치와 함께 보여준다는 점에서, 온디바이스 AI를 고민하는 개발자에게 충분히 들여다볼 가치가 있다.