Open Dreamer: Implementasi Open-Source Pipeline Dreamer 4 di JAX/Flax dengan Resep Pelatihan Lengkap

Peneliti Francesco Sacco, Diego Martí, dan Edward Hu, yang disponsori oleh Reactor, merilis Open Dreamer, implementasi terbuka dari Dreamer 4 dalam JAX/Flax NNX. Dua repositori dirilis: satu untuk pelatihan (tokenizer, model dinamika, rollout, pemeringkatan FVD) dan satu untuk inferensi lokal. Demo browser mengalirkan dunia Minecraft yang dihasilkan dan memungkinkan beralih antara game nyata dan model dunia. Arsitektur ini menggunakan backbone transformer blok-kausal untuk tokenizer dan model dinamika. Tokenizer adalah Masked Autoencoder berbasis transformer yang mencapai kompresi sekitar 100× tanpa kerugian KL atau adversial. Model dinamika (1,6 miliar parameter) melakukan prediksi bingkai berikutnya menggunakan diffusion forcing, flow matching, dan model pintasan, serta memprediksi tindakan berikutnya. Pelatihan berjalan selama 200.000 langkah dengan pengoptimal Muon, jadwal WSD, tingkat pembelajaran puncak 3e-4, dan peluruhan EMA 0,999. Tim mencapai pemanfaatan FLOP model sebesar 57-58% pada GPU B200, dengan status model muat dalam ~24 GiB. Aktivasi adalah biaya memori utama, yang diatasi dengan paralelisme data dan checkpointing aktivasi. Masalah stabilitas adalah tantangan terbesar; enam perbaikan termasuk pengoptimal Muon, bobot EMA wajib, presisi campuran yang hati-hati, pembobotan kerugian v-space, transportasi optimal barycentric minibatch, dan tidak menekankan parameterisasi-μ. Repositori tidak menyertakan loop pelatihan peniruan perilaku atau RL, dan skor FVD tidak dipublikasikan.