라즈베리파이 피코에서 신경망 게임 엔진 구현
원제 A tiny world model running games on a Raspberry Pi Pico 2 W
148 포인트댓글 8
Key Point
라즈베리파이 같은 극도로 제한된 임베디드 환경에서 신경망으로 게임을 생성하고 플레이하는 새로운 방식이 작동한다는 것을 보여주며, 마이크로컨트롤러에서 머신러닝 모델을 실행하는 방법의 한계를 재정의한다.
핵심 요약
- 임베디드 마이크로컨트롤러인 라즈베리파이 피코 2 W에서 신경망 기반 게임 엔진을 구현했다. 세계 모델 네트워크가 현재 프레임과 액션에서 다음 프레임과 보상을 예측하고, 액터 네트워크가 그 생성된 프레임에서 액션을 선택한다.
- 반복 트랜스포머(looped transformer) 구조를 사용해 같은 가중치를 1~4회 반복 적용함으로써 메모리 효율을 높였다. 26k~131k 매개변수의 세계 모델과 9k~16k 매개변수의 액터 네트워크를 탑재했다.
- 브레이크아웃(정확도 99.99%), 플래피 버드(100%), 팩맨 라이트(86.4%) 등 5개 게임을 구현했으며, 각 게임은 토큰 기반으로 분할된 프레임을 처리한다. 세계 모델은 배경 셀은 독립적으로 예측하되, 공이나 플레이어 같은 단일체는 한 번만 배치되도록 디코딩 규칙을 적용했다.
- 모든 가중치를 int8 양자화로 압축했다. 행렬당 출력 행마다 스케일 팩터 하나씩만 저장하고, 활성화도 int8으로 처리하되 루프 인덱스마다 별도의 스케일을 사용했다. 훈련 중 QAT(양자화 인식 훈련)를 적용해 int8 모델이 훈련된 모델이 되도록 했다.
- 스트림 강화학습 에이전트(Stream AC(λ))를 도입했으며, 리플레이 버퍼나 배치 없이 한 번에 한 전이만 학습한다. 에이전트는 프레임에서 읽은 특성 벡터(예: 공의 위치와 방향)를 입력받아 두 개의 64유닛 숨겨진 계층을 거친다.
- 보드 위에서 자체 생성한 꿈 안에서 에이전트가 처음부터 배우도록 할 수 있다. 피코의 코어 1에서 게임을 실행하고 코어 0에서 USB와 디스플레이를 처리해 프레임 생성이 화면 출력을 기다리지 않도록 설계했다.
- 메모리 할당이 런타임에 일어나지 않으며, 가중치는 플래시의 정적 배열에 저장되고 모든 버퍼는 고정 크기 SRAM이다. 2루프 기준 프레임당 2~4M 곱셈-누산이 필요하고, 데스크톱 CPU에서 같은 Rust 코드는 프레임당 0.7~1.5ms가 소요된다.
- 브레이크아웃은 361KB 플래시, 299KB SRAM을 사용하고, 가장 큰 팩맨은 600KB 플래시, 460KB SRAM을 사용한다. 피코 2 W의 4MB 플래시와 520KB SRAM 범위 안에 모든 게임이 들어간다.
- Int8 닷 곱셈이 대부분의 시간을 소비하며, 32비트 로드로 4개 가중치를 가져오고 SXTB16으로 언팩한 후 SMLAD 명령어로 두 개의 곱셈-누산을 한 번에 수행한다. 노름, 지수 테이블 조회, 최종 스케일은 FPU에서 float32로 유지된다.
- 같은 int8 엔진이 WebAssembly로도 컴파일되며, 브라우저에서 동일한 시드로 실행하면 네이티브 빌드와 프레임 단위로 정확히 동일한 에피소드를 생성한다. 게임당 1~5ms의 프레임 시간이 소요된다.
- 소스는 Python 훈련 코드(데이터 수집, 세계 모델 학습, 꿈 검증)와 재사용 가능한 PyTorch-to-Rust 내보내기 도구(QAT 계층, int8 스펙, 코드 생성)로 구성된다. Rust 펌웨어와 웹어셈블리 빌드도 포함되어 있다.
- 기존 방식과 달리 온디바이스 양자화와 반복 트랜스포머를 통해 메모리를 극도로 절약하면서도, 에이전트가 보드 위에서 처음부터 학습할 수 있고, 사용자가 버튼으로 언제든 조종권을 잡을 수 있다. 양자화된 int8 모델은 정확하게 검증되며, Rust와 Python 구현이 비트 단위로 동일함을 보장한다.