tinyworldformer 는 작은 신경망 두 개로 게임 생성과 플레이를 동시에 수행함. 월드 모델이 현재 프레임과 행동으로 다음 프레임과 보상을 예측하고, 행동 신경망이 생성된 화면을 보고 다음 행동을 선택함
루프형 트랜스포머 는 같은 블록의 가중치를 반복 사용해 메모리를 절약함. 월드 모델은 게임에 따라 2만 6,000~13만 1,000개 파라미터로 구성되며, Pico에서는 기본적으로 두 번 반복해 실행함
Breakout, Doom-lite, Flappy Bird, Pong, Pac-Man lite를 지원하며, 칩에서는 실제 게임 규칙 없이 두 신경망만으로 플레이함. 버튼을 누르면 사람이 직접 모델이 생성하는 게임을 플레이할 수도 있음
Stream AC(λ) 기반 에이전트는 재생 버퍼나 배치 없이 전이 하나씩 학습하므로 보드 안에서도 처음부터 학습할 수 있음. 다만 에이전트가 월드 모델의 오류를 악용할 수 있어 실제 규칙과 대조하고 데이터를 추가 수집해 보정함
int8 양자화와 no_std
Rust 구현으로 가장 큰 Pac-Man도 플래시 600KB, SRAM 460KB에 들어감. 같은 엔진은 WebAssembly에서도 동작하지만, Pico 자체의 프레임률은 아직 측정하지 않음
두 신경망으로 생성하고 플레이하는 게임
월드 모델 은 프레임 + 행동 → 다음 프레임 + 보상
, 행동 신경망 은 프레임 → 행동
을 담당함
월드 모델은 int8 루프형 트랜스포머로, 게임별 파라미터 수는 2만 6,00013만 1,000개, 프레임당 연산량은 약 200만800만 회의 곱셈-누산임
행동 신경망은 은닉층 두 개를 갖춘 float32 MLP로, 파라미터와 프레임당 곱셈-누산이 각각 약 9,0001만 6,000개임
가중치 재사용 으로 월드 모델을 마이크로컨트롤러에 넣을 수 있는지 확인하는 실험에서 출발함
Pico 2 W는 저렴하고 FPU가 있어 선택했으며, 작은 게임은 NumPy로 실제 규칙을 구현해 생성 프레임과 비교할 수 있어 선택함
이후 스트리밍 강화학습을 추가해 보드가 자체 생성 환경 안에서 학습하도록 확장함
칩에서는 어느 신경망도 실제 게임을 사용하지 않음
행동 신경망은 월드 모델이 만든 프레임만 보고, 월드 모델은 행동 신경망이 선택한 행동을 받음
두 신경망은 RP2350의 코어 1에서, USB와 디스플레이는 코어 0에서 실행함
사전 학습된 행동 신경망 대신 새 에이전트 를 넣어 보드에서 처음부터 학습할 수 있으며, 사람이 버튼을 누르면 직접 조작으로 전환됨
지원 게임과 프레임 정확도
각 게임은 수백 줄의 NumPy 코드 로 구현되며, 작은 토큰 ID 격자로 렌더링함
한 프레임에 게임 상태 전체가 들어가야 하므로 공의 진행 방향은 토큰에 포함하고, Doom-lite는 미니맵을 표시함
실제 규칙으로 동작하는 게임은 터미널에서도 플레이할 수 있음
Breakout 은 20×16 격자에 패들, 공, 벽돌 56개, 목숨 5개를 사용함
공을 놓친 뒤 서브가 무작위이며, 월드 모델은 파라미터 2만 6,000개와 430만 MACs를 사용함
2회 루프 기준 평가 프레임 완전 일치율은 99.99% 임
Doom-lite 는 32×32 격자로 8×8 미로의 레이캐스팅 1인칭 화면을 표현하며, 최단 경로로 접근하는 적을 사격함
적의 재생성이 무작위이며, 파라미터 5만 5,000개와 420만 MACs를 사용함
프레임 완전 일치율은 83.6% 지만, 768셀짜리 3D 화면에서 벽의 음영 픽셀 하나만 틀려도 프레임 전체를 오답으로 처리하므로 실제 플레이 상태 정확도보다 낮게 나타남
Flappy Bird 는 24×24 격자에서 중력, 날갯짓, 이동하는 파이프를 구현하며 다음 통로 높이가 무작위임
파라미터 4만 6,000개와 200만 MACs를 사용하고, 프레임 완전 일치율은 100% 임
Pong 은 20×32 격자에서 반응이 늦은 CPU 패들과 대결하며 먼저 5점을 얻으면 승리함
득점 뒤 서브가 무작위이며, 파라미터 4만 6,000개와 230만 MACs를 사용함
프레임 완전 일치율은 99.45% 임
Pac-Man lite 는 48×48 화면에 23×21 미로, 먹이 249개, 최단 경로로 이동하는 유령 두 마리, 목숨 3개를 구현함
유령의 재생성이 무작위이며, 파라미터 13만 1,000개와 820만 MACs를 사용함
프레임 완전 일치율은 86.4% 임
프레임 완전 일치율 은 학습에 쓰지 않은 전이에서 예측 프레임의 모든 셀이 실제 게임과 일치하는 비율임
무작위 사건 직후의 프레임은 예측이 아니라 샘플링 대상이므로 평가에서 제외함
패치 임베딩과 루프형 트랜스포머
화면을 정사각형 패치 로 나눠 게임별 3680개 토큰으로 변환함
Breakout은 2×2셀, Doom-lite/Flappy Bird/Pong은 4×4셀, Pac-Man은 6×6셀 패치를 사용함
패치 임베딩은 각 셀의 학습 벡터를 더해 만들며, 패치 내부 위치별 테이블을 따로 둬 셀 순서를 보존함
각 토큰에는 위치 벡터와 행동 벡터도 더함
어텐션과 ReLU MLP로 구성된 트랜스포머 블록 하나 를 같은 가중치로 14회 적용함
매 반복 전에 입력을 다시 더하고, 각 반복의 출력을 다음 프레임에 맞춰 학습함
실행 시 반복 횟수로 속도와 정확도를 조절하며, Pico의 기본값은 2회임
Doom-lite의 한 평가 사례에서는 반복을 늘리면서 존재하지 않는 적의 잘못된 셀이 92개 → 35개 → 12개 → 0개 로 줄어듦
반복 횟수를 늘린다고 전체 정확도가 항상 높아지지는 않음
Doom-lite의 완전 일치율은 14회에서 각각 59.8%, 83.6%, 85.9%, 84.7%임
Pac-Man은 각각 79.6%, 86.4%, 85.7%, 84.4%로 2회에서 가장 높았음
Breakout과 Pong은 3회에서 각각 100%, 99.60%였고, Flappy Bird는 2회부터 100%였음
디코딩으로 게임의 객체 제약 유지
모델이 출력한 셀별 토큰 점수를 화면으로 바꾸는 디코딩 규칙 이 모델 자체에 거의 맞먹을 만큼 중요했음
벽, 벽돌, 먹이, 바닥 같은 배경은 각 셀에서 가장 확률이 높은 토큰을 독립적으로 선택함
공, 플레이어, 각 유령처럼 하나만 있어야 하는 객체는 전체 (셀, 토큰)
쌍에서 한 번만 선택해 정확히 하나를 배치함
무작위 사건 뒤에는 최고 점수와 2.2 로짓 이내 인 후보에서 샘플링해 새 서브나 재생성을 생성함
Pac-Man과 유령의 2×2 스프라이트 는 상태를 담은 셀 하나만 예측하고 나머지 세 셀을 찍어 넣음
Pac-Man의 입은 진행 방향을, 유령의 왼쪽 아래 셀은 밑에 먹이가 있는지를 담음
초기 Breakout 모델은 공 두 개를 동시에 그렸으며, 전역 단일 객체 선택 을 적용한 뒤 완전 일치율이 91.6%에서 99.7%로 높아짐
실제 규칙 대조와 모델 오류 악용
기록 데이터의 정확도만으로는 충분하지 않음. 에이전트가 학습 데이터에 없던 상태로 진입 하고 그곳의 모델 오류를 악용할 수 있기 때문임
주요 검증은 생성 프레임에서 상태를 읽고, 그 상태와 행동에 실제 게임 규칙 을 적용한 결과를 다음 생성 프레임 및 보상과 비교함
배포 에이전트와 2회 루프로 게임별 5,000스텝 을 검사한 결과는 다음과 같음
Breakout은 상태를 역추출하지 않았으며, 불가능한 토큰이 있는 프레임은 0개였지만 벽돌이 깨지지 않았는데 보상을 지급한 경우가 6회였음
Doom-lite는 상태 오류 1회, 보상 오류 0회, 스텝당 잘못된 셀 1.50개였음
Flappy Bird는 상태, 보상, 셀 오류가 모두 0이었음
Pong은 상태 오류 1회, 보상 오류 0회였으며 스텝당 셀 오류는 표에서 0.00으로 기록됨
Pac-Man은 상태 오류 742회 , 보상 오류 38회, 스텝당 잘못된 셀 0.39개였음
에이전트가 발견한 악용 사례 를 실제 게임 데이터 재수집과 미세조정으로 보정함
Breakout에서는 모서리 벽돌을 맞혀도 제거하지 않는 오류 때문에 벽돌 26개가 남은 채 154점을 얻었고, Doom-lite에서는 5셀 밖에서 저격하는 행동을 학습함
생성 환경에서만 학습한 에이전트는 실제 최대 점수가 56점인 게임에서 생성 환경 점수 210점, 실제 게임 점수 28점을 기록함
이는 Ha와 Schmidhuber의 World Models 에서 예상한 문제와 같으며, 수집과 미세조정을 반복할 때마다 해당 에이전트가 찾은 악용 경로를 막음
데이터 수집과 학습 중 수정한 문제
실제 게임에서 에이전트가 플레이한 프레임, 행동, 다음 프레임, 보상 을 게임당 100만300만 개 수집함
일부 에피소드는 구멍 난 벽돌 배치나 무작위 위치 같은 특이한 상태에서 시작해 데이터 범위를 넓힘
셀별 교차 엔트로피 손실 로 다음 프레임을 학습하며, 기본 과정은 3만 스텝임
처음 3,0004,000스텝은 부동소수점으로 실행한 뒤 int8 반올림을 활성화함
학습한 모델은 Rust 소스로 내보내며, 부동소수점 값은 정확한 비트값으로 저장함
생성 환경에서 학습한 에이전트로 실제 게임 데이터를 다시 수집하고 전체 데이터로 미세조정함
손실에 변경된 셀 여부 로 가중치를 주자 변화를 과도하게 예측해 공 두 개와 벽돌 두 개 파괴를 생성함
입력 프레임에서 사건이 일어날 수 있는 위치를 기준으로 가중치를 주고, 단일 객체 선택을 적용해 수정함
Doom-lite에서 적이 사라진 뒤 돌아오지 않는 문제는 미니맵을 바탕으로 다시 그리지 않고 이전 화면을 복사하면서 발생함
GameNGen처럼 입력 화면에 무작위 픽셀을 넣고 때로 적 스프라이트를 지우는 노이즈 증강 으로 수정함
넓어진 Breakout 데이터는 d=32
로 학습하기 어려웠지만 d=48
로는 가능했으며, Pac-Man의 48×48 화면은 d=64
가 필요했음
스트리밍 강화학습 에이전트
Stream AC(λ) 는 Elsayed, Vasan, Mahmood의 2024년 알고리듬으로, 전이 하나씩 학습하며 재생 버퍼, 배치, 타깃 네트워크를 사용하지 않음
적격성 추적(eligibility traces)이 시간상 이전 행동에 학습 신호를 전달함
ObGD 는 시간차(TD) 타깃을 지나칠 수 있는 업데이트의 크기를 줄임
정책과 가치 네트워크는 각각 LayerNorm과 64유닛 은닉층 두 개를 사용하며, Pico에서는 32유닛으로 줄임
입력은 픽셀 전체가 아니라 프레임에서 읽은 특징 벡터 임
Breakout에서는 공의 행, 열, 진행 방향을 사용함
Pac-Man에서는 각 방향에서 가장 가까운 먹이까지의 최단 거리와 유령의 근접도를 사용함
실제 게임에서 게임별 한 번씩 50만 스텝 학습하고, 각 시점의 탐욕적 정책을 10개 에피소드 평균으로 평가함
최종 점수는 Breakout 45.6/56, Doom-lite 82.9, Flappy Bird 332/332, Pac-Man 217.4/249, Pong 1.1/5였음
시각화는 미학습 상태, 최종 상태, 점수 향상의 25%, 50%, 75%, 95%에 도달한 시점을 같은 시작 상태에서 비교함
Doom-lite는 보상 조정이 없으면 적이 다가오는 시작 모서리에 머물며 사격함
최근 20스텝에서 6회 넘게 방문한 셀에 머무는 매 스텝에 0.1 페널티 를 부과하면 순찰하며 처치 수가 두 배로 늘어남
위치를 미니맵에서 읽으므로 생성 환경에서도 같은 페널티를 적용할 수 있음
no_std
Rust 학습기는 역전파를 직접 구현 했으며, PyTorch의 200회 업데이트와 오차 3.5e-6 이내로 일치함
호스트의 Rust int8 Flappy Bird 생성 환경에서 새 에이전트는 2만 5,000스텝 후 에피소드당 평균 파이프 6개, 10만 스텝 후 121개, 25만 스텝 후 300개 이상을 통과함
int8 양자화와 재현 가능한 Rust 내보내기
월드 모델의 모든 행렬은 int8 이며, 출력 행마다 max|w| / 127
스케일을 사용함
활성값도 int8이며, 반복별 값 범위가 달라 텐서와 루프 인덱스별로 스케일을 둠
int8 곱셈 결과를 int32로 누산하고 두 스케일의 곱을 부동소수점으로 적용함
Softmax는 1,024항 지수 테이블을 사용하고 확률을 uint8로 저장함
정규화, 잔차 연결, 임베딩은 FPU를 활용하는 float32로 유지함
양자화 인식 학습(QAT) 에 순전파와 같은 반올림 및 straight-through estimator를 사용해 실제 배포할 int8 모델을 학습함
NumPy 정수 참조 구현은 부동소수점 연산 순서를 고정하며, Rust 엔진은 이를 비트 단위로 재현 함
내보낸 환경마다 기준 벡터를 생성하고 cargo test
로 검증함
src/tinyworldformer/rustexport/
는 NumPy와 PyTorch만 의존 하며 다른 프로젝트에서도 재사용할 수 있음
QAT 레이어로 모델 구성, 부동소수점 준비 학습 후 양자화 활성화, 정수 모델 변환, Rust 코드 생성, 기준 벡터 검증으로 이어짐
Pico의 메모리와 실행 구조
펌웨어는 embassy 기반 no_std
Rust 이며 실행 중 메모리를 동적 할당하지 않음
가중치는 플래시의 static
배열, 버퍼는 RAM의 고정 크기 static
으로 유지함
코어 1의 프레임 루프 는 행동 선택, 월드 모델 추론과 디코딩, 선택적 학습, 화면 전달 순서로 동작함
행동은 사전 학습 모델, 학습 중인 에이전트, 사용자 입력 중에서 선택함
무작위 사건 샘플링에는 xorshift 난수 생성기를 쓰고, 보상은 보상 헤드의 최댓값으로 선택하며 목숨과 에피소드 종료는 별도 카운터로 관리함
학습 시 해당 전이로 한 번 업데이트한 뒤 버림
코어 0으로 보내는 2슬롯 채널이 바쁘면 프레임을 버려 추론이 화면 출력을 기다리지 않게 함
사람이 플레이할 때는 게임에 따라 프레임당 66150ms 로 속도를 제한하고, 에이전트는 칩이 허용하는 최고 속도로 실행함
코어 0은 USB와 선택적 Display Pack 을 비동기 작업으로 처리함
USB로 컬러 ANSI 프레임과 모델/행동 신경망 실행 시간, fps, 환경 및 에이전트 이름을 출력하고 키 입력을 받음
Display Pack은 320×240 ST7789이며 최대 62.5MHz SPI를 사용하고 변경된 셀만 다시 그림
버튼은 20ms마다 확인하며, 두 코어는 프레임 채널과 조작 주체, 키, 학습 여부, 루프 횟수에 관한 소수의 원자적 변수만 공유함
Pico 2 W는 플래시 4MB, SRAM 520KB 를 갖추며 주 메모리 맵에는 512KiB가 있음
Display Pack 드라이버와 온디바이스 학습기를 포함한 펌웨어의 플래시/RAM 사용량은 Breakout 361/299KB, Doom-lite 471/357KB, Flappy Bird 409/227KB, Pong 414/244KB임
가장 큰 Pac-Man도 600/460KB 로 들어감
메모리에서 큰 비중을 차지하는 부분은 float32 임베딩과 추론 임시 버퍼 임
월드 모델의 플래시 데이터는 52138KB이며, 임베딩 테이블은 int8로 바꿀 여지가 있음
추론 버퍼는 Flappy Bird 114KB부터 Pac-Man 368KB까지이며, Pac-Man은 셀별 출력 점수만 184KB임
행동 신경망은 플래시 5161KB와 작업 RAM 약 1KB, 학습기는 특징 수에 따라 RAM 49106KB를 사용함
코드, USB, 디스플레이 드라이버는 약 150KB이며 코어 1 스택은 16KB임
실행 시간 대부분은 int8 내적 에 사용됨
32비트 로드로 가중치 4개를 읽고 SXTB16
으로 분리한 뒤 SMLAD
한 명령으로 곱셈-누산 두 번을 수행함
uint8 어텐션 확률과 int8 값을 곱할 때는 UXTB16
을 사용함
같은 Rust 코드의 데스크톱 프레임 생성 시간은 Flappy Bird 0.7ms~Breakout 1.5ms이며, Pico 프레임률은 미측정 임
실행 방식과 브라우저 지원
게임별 월드 모델과 에이전트는 빌드 시 선택 하며, 다른 게임용 에이전트 조합은 거부함
보드는 기본적으로 새 에이전트의 온디바이스 학습 을 시작함
USB의 l
또는 Display Pack B 버튼 1초 누르기로 사전 학습 모델과 전환할 수 있음
사용자 조작 후 에이전트에 제어를 돌려주거나 루프 횟수를 1~4회로 조정하고 학습을 재시작할 수 있음
같은 no_std
int8 엔진 을 WebAssembly로도 빌드함
외부 임포트와 wasm-bindgen
없이 게임 생성, 행동 선택, 실시간 Stream-x 학습을 수행하고 정적 페이지의 캔버스에 표시함
같은 시드에서는 네이티브 빌드와 프레임 단위로 동일한 에피소드를 생성하며, 게임에 따라 프레임당 1~5ms가 걸림
게임별 모듈에 기본 에이전트를 포함하며, 키보드와 화면 버튼을 지원하고 최대 속도 실행으로 학습을 빠르게 진행할 수 있음
코드는 게임 규칙 , 학습/검증, 재사용 가능한 Rust 내보내기, Rust 엔진/학습기, Pico 펌웨어, WebAssembly 페이지로 구분됨
전체 재현 흐름은 실제 게임 데이터 수집, 월드 모델 학습, 에이전트 학습과 생성 환경 검증, Rust 내보내기, 기준 벡터 테스트 및 실행으로 이어짐
댓글과 토론