라즈베리 파이 피코 2 W에서 루프형 월드 모델과 RL 액터를 실행하는 TinyWorldFormer

tinyworldformer는 SheepSeb의 오픈소스 프로젝트로, 라즈베리 파이 피코 2 W에서 두 개의 작은 신경망을 실행합니다: 게임 프레임을 생성하는 월드 모델과 그 프레임을 플레이하는 사전 훈련된 액터입니다. 두 네트워크 모두 칩에서 실제 게임을 보지 않습니다 — 액터는 월드 모델이 생성한 프레임만 읽고, 월드 모델은 액터로부터 액션만 받습니다.

아키텍처

• 월드 모델: int8 양자화된 루프 트랜스포머. 26k~131k 파라미터, 프레임당 2~8M 곱셈-덧셈. 입력은 현재 프레임 + 액션, 출력은 다음 프레임 + 보상.
• 액터: 두 개의 은닉층을 가진 MLP, float32. 9k~16k 파라미터, 프레임당 9k~16k 곱셈-덧셈. 입력은 프레임, 출력은 액션.
• 둘 다 RP2350의 코어 1에서 실행되며, 코어 0은 USB와 디스플레이를 처리합니다.
• 사전 훈련된 액터는 보드에서 처음부터 학습하는 새로운 stream-x RL 에이전트로 교체할 수 있습니다. 버튼을 누르면 인간 플레이어에게 제어권이 넘어갑니다.

게임 및 벤치마크 결과

각 게임은 src/tinyworldformer/games/에 있는 수백 줄의 NumPy 코드로, 작은 토큰 ID 그리드로 렌더링됩니다. 한 프레임이 전체 게임 상태를 담으므로 숨겨진 상태가 프레임에 그려집니다(공의 진행 방향은 토큰의 일부이며, 둠-라이트에는 미니맵이 있습니다).

• 브레이크아웃 — 20×16 프레임, 26k 파라미터, 4.3M MAC, 99.99% 정확한 프레임.
• 둠-라이트 — 8×8 미로의 32×32 레이캐스트 뷰, 55k 파라미터, 4.2M MAC, 83.6% 정확한 프레임(음영 처리된 벽 픽셀 하나만 틀려도 전체 프레임이 틀린 것으로 간주).
• 플래피 버드 — 24×24 프레임, 46k 파라미터, 2.0M MAC, 100% 정확한 프레임.
• 퐁 — 20×32 프레임, 46k 파라미터, 2.3M MAC, 99.45% 정확한 프레임.
• 팩맨 라이트 — 48×48 프레임, 23×21 미로, 249개 펠릿, 131k 파라미터, 8.2M MAC, 86.4% 정확한 프레임.

"정확한 프레임"은 예측된 프레임의 모든 셀이 실제 게임과 일치하는 홀드아웃 전환의 비율입니다(무작위 이벤트 직후의 프레임은 제외).

월드 모델의 작동 방식

프레임은 정사각형 패치로 분할됩니다(브레이크아웃은 2×2 셀, 둠-라이트, 플래피 버드, 퐁은 4×4, 팩맨은 6×6). 게임당 36~80개의 토큰이 생성됩니다. 각 패치 임베딩은 셀당 하나의 학습된 벡터와 패치 내 위치별 별도 테이블, 그리고 위치 및 액션 벡터를 합산합니다. 모델은 단일 트랜스포머 블록(어텐션 + ReLU MLP)으로, 하나의 가중치 세트를 여러 번 적용하는 루프 트랜스포머로, 메모리를 마이크로컨트롤러에 맞게 낮게 유지합니다.

실제 게임 실행

uv run python -m tinyworldformer.games.play --game pacman

breakout, doomlite, flappy, pong으로 대체할 수 있습니다. 저장소에는 tinyworldformer-rs, 문서, 테스트, 브라우저 데모도 포함되어 있습니다.

엣지 추론이나 작은 월드 모델을 연구하는 개발자에게 실질적인 시사점은 메모리 대 연산 트레이드오프입니다: 하나의 트랜스포머 블록 가중치를 재사용하면 26k 파라미터 모델이 프레임당 수백만 MAC으로 실행되면서도 피코 2 W의 한계 내에 머물 수 있습니다.