Made refactoring, added readme.md

Author: juraam

.gitignore

@@ -3,8 +3,10 @@ __pycache__/
 *.py[cod]
 *$py.class
 
-training_data
-test_models
+trained-models
+papers
+analysis
+runpod
 
 # Distribution / packaging
 dist/

README.md

@@ -0,0 +1,52 @@
+# Snake Diffusion model
+
+It is an educational repo to build realtime snake game based on Diffusion model. It was inspired by great papers:
+* Doom Diffusion from Google ([paper](https://arxiv.org/html/2408.14837v1))
+* Oasis ([github](https://github.com/etched-ai/open-oasis))
+* Diamond ([paper](https://arxiv.org/pdf/2405.12399))
+
+My goal was to build something similar and I have choosen Snake game for simple logic. It took near 2 months of different experiments to get a ready-to-play model.
+
+If you don't have GPU you can use [runpod.io](runpod.io)(it is paid). Also I created a [Google colab](https://colab.research.google.com/drive/1OxneGBeb4B1U5dszVf_2UDYZLHshJQ5T?usp=sharing) for playing.
+
+## Model scheme 
+
+After couple of experiments I chose EDM diffusion model, because it shows high performance on small sample steps. DDIM requires much more steps to generate the same quality.
+
+![Model scheme](assets/scheme.png)
+
+## Install requirements
+
+```shell
+pip install -r requirements.txt
+```
+
+## Training
+
+To train a new model, you should have a dataset. You can download it running a script:
+```shell
+bash scripts/download-dataset.sh
+```
+
+Or generate manually:
+
+```shell
+python src/generate_dataset.py --model agent.pth --dataset training_data --record
+```
+
+Then you can start training with command:
+```shell
+python src/train.py --model-type edm --output models --loader loader.pkl --gen-val-images --config Diffusion.yaml
+```
+
+I trained my model on [runpod.io](runpod.io). It had 32 epochs, took ~27 hours and the cost was 10$.
+
+## Inference
+
+You can download my ready-to-use model:
+```shell
+git clone https://huggingface.co/juramoshkov/snake-diffusion models
+```
+Then run [Play.ipynb](src/play.ipynb), where you can play Snake with 1 fps 🤓.
+
+Another way to play is to use [Google colab](https://colab.research.google.com/drive/1OxneGBeb4B1U5dszVf_2UDYZLHshJQ5T?usp=sharing)

assets/scheme.png

@@ -0,0 +1,35 @@
+training:
+  epochs: 30
+  batch_size: 4
+  num_workers: 2
+  save_every_epoch: 2
+
+generation:
+  image_size: 64
+  input_channels: 3
+  output_channels: 3
+  context_length: 4
+  actions_count: 5
+
+edm:
+  p_mean: -1.2
+  p_std: 1.2
+  sigma_data: 0.5
+  sigma_min: 0.002
+  sigma_max: 80
+  rho: 7
+  unet:
+    __type__: models.gen.blocks.UNetConfig
+    steps: [2, 2, 2, 2]
+    channels: [64, 64, 64, 64]
+    cond_channels: 256
+    attn_step_indexes: [false, false, false, false]
+
+ddpm:
+  T: 1000
+  unet:
+    __type__: models.gen.blocks.UNetConfig
+    steps: [2, 2, 2, 2]
+    channels: [64, 64, 64, 64]
+    cond_channels: 256
+    attn_step_indexes: [false, false, false, false]

config/Diffusion.yaml

@@ -0,0 +1,19 @@
+q_agent:
+  max_memory: 100000
+  batch_size: 1000
+  lr: 1.0e-3
+  hidden_state: 256
+  value_for_end_game:
+    __type__: q_agent.ValueForEndGame
+    value: "last_action"
+  iterations: 80000
+  min_deaths_to_record: 60
+
+env: 
+  __type__: game.snake.env.GameEnvironment
+  game:
+    __type__: game.snake.game.SnakeGame
+    width: 64
+    height: 64
+    speed: 240
+    block_size: 5