یک مدل جهانی کوچک که بازی ها را روی Raspberry Pi Pico 2 W اجرا می کند

با استفاده از 2 شبکه عصبی کوچک برای تولید و اجرای بازی به طور همزمان روی Raspberry Pi Pico 2 W. اولین مورد بازی را ایجاد می کند: با توجه به فریم فعلی و عملکرد، فریم بعدی و پاداش را پیش بینی می کند. دومی آن بازی را انجام می دهد: از فریم تولید شده استفاده می کند و اقدامی را انتخاب می کند.

هر فریم بالا توسط شبکه مدل جهانی نمونه برداری شد و اقدامات توسط شبکه بازیگر انجام شد.

من به ترانسفورماتورهای حلقه ای علاقه مند شدم، جایی که یک بلوک ترانسفورماتور چندین بار با وزن های یکسان به جای انباشتن لایه های مختلف اعمال می شود. استفاده مجدد از وزنه ها برای حافظه ارزان است، بنابراین می خواستم ببینم آیا یک مدل جهان حلقه ای می تواند به اندازه کافی کوچک باشد که روی یک میکروکنترلر اجرا شود یا خیر. من Pico 2 W را انتخاب کردم زیرا ارزان است و دارای FPU است و بازی های کوچک را به این دلیل انتخاب کردم که می توانستم نسخه های واقعی را مستقیما در numpy بنویسم و ​​هر فریم تولید شده را با آنها بررسی کنم. بعدا یک عامل پخش جریانی RL اضافه کردم تا هیئت مدیره نیز بتواند در رویای خود بازی کند.

محتویات: دو شبکه · بازی ها · مدل جهانی · عوامل · کمی سازی · روی تراشه · چشمک زن · مرورگر · تولید مجدد · طرح بندی کد

اندازه به بازی بستگی دارد، برخی از بازی ها برای یادگیری آنها به پارامترهای بیشتری نیاز دارند تا قابل بازی باشند، ما سعی کردیم تعداد پارامترها را به اندازه کافی کوچک نگه داریم.

در تراشه هیچ یک از شبکه‌ها هرگز از بازی واقعی استفاده نمی‌کنند: بازیگر فقط فریم‌های مدل جهان را می‌بیند و مدل جهانی فقط اقداماتی را از بازیگر دریافت می‌کند. هر دو روی هسته 1 RP2350 کار می کنند، در حالی که هسته 0 USB و صفحه نمایش را کنترل می کند.

بازیگری که از قبل آموزش دیده را می توان با یک عامل جدید استریم-x که از ابتدا روی تخته یاد می گیرد تعویض کرد. و هر دکمه ای که با دست ها کنترل می شود، به این ترتیب می توانید رویا را خودتان بازی کنید.

هر بازی چند صد خط numpy در src/tinyworldformer/games/ است. به شبکه کوچکی از شناسه‌های توکن رندر می‌شود، و یک فریم همیشه کل حالت بازی را نگه می‌دارد، بنابراین حالت مخفی باید به داخل کادر کشیده شود: عنوان توپ بخشی از نشانه آن است و Doom-lite یک نقشه کوچک دارد.

«فریم‌های دقیق» سهم انتقال‌هایی است که در آن هر سلول از فریم پیش‌بینی‌شده با بازی واقعی مطابقت دارد. فریم‌ها درست بعد از یک رویداد تصادفی کنار گذاشته می‌شوند، زیرا فقط می‌توان آنها را نمونه‌برداری کرد، نه پیش‌بینی کرد. نمای سه‌بعدی Doom-lite دارای 768 سلول است و یک پیکسل دیوار سایه‌دار تمام فریم را اشتباه حساب می‌کند، بنابراین تعداد آن بدتر از نمایش آن به نظر می‌رسد. وضعیت بازی تقریبا همیشه درست است (به بررسی رویا مراجعه کنید).

می‌توانید بازی‌های واقعی را در ترمینال بازی کنید: uv run python -m tinyworldformer.games.play --game pacman (یا breakout، doomlite، flappy، pong).

قاب به تکه های مربعی برش داده شده است: سلول های 2×2 در Breakout، 4×4 در Doom-lite، Flappy Bird و Pong، 6×6 در Pac-Man. که 36 تا 80 توکن در هر بازی می دهد. تعبیه یک پچ مجموع یک بردار آموخته شده در هر سلول است، با یک جدول جداگانه برای هر موقعیت در داخل پچ تا ترتیب سلول ها از بین نرود. هر نشانه همچنین یک بردار موقعیت و بردار عمل دریافت می کند.

