Tek bir GPU'da eğitilen JEPA dünya modeli, Pokémon Red'de başlangıç Pokémonunu seçmek için plan yapıyor

nostalgia.dev'de yayımlanan ve Hacker News'in ana sayfasında öne çıkan bir proje, model tabanlı planlama alanında küçük ama öğretici bir deneyi belgeliyor: bir dünya modeline Pokémon Red'in açılışını oynamayı öğretmek ve bunu, modelin Squirtle'ı başlangıç Pokémonu olarak seçen bir tuş dizisini planlamasıyla taçlandırmak.

Kasıtlı olarak dar bir görev

Yazıya göre özgün hedef, Profesör Oak'ın laboratuvarına ulaşmak, diyalogunu bitirmek, bir başlangıç Pokémonu seçmek, laboratuvardan çıkmak ve Oak'ın torununu yenmakti. Yazar, bunun ilk deney için fazla iddialı olduğuna hızla kanaat getirip hedefi daralttı: Oak'ın laboratuvarının içindeki kayıtlı bir durumdan başlayarak üç başlangıç Pokémonundan herhangi birini — Bulbasaur, Charmander veya Squirtle — başarıyla edinmek başarı sayılacaktı.

Bu daraltılmış görev bile kolay değil. Kayıtlı konumdan A tuşuna üst üste on iki kez basmak seçimi tamamlamak için yeterli. Ancak oyun hakkında hiçbir bilgisi olmayan bir ajan, laboratuvarda sonsuza dek dolaşabileceği gibi, ortada B'ye basarak kendi ilerlemesini de iptal edebilir.

Model gerçekte neyi öğreniyor

Yazar, dünya modelini mevcut bir gözlemle bir aksiyonu bir sonraki gözleme eşleyen bir fonksiyon olarak tanımlıyor ve bunu şematik olarak, bir sonraki durumun mevcut durum ile basılan tuşun bir fonksiyonuna yaklaşımları olarak yazıyor. Buradaki asıl nokta sola basmanın bir belirli odada işe yaradığını ezberlemek değil, sola basmanın oyundaki her yerde oyuncuyu sola hareket ettirdiği sezgisini genelleştirmek.

Kurulum, pekiştirmeli öğrenmeye (reinforcement learning) benziyor ama yazının altını çizdiği çok önemli bir farkla: dünya modeli, durum dinamiğini herhangi bir ödül sinyali olmadan öğreniyor; yani ödülsüz (reward-free).

Model, ham ekran görüntülerini tahmin etmek yerine latent uzayda çalışıyor. Bir encoder her ekran görüntüsünü 192 sayılık bir embedding'e — koordinatları tek tek insan tarafından anlam yüklenmemiş, öğrenilmiş bir vektör temsiline — dönüştürüyor. Ayrı bir predictor, mevcut embedding ile seçilen aksiyonu alıp tahmini bir sonraki embedding'i üretiyor; bu da ortalama kare hata ile gerçekten ardından gelen ekran görüntüsünün embedding'iyle karşılaştırılıyor. Tahmin edilen embedding hiçbir zaman görüntüye geri çözülmüyor; her şey öğrenilmiş temsil içinde gerçekleşiyor.

Çöküş problemi

Encoder ve predictor'ı birlikte eğitmek, yazının "tehlikeli dejenere durum" dediği durumu gizliyor. Aynı encoder hem predictor'ın girdisini hem de gerçek hedefi ürettiği için sistem, tahmin kaybını embedding'i yararlı bilgiden yoksun bırakacak bir biçimde tatmin edebiliyor. Yazı bunun üstesinden SIGReg ile geliyor; embedding uzayını kullanışlı tutmakla övülen bu regularisation tekniği, JEPA gibi latent tahmin mimarilerinin eğitilebilir olması için çözmek zorunda olduğu klasik temsil çöküşü (representation collapse) problemi.

Mimarinin kendisi, Yann LeCun'un JEPA'sının (Joint Embedding Predictive Architecture) bir varyantı; yazarın çekici bulduğu neden ise, modelin tek bir RTX 3080 Ti'da yerel olarak eğitebilecek kadar küçük ve birçok önceki JEPA modelinden tasarımda daha basit olması.

Başarıdan önceki başarısızlıklar

Deney yolunda gitmedi. Yazar, model daha basit testlerde umut vaat etmiş olsa da beklenenden daha fazla zorluk ve aksilik yaşadığını belirtiyor. Yazının yapısı, ilk planlama girişiminin başarısız olduğunu ve ikinci denemeden önce bir tur rollout fine-tuning yapıldığını gösteriyor: model artık yalnızca tek adımlı tahminlerle değil, kendi ürettiği çok adımlı dizilerle de eğitiliyor. Yazıda gösterilen nihai sonuç, eğitilmiş modelin Oak'ın laboratuvarındaki masadan Squirtle'ı seçen tuşları planlaması.

Neden önemli

Dünya modelleri, yapay zekâ araştırmalarının en yakından izlenen yönlerinden biri ve bu yazı bunun nedenini yerelden gösteren faydalı bir örnek. Tüketici sınıfı bir GPU'ya sığacak kadar küçük bir JEPA tarzı modelin, yalnızca ekran görüntüleri ve tuş girdilerinden kullanılabilir ortam dinamiği öğrenebildiğini ve ardından kendisinden bir plan istenebildiğini kanıtlıyor — ne ödül fonksiyonu var, ne elle yazılmış oyun mantığı.

En az bunun kadar değerli olanı, başarısızlık modlarının dürüstçe dökümü: latent tahmin temsil çöküşüne davetiye çıkarıyor ve bir adımı doğru tahmin eden bir model, planlamanın gerektirdiği uzun rollout'lar boyunca yine dağılabiliyor. Bunlar, alanın büyük ölçekte üzerinde çalıştığı gerilimlerin aynısı; burada 1996 tarihli bir Game Boy dünyasında yeniden üretilmişler — bu da yazıyı, model tabanlı ajanları takip eden herkes için araştırma makaleleriyle pratik arasında okunabilir bir köprü yapıyor.