Mô hình thế giới siêu nhỏ chạy game trên Raspberry Pi Pico 2 W

Mô hình thế giới tinyworldformer

Tóm tắt nhanh: Dự án tinyworldformer của lập trình viên SheepSeb đưa một mô hình thế giới và một tác nhân học tăng cường vào vi mạch Raspberry Pi Pico 2 W chỉ với 4 MB flash và 520 KB SRAM. Hai mạng nơ-ron nhỏ phối hợp để vừa tạo khung hình game vừa chơi, tất cả diễn ra trong "giấc mơ" do chính AI sinh ra — không cần chạy game thật.

Ý tưởng đằng sau dự án

Điểm khởi đầu là mối quan tâm tới looped transformer — kiến trúc áp dụng lặp lại một khối transformer duy nhất với cùng bộ trọng số, thay vì xếp chồng nhiều lớp khác nhau. Việc tái sử dụng trọng số giúp tiết kiệm bộ nhớ, nên tác giả đặt câu hỏi: liệu một mô hình thế giới dạng này có đủ nhỏ để chạy trên vi điều khiển hay không?

Tác giả chọn Pico 2 W vì giá rẻ và có FPU (đơn vị dấu phẩy động). Các game được chọn đều đơn giản để có thể viết lại bằng numpy phía Python, từ đó đối chiếu từng khung hình do AI sinh ra với luật game thật.

Hai mạng nơ-ron, hai nhiệm vụ

Hệ thống gồm hai mạng với vai trò tách biệt:

• Mạng thế giới (world model): nhận khung hình hiện tại và hành động, dự đoán khung hình tiếp theo cùng phần thưởng. Đây là transformer lặp, định lượng int8, với 26k–131k tham số tùy game.
• Mạng tác nhân (actor): nhận khung hình và chọn hành động. Đây là MLP hai lớp ẩn, float32, chỉ 9k–16k tham số.

Điều thú vị là trên chip, không mạng nào từng chạm vào game thật. Tác nhân chỉ thấy khung hình do mô hình thế giới tạo ra, còn mô hình thế giới chỉ nhận hành động từ tác nhân. Cả hai chạy trên lõi 1 của RP2350, trong khi lõi 0 xử lý USB và màn hình.

Năm tựa game trong mơ

Mỗi game chỉ vài trăm dòng numpy, hiển thị thành lưới nhỏ các token ID:

Game	Khung hình	Tham số mô hình	Tỷ lệ khung chính xác
Breakout	20×16	26k	99,99%
Doom-lite	32×32	55k	83,6%
Flappy Bird	24×24	46k	100%
Pong	20×32	46k	99,45%
Pac-Man lite	48×48	131k	86,4%

Giấc mơ Breakout do AI tạo ra

Một khung hình luôn chứa toàn bộ trạng thái game, nên các trạng thái ẩn phải được vẽ trực tiếp vào khung — ví dụ hướng bay của bóng nằm trong token của bóng, còn Doom-lite có thêm bản đồ nhỏ (minimap).

Kỹ thuật then chốt: giải mã và lượng tử hóa

Theo tác giả, bước giải mã (decoding) quan trọng gần ngang với bản thân mô hình. Ba quy tắc được áp dụng:

• Ô nền (tường, gạch, thức ăn) lấy token có xác suất cao nhất một cách độc lập.
• Các vật thể đơn (bóng, người chơi, ma) được đặt bằng một lựa chọn toàn cục trên mọi cặp (ô, token), đảm bảo luôn có đúng một thực thể. Sau biến cố ngẫu nhiên, lựa chọn được lấy mẫu để tạo ra giao bóng hay hồi sinh mới thay vì lấy trung bình.
• Với Pac-Man và ma, mỗi sprite 2×2 chỉ một ô mang trạng thái, ba ô còn lại được đóng dấu từ ô đó.

Quy tắc thứ hai ra đời sau khi mô hình Breakout đầu tiên vẽ hai quả bóng cùng lúc — sửa lỗi này đã đưa độ chính xác khung hình từ 91,6% lên 99,7%.

Toàn bộ ma trận trong mô hình thế giới được lượng tử hóa int8 với một scale cho mỗi hàng đầu ra. Phép nhân ma trận cộng int8 × int8 vào int32 — hai phép nhân-cộng mỗi lệnh SMLAD trên RP2350 — rồi nhân một lần với hai scale. Quá trình huấn luyện dùng kỹ thuật QAT (lượng tử hóa khi huấn luyện) nên mô hình int8 chính là mô hình đã được huấn luyện.

Tác nhân stream-x: học từng bước một

Tác nhân dùng thuật toán Stream AC(λ) của Elsayed, Vasan và Mahmood (2024): học từ một chuyển trạng thái tại một thời điểm, không có replay buffer, không theo lô, không mạng đích. Nhờ vậy thuật toán đủ nhẹ để chạy trong môi trường no_std Rust trên chip.

Giấc mơ Flappy Bird

Tác giả cũng thừa nhận các tác nhân từng "khai thác lỗi" của mô hình thế giới, ví dụ một tác nhân Breakout trong mơ đạt 154 điểm (trong khi mức tối đa thật chỉ là 56), hay tác nhân Doom-lite học cách bắn tỉa từ khoảng cách 5 ô. Mỗi vòng thu thập dữ liệu thật và tinh chỉnh lại bịt được một lỗ hổng như vậy.

Chạy trên chip và trong trình duyệt

Firmware là Rust no_std chạy trên embassy, không cấp phát bộ nhớ động: trọng số là mảng tĩnh trong flash, mọi buffer có kích thước cố định. Doom-lite dùng 471 KB flash và 357 KB RAM; Pac-Man lớn nhất cần khoảng 600 KB flash và 460 KB SRAM — vẫn vừa với cấu hình của Pico 2 W.

Đáng chú ý, cùng engine int8 đó có thể biên dịch sang WebAssembly và chạy trong trình duyệt, cho ra kết quả giống hệt bản native từng khung hình, với tốc độ 1–5 ms mỗi khung.

Người dùng có thể chuyển đổi giữa tác nhân đã huấn luyện và một tác nhân học từ đầu ngay trên board, hoặc nhấn nút để tự chơi giấc mơ của AI. Với những ai quan tâm tới AI biên (edge AI) và mô hình thế giới, đây là một dự án mã nguồn mở đáng để thử nghiệm — đặc biệt khi các thiết bị nhúng giá rẻ ngày càng được kỳ vọng chạy được những mô hình từng chỉ dành cho máy chủ.

Toàn bộ mã nguồn và tài liệu tái lập thí nghiệm có tại kho GitHub của dự án: SheepSeb/tinyworldformer.