این مدل یک بلوک ترانسفورماتور منفرد (توجه به اضافه یک ReLU MLP) با یک مجموعه وزن است که 1 تا 4 بار اعمال می شود و ورودی قبل از هر پاس اضافه می شود. خروجی پس از هر پاس در برابر فریم بعدی آموزش داده می شود، بنابراین در زمان اجرا تعداد حلقه سرعت را با دقت عوض می کند. در Pico برای حفظ نرخ فریم مناسبی که استفاده می کنیم از دو حلقه استفاده می کنیم.

این انیمیشن یک مرحله Doom-lite را نشان می دهد. پس از یک پاس، مدل دشمنی را ترسیم می‌کند که آنجا نیست، و هر پاس بعدی بخشی از آن را حذف می‌کند: 92، 35، 12، سپس 0 سلول اشتباه. میله‌های سمت راست سهم فریم‌های نگه‌داشته‌شده هستند که تعداد حلقه‌ها دقیقا درست است.

در اینجا یک مرحله Doom-lite رویایی است که پس از 1، 2، 3 و 4 حلقه رمزگشایی شده است، با جعبه های قرمز رنگ روی سلول ها که با قوانین واقعی متفاوت است:

مدل برای هر سلول و هر توکن امتیازی را به دست می دهد. تبدیل آن امتیازها به یک قاب تقریبا به اندازه مدل اهمیت داشت و به سه قانون ختم شد (در بالا در Pong نشان داده شده است):

ما قانون 2 را اجرا کردیم زیرا اولین مدل Breakout دو توپ را همزمان ترسیم کرد. یک انتخاب جهانی فریم های دقیق را از 91.6٪ به 99.7٪ انتخاب کرد.

دقت در داده های ثبت شده کافی نبود. یک عامل رویا را به حالت هایی هدایت می کند که داده ها هرگز نداشته اند، و هر اشتباهی چیزی است که می تواند از آن سوء استفاده کند. بنابراین چک اصلی در حین بازی اجرا می شود: حالت بازی را از فریم رویایی برگردانید ( games/<game>.py: read )، قوانین واقعی را در آن حالت و عملکرد عامل اعمال کنید و نتیجه را با فریم رویایی بعدی و پاداش مقایسه کنید.

5000 مرحله رویایی در هر بازی در 2 حلقه، با نماینده ارسال شده ( training/evalwm.py --agent ):

اول، عوامل بازی واقعی را انجام می دهند و هر مرحله به عنوان یک مثال (فریم، اکشن، فریم بعدی، پاداش)، 1 تا 3 میلیون در هر بازی ثبت می شود. برخی از قسمت‌ها از حالت‌های غیرعادی شروع می‌شوند، مانند آجرهایی با سوراخ یا موقعیت‌های تصادفی، بنابراین داده‌ها بیشتر قسمت بازی را پوشش می‌دهند.

سپس مدل جهانی یاد می‌گیرد که فریم بعدی را با از دست دادن آنتروپی متقاطع در هر سلول پیش‌بینی کند: 30 هزار گام، با 3 تا 4 هزار گام اول به صورت شناور، سپس کلید گرد کردن int8 روشن می‌شود (QAT).

وزن‌های int8 در منبع Rust نوشته می‌شوند و هر شناور به‌عنوان بیت‌های دقیق آن ذخیره می‌شود. ما آنها را با تست بار بررسی می کنیم تا بیت مشخصات عدد صحیح پایتون را برای بیت بازتولید کنیم و سپس سیستم عامل روی Pico اجرا می شود.

یک عامل آموزش دیده در مدل جهانی به حالت هایی می رسد که داده های ضبط شده هرگز نداشته است، بنابراین بازی های واقعی را با آن نماینده ضبط کنید و مدل جهانی را روی همه چیز به خوبی تنظیم کنید.

عامل در بازی های مختلف اظهار داشت که از آنها سوء استفاده می کند، بنابراین چند راه حل برای آن وجود دارد:

عوامل از Stream AC(λ) (Elsayed, Vasan & Mahmood, 2024) استفاده می کنند که از یک انتقال در یک زمان یاد می گیرد: بدون بافر پخش مجدد، بدون دسته، بدون شبکه هدف. ردیابی واجد شرایط بودن اعتبار را به گذشته منتقل می کند و ObGD هر مرحله ای را که از هدف TD فراتر می رود کوچک می کند. شبکه های خط مشی و ارزش هر کدام دارای دو لایه مخفی 64 واحدی (32 واحد در Pico) با LayerNorm هستند.

