subbonan/nitrogen-pizza-tower-finetune

NitroGen Fine-Tune

This is a fine-tuned derivative of nvidia/NitroGen, adapted for any game from locally recorded gameplay (example: Pizza Tower).

• GitHub repository: https://github.com/subbon/NitroGen-Full

Base Model

• Base model: nvidia/NitroGen
• Original paper: https://arxiv.org/abs/2601.02427
• Original repository: https://github.com/MineDojo/NitroGen

What This Version Changes

Compared to the original NitroGen release, this fine-tuned version adds a complete pipeline for adapting the model to any target game using custom gameplay recordings.

Main practical changes:

• Direct local inference without the original server/client split
• Custom gameplay dataset recording from Windows
• NitroGen post-training with EMA, BF16 AMP, and torch.compile
• Stronger loss weighting for active action targets

In the current training code, non-neutral action targets are upweighted by x3.5 in the flow-matching loss to make rare button presses and meaningful actions matter more during fine-tuning.

Training Summary

Latest Training (1.0x weights — final_model.pt)

• Base checkpoint: ng.pt
• Dataset: 34 runs, 681,203 samples
• Train / val split: 613,083 / 68,120
• Trainable parameters: 177,675,289 / 493,631,513
• Training environment: WSL + RTX 4070 Ti Super
• Schedule: OneCycle
• Epochs: 10
• Best validation loss: 0.0160 (EMA, epoch 10)
• Best Macro F1: 44.61% (per benchmark)
• Date: March 9, 2026

Previous Training (3.5x weights — final_model_35.pt)

• Loss Weight: 3.5x for active actions
• Best validation loss: 0.0209
• Best Macro F1: 45.49% (per benchmark)
• Date: March 6, 2026

Model Versions

Model	Loss Weight	Loss	Macro F1	Description
final_model.pt	1.0x	0.0158	44.61%	Recommended — Standard loss weights, best stability for gameplay
final_model_35.pt	3.5x	0.0209	45.49%	Best Macro F1 — upweighted actions, potentially less stable
ng.pt (base)	—	0.0491	7.24%	Original foundation model (for comparison)

Benchmarked on validation set (200 batches, 2 runs per model)

Recommendation: Use final_model.pt (1.0x) for actual gameplay — it provides more consistent behavior. The 3.5x model upweights rare actions heavily which can lead to unpredictable button spam.

Gameplay Examples

1.0x Model (final_model.pt) — Recommended

3.5x Model (final_model_35.pt)

License and Restrictions

This model is derived from NitroGen and should be shared under terms compatible with the original license and notices.

The original repository license is the NVIDIA License and includes a non-commercial research use limitation. Review the original project license before redistribution or downstream use:

• https://huggingface.co/nvidia/NitroGen
• https://github.com/MineDojo/NitroGen

Files

• final_model.pt — Recommended. Model with 1.0x loss weighting. Best stability for gameplay (Loss: 0.0158, Macro F1: 44.61%).
• final_model_35.pt — Model with 3.5x loss weighting. Higher Macro F1 (45.49%) but higher loss (0.0209), potentially less predictable behavior.

Inference

This model uses the companion code in this repository fork.

Recommended (1.0x model — more stable):

python play.py path\to\final_model.pt --process PizzaTower.exe

Alternative (3.5x model — higher F1 but less stable):

python play.py path\to\final_model_35.pt --process PizzaTower.exe

Note: If speedhack fails with Error 5, run your terminal as Administrator.

Citation

Please cite the original NitroGen paper if you use this model:

@misc{magne2026nitrogen,
  title={NitroGen: An Open Foundation Model for Generalist Gaming Agents},
  author={Loic Magne and Anas Awadalla and Guanzhi Wang and Yinzhen Xu and Joshua Belofsky and Fengyuan Hu and Joohwan Kim and Ludwig Schmidt and Georgia Gkioxari and Jan Kautz and Yisong Yue and Yejin Choi and Yuke Zhu and Linxi "Jim" Fan},
  year={2026},
  eprint={2601.02427},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2601.02427}
}