Update README.md

README.md

@@ -1,3 +1,28 @@
----
-license: mit
----
+---
+license: mit
+datasets:
+- reeeemo/jigsaw_puzzle
+pipeline_tag: reinforcement-learning
+library_name: stable-baselines3
+---
+
+Reinforcement-Learning model that utilizes MaskablePPO to guide assembly of a jigsaw puzzle (puzzle pieces with irregular, convex boundaries).
+
+Code to train/create this custom environment can be found in the "How Puzzling!" [github repo](https://github.com/reeeeemo/how-puzzling).
+
+Initialization Parameters:
+```
+def __init__(self, images, seg_model_path, max_steps=100, device="cpu")
+```
+- *images*: **list[np.ndarray]**
+  - List of images that **MUST BE SAME SAME, SAME ORIENTATION (3x3, 4x4 puzzles, etc.)**.
+  - Each will be randomly initialized every `.reset()` of the environment
+- *seg_model_path*: **string**
+  - Path to [image segmentation weight folder](https://huggingface.co/reeeemo/puzzle-segment-model)
+  - Note that any image segmentation that segments puzzle pieces will work, however they must be axis aligned
+  - The segmentation model is used to initialize another custom model found in the same github repo, [here](https://github.com/reeeeemo/how-puzzling/blob/main/model/model.py)
+- *max_steps*: **int**
+  - Max number of steps allowed
+  - Using MaskablePPO solves the issue of infinite actions, but if you decide to use PPO, ensure `max_steps` is set to the max number of puzzle pieces
+- *device*: **string**
+  - CPU or GPU usage