ورودی یک بردار ویژگی است که از کادر خوانده می‌شود (training/streamx.py: Observer_*، در Rust در tinyworldformer-rs/src/agent.rs منعکس شده است). به عنوان مثال، ردیف، ستون و عنوان توپ در Breakout، یا در Pac-Man، برای هر جهت، کوتاه‌ترین فاصله تا نزدیک‌ترین گلوله و اینکه هر شبح چقدر نزدیک است.

در زیر یک تمرین در هر بازی (بازی واقعی، 500 هزار قدم) در شش مرحله آورده شده است: آموزش ندیده، پایان، و عکس های فوری که امتیاز 25، 50، 75 و 95 درصد از راه را پوشش داده است. هر پانل از همان شروع حریصانه بازی می کند و در هر فریم چند مرحله را رد می کند. "try" اپیزودهای پخش شده تا کنون را می شمارد تا بتوانید خرابی ها را ببینید.

بازگشت حریصانه در بازی واقعی، میانگین 10 قسمت، در عکس های فوری از همان اجرا:

بدون شکل دادن، عوامل Doom-lite در گوشه شروع اردو زده و شلیک می کنند، زیرا دشمن همیشه به سمت آنها می آید. -- idle-penalty 0.1 برای هر مرحله سپری شده در سلولی که بیش از 6 بار در 20 مرحله آخر بازدید شده است 0.1 هزینه دارد. موقعیت از روی نقشه کوچک خوانده می شود، بنابراین پنالتی در رویا نیز کار می کند. با آن مامور گشت می زند و دو برابر کشته می شود.

tinyworldformer-rs/src/learner.rs همان الگوریتمی در no_std Rust است که با گذر به عقب با دست نوشته شده است. 200 به روز رسانی PyTorch را با 3.5e-6 مطابقت می دهد. در میزبان، یک عامل تازه داخل رویای Rust int8 Flappy Bird میانگین 6 پیپ در هر قسمت بعد از 25 هزار مرحله، 121 بعد از 100k و 300+ بعد از 250k را دارد.

هر ماتریس در مدل جهانی int8 با یک مقیاس در هر ردیف خروجی است (مقیاس = max|w| / 127). فعال‌سازی‌ها نیز int8 هستند، با یک مقیاس برای هر تانسور و هر شاخص حلقه، زیرا حلقه‌ها محدوده‌های مختلفی را می‌بینند. Matmuls int8 × int8 را به int32 جمع می کند (دو ضرب-افزودن در هر دستورالعمل SMLAD در RP2350) و سپس یک float ضرب در دو مقیاس اعمال می کند. Softmax از یک جدول exp 1024 ورودی استفاده می کند که احتمالات آن به عنوان uint8 ذخیره می شود. هنجارها، باقیمانده ها و جاسازی ها در float32 باقی می مانند، زیرا RP2350 دارای یک FPU است.

آموزش همان گرد کردن (QAT با برآوردگر مستقیم) را اعمال می کند، بنابراین مدل int8 مدلی است که آموزش داده شده است. rustexport/intref.py مدل int8 در numpy با هر عملیات شناور به ترتیب ثابت است. موتور Rust آن را ذره ذره بازتولید می کند و آزمایش محموله بردارهای طلایی را برای هر env صادر شده بررسی می کند.

فلش و RAM از میان‌افزار ELF (اندازه -A) اندازه‌گیری می‌شوند که با درایور Display Pack ساخته شده است، بنابراین شامل کد، USB، درایور نمایشگر و یادگیرنده روی دستگاه می‌شود. Pico 2 W دارای 4 مگابایت فلش و 520 کیلوبایت SRAM است که 512 کیلوبایت آن در نقشه حافظه اصلی است. جداول تعبیه شده float32 بزرگترین قسمت هستند و می توانند int8 نیز باشند.

src/tinyworldformer/rustexport/ فقط numpy و مشعل را وارد می‌کند (یک آزمایش آن را حفظ می‌کند)، بنابراین پروژه‌های دیگر می‌توانند دوباره از صادرات استفاده کنند: یک مدل از rustexport.qat بسازید (QLinear، LoopedBlock)، آن را با ActQuant.enabled آموزش دهید = True پس از یک float گرم کردن rustemoop، سپس آن را بنویسید. منبع Rust با int_model_arrays، mlp_arrays و codegen.write_module. rustexport.Golden بردارها را می نویسد و تست زنگ بررسی می کند.

سیستم عامل (tinyworldformer-rs/firmware) no_std Rust on embassy است. هیچ چیزی در زمان اجرا تخصیص داده نمی شود: وزن ها آرایه های ثابت در فلش هستند و هر بافر یک استاتیک با اندازه ثابت در RAM است.

Core 1 بازی را اجرا می کند. یک پاس از حلقه آن (dream_loop در firmware/src/main.rs) یک فریم است:

هنگامی که بازی می کنید، هسته 1 به سرعت قابل بازی کاهش می یابد (بسته به بازی 66 تا 150 میلی ثانیه در هر فریم). عامل به همان سرعتی که تراشه اجازه می دهد اجرا می شود.

Core 0 وظایف async سفارت را برای هر چیز دیگری اجرا می کند. از طریق سریال USB، هر فریم به صورت متن ANSI رنگی با خط وضعیت (زمان مدل، زمان بازیگر، فریم در ثانیه، env و نام عامل) بیرون می‌آید و کلیدها باز می‌گردند. بسته نمایشی اختیاری (نمایش ویژگی‌ها) یک ST7789 320×240 روی SPI تا حداکثر 62.5 مگاهرتز است و تنها سلول‌های آخرین فریم تغییر می‌کنند. دکمه های آن هر 20 میلی ثانیه نظرسنجی می شوند. دو هسته فقط کانال فریم و چند اتمی مشترک دارند: چه کسی در حال بازی است، کلید فشرده، روشن/خاموش یادگیری و تعداد حلقه.

هر بازی در 4 مگابایت فلش و 512 کیلوبایت SRAM جای می گیرد. Breakout از 361 کیلوبایت فلش و 299 کیلوبایت SRAM استفاده می کند. Pac-Man، بزرگترین، از 600 کیلوبایت و 460 کیلوبایت استفاده می کند.

بیشتر اوقات به محصولات int8 dot اختصاص می یابد (dot_i8 در src/lib.rs). یک بار 32 بیتی 4 وزن را واکشی می کند، SXTB16 بایت های 0/2 و 1/3 را به نصف های 16 بیتی باز می کند و SMLAD دو ضرب-افزودن را در یک دستورالعمل انجام می دهد. توجه احتمالات uint8 را در مقادیر int8 ضرب می کند، بنابراین به جای آن از UXTB16 استفاده می کند. هنجارها، جستجوی جدول exp و مقیاس نهایی هر خروجی در float32 در FPU باقی می مانند. یک فریم در 2 حلقه 2-4M ضرب-افزودن است (8.2M برای Pac-Man). در یک CPU دسکتاپ، همان کد Rust رویای فریمی را در 0.7 ms (Flappy Bird) تا 1.5 ms (Breakout) می بیند. من هنوز نرخ فریم را روی خود پیکو اندازه گیری نکرده ام.

بوت هیئت مدیره با یک عامل تازه در دستگاه یاد می گیرد. l را روی USB فشار دهید، یا B را به مدت 1 ثانیه در صفحه نمایش بسته نگه دارید تا به بازیگر از قبل آموزش دیده بروید، و دوباره برای برگشتن.

هر صادراتی در کنار هم در جعبه قرار می گیرد: مدل های جهانی در tinyworldformer-rs/models/envs/<env>/ , agents in models/agents/<game>/<agent>/ . شما هر دو را در زمان ساخت انتخاب می کنید.

بدون --agent، عاملی برای بازی env که در مدل‌ها/پیش‌فرض فهرست شده است، در صورت وجود استفاده می‌شود، در غیر این صورت اولین نام آن است. یک نماینده برای یک بازی دیگر با لیستی از بازی های معتبر رد می شود. خط وضعیت روی صفحه نشان می دهد که کدام env و عامل در حال اجرا هستند.

کلیدهای روی USB: Breakout a/d/s، Doom-lite w/a/d/space/s، Flappy space، Pac-Man w/a/s/d، Pong w/s. سپس دست را به نماینده برگردانید، 1 تا 4 تعداد حلقه را تنظیم کنید، یادگیری روی دستگاه را تغییر دهید و r آن را مجددا راه اندازی کنید.

همان موتور int8 no_std نیز در WebAssembly (tinyworldformer-rs/web/) کامپایل می شود. این ماژول هیچ وارداتی و هیچ wasm-bindgen ندارد، و حلقه سفت‌افزار را اجرا می‌کند: dream، agent و stream-x Learning در حال پرواز. یک صفحه استاتیک فریم ها را روی بوم می کشد.

با همان seed، wasm build همان اپیزودهای ساخت بومی را، فریم به فریم، با سرعت 1 تا 5 میلی‌ثانیه در هر فریم بسته به بازی، رویا می‌کند. هر ماژول دارای عامل پیش فرض بازی خود است. TWF_AGENT=... ./build.sh <env> دیگری را انتخاب می کند. کلیدها مانند Pico هستند، به علاوه کلیدهای جهت‌نما و دکمه‌های روی صفحه، و f رویا را تا جایی که می‌تواند سریع اجرا می‌کند، که به یادگیری کمک می‌کند ( l ).