diff --git a/README.md b/README.md
index 4e7e315bc87854f056c643808796056629f27e75..7da169c9bb2e36fc85f6e839603ffa27301aadaa 100644
--- a/README.md
+++ b/README.md
@@ -10,7 +10,7 @@ The real-time overcooked simulation for a cognitive cooperative system.
 
 You have two options to install the environment. Either clone it and install it locally or install it in your
 site-packages.
-You need a Python 3.10 or higher environment. Either conda or PyEnv.
+You need a Python 3.10 or newer environment. Either conda or PyEnv.
 
 ### Local Editable Installation
 
@@ -73,7 +73,8 @@ can be cooked/created.
 ### Layout Config
 
 You can define the layout of the kitchen via a layout file. The position of counters are based on a grid system, even
-when the players do not move grid steps but continuous steps. Each character defines a different type of counter.
+when the players do not move grid steps but continuous steps. Each character defines a different type of counter. Which
+character is mapped to which counter is defined in the Environment config.
 
 ### Environment Config
 
diff --git a/overcooked_simulator/__init__.py b/overcooked_simulator/__init__.py
index cfd15d482adac935144d6b73af50657fec6b67c4..5d7acaeac21d083a862405fb39c9060ac5365e30 100644
--- a/overcooked_simulator/__init__.py
+++ b/overcooked_simulator/__init__.py
@@ -5,8 +5,8 @@ This is the documentation of the Overcooked Simulator.
 # About the package
 
 The package contains an environment for cooperation between players/agents. A PyGameGUI visualizes the game to
-human or visual agents in 2D. A 3D web-enabled version (for example for online studies, currently under development)
-can be found [here](https://gitlab.ub.uni-bielefeld.de/scs/cocosy/godot-overcooked-3d-visualization)
+human or virtual agents in 2D. A 3D web-enabled version (for example for online studies, currently under development)
+can be found [here](https://gitlab.ub.uni-bielefeld.de/scs/cocosy/godot-overcooked-3d-visualization).
 
 # Background / Literature
 The overcooked/cooking domain is a well established cooperation domain/task. There exists
@@ -19,9 +19,23 @@ With this overcooked-simulator, we want to bring both worlds together: the reinf
 environment with an appealing visualisation. Enable the potential of developing artificial agents that play with humans
 like a "real" cooperative / human partner.
 
+# Installation
+
+You need a Python **3.10** or newer environment.
+```bash
+pip install overcooked-environment@git+https://gitlab.ub.uni-bielefeld.de/scs/cocosy/overcooked-simulator@main
+```
+Or clone it and install it as an editable library which allows you to use all the scripts directly.
+```bash
+git clone https://gitlab.ub.uni-bielefeld.de/scs/cocosy/overcooked-simulator.git
+cd overcooked_simulator
+pip install -e .
+```
+
 # Usage / Examples
-Our overcooked simulator is designed for real time interaction but also with reinforcement learning in mind (gymnasium environment).
-It focuses on configurability, extensibility and appealing visualization options.
+ Our overcooked simulator is designed for real time interaction but also for reinforcement
+learning (gymnasium environment) with time independent step function calls. It focuses on configurability, extensibility and appealing visualization
+options.
 
 ## Human Player
 Run it via the command line (in your pyenv/conda environment):
@@ -30,7 +44,7 @@ Run it via the command line (in your pyenv/conda environment):
 overcooked-sim  --url "localhost" --port 8000
 ```
 
-_The arguments are the defaults. Therefore, they are optional._
+*The arguments shown are the defaults.*
 
 You can also start the **Game Server** and the **PyGame GUI** individually in different terminals.
 
@@ -41,12 +55,261 @@ python3 overcooked_simulator/gui_2d_vis/overcooked_gui.py --url "localhost" --po
 ```
 
 ## Connect with agent and receive game state
-...
+Or you start a game server, create an environment and connect each player/agent via a websocket connection.
+
+To start a game server see above. Your own manager needs to create an environment.
+```python
+import requests
+
+from overcooked_simulator import ROOT_DIR
+from overcooked_simulator.game_server import CreateEnvironmentConfig
+from overcooked_simulator.server_results import CreateEnvResult
+
+
+with open(ROOT_DIR / "game_content" / "item_info.yaml", "r") as file:
+    item_info = file.read()
+with open(ROOT_DIR / "game_content" / "layouts" / "basic.layout", "r") as file:
+    layout = file.read()
+with open(ROOT_DIR / "game_content" / "environment_config.yaml", "r") as file:
+    environment_config = file.read()
+
+create_env = CreateEnvironmentConfig(
+    manager_id="SECRETKEY1",
+    number_players=2,
+    environment_settings={"all_player_can_pause_game": False},
+    item_info_config=item_info,
+    environment_config=environment_config,
+    layout_config=layout,
+).model_dump(mode="json")
+
+env_info = requests.post("http://localhost:8000/manage/create_env", json=create_env)
+if env_info.status_code == 403:
+    raise ValueError(f"Forbidden Request: {env_info.json()['detail']}")
+env_info: CreateEnvResult = env_info.json()
+```
+
+Connect each player via a websocket (threaded or async).
+```python
+import json
+import dataclasses
+from websockets import connect
+
+from overcooked_simulator.overcooked_environment import Action, ActionType, InterActionData
+from overcooked_simulator.utils import custom_asdict_factory
+
+
+p1_websocket = connect("ws://localhost:8000/ws/player/" + env_info["player_info"]["0"]["client_id"])
+
+# set player "0" as ready
+p1_websocket.send(json.dumps({"type": "ready", "player_hash": env_info["player_info"]["0"]["player_hash"]}))
+assert json.loads(websocket.recv())["status"] == 200, "not accepted player"
+
+
+# get the state for player "0", call every frame/step
+p1_websocket.send(json.dumps({"type": "get_state", "player_hash": env_info["player_info"]["0"]["player_hash"]}))
+state = json.loads(websocket.recv())
+
+# send an action for player "0"
+# --- movement ---
+action = Action(
+    player="0",
+    action_type=ActionType.MOVEMENT,
+    action_data=[0.0, 1.0],  # direction (here straight up)
+    duration=0.5  # seconds
+)
+# --- pickup/drop off ---
+action = Action(
+    player="0",
+    action_type=ActionType.PUT,
+    action_data="pickup",
+)
+# --- interact ---
+action = Action(
+    player="0",
+    action_type=ActionType.INTERACT,
+    action_data=InterActionData.START  # InterActionData.STOP when to stop the interaction
+)
+
+p1_websocket.send(json.dumps({
+    "type": "action",
+    "player_hash": env_info["player_info"]["0"]["player_hash"],
+    "action": dataclasses.asdict(
+        action, dict_factory=custom_asdict_factory
+    ),
+}))
+websocket.recv()
+
+```
+
+Stop the environment if you want the game to end before the time is up.
+```python
+requests.post(
+    "http://localhost:8000/manage/stop_env",
+    json={
+        "manager_id": "SECRETKEY1",
+        "env_id": env_info["env_id"],
+        "reason": "closed environment",
+    },
+)
+```
+
+## Direct integration into your code
+You can use the `overcooked_simulator.overcooked_environment.Environment` class
+and call the `step`, `get_json_state`, and `perform_action` methods directly.
+
+```python
+from datetime import timedelta
+
+from overcooked_simulator import ROOT_DIR
+from overcooked_simulator.overcooked_environment import Action, Environment
+
+env = Environment(
+    env_config=ROOT_DIR / "game_content" / "environment_config.yaml",
+    layout_config=ROOT_DIR / "game_content" / "layouts" / "basic.layout",
+    item_info=ROOT_DIR / "game_content" / "item_info.yaml"
+)
+
+env.add_player("0")
+env.add_player("1")
+
+while True:
+    # adapt this to real time if needed with time.sleep etc.
+    env.step(timedelta(seconds=0.1))
+
+    player_0_state = json.loads(env.get_json_state("0"))
+    if player_0_state["ended"]:
+        break
+
+    action = ...  # See above but use np.array instead of list for the movement direction vector
+    env.perform_action(action)
+```
+
+# JSON State
+The JSON Schema for the state of the environment for a player can be generated by running the `state_representation.py`
+```bash
+python state_representation.py
+```
+Should look like
+```
+{'$defs': {'CookingEquipmentState': {'properties': {'content_list': {'items': {'$ref': '#/$defs/ItemState'}, 'title': 'Content List', 'type': 'array'}, 'content_ready': {'anyOf': [{'$ref': '#/$defs/ItemState'}, {'type': 'null'}]}}, 'required': ['content_list', 'content_ready'], 'title': 'CookingEquipmentState', 'type': 'object'}, 'CounterState': {'properties': {'id': {'title': 'Id', 'type': 'string'}, 'category': {'const': 'Counter', 'title': 'Category'}, 'type': {'title': 'Type', 'type': 'string'}, 'pos': {'items': {'type': 'number'}, 'title': 'Pos', 'type': 'array'}, 'occupied_by': {'anyOf': [{'items': {'anyOf': [{'$ref': '#/$defs/ItemState'}, {'$ref': '#/$defs/CookingEquipmentState'}]}, 'type': 'array'}, {'$ref': '#/$defs/ItemState'}, {'$ref': '#/$defs/CookingEquipmentState'}, {'type': 'null'}], 'title': 'Occupied By'}}, 'required': ['id', 'category', 'type', 'pos', 'occupied_by'], 'title': 'CounterState', 'type': 'object'}, 'ItemState': {'properties': {'id': {'title': 'Id', 'type': 'string'}, 'category': {'anyOf': [{'const': 'Item'}, {'const': 'ItemCookingEquipment'}], 'title': 'Category'}, 'type': {'title': 'Type', 'type': 'string'}, 'progress_percentage': {'anyOf': [{'type': 'number'}, {'type': 'integer'}], 'title': 'Progress Percentage'}}, 'required': ['id', 'category', 'type', 'progress_percentage'], 'title': 'ItemState', 'type': 'object'}, 'KitchenInfo': {'properties': {'width': {'title': 'Width', 'type': 'number'}, 'height': {'title': 'Height', 'type': 'number'}}, 'required': ['width', 'height'], 'title': 'KitchenInfo', 'type': 'object'}, 'OrderState': {'properties': {'id': {'title': 'Id', 'type': 'string'}, 'category': {'const': 'Order', 'title': 'Category'}, 'meal': {'title': 'Meal', 'type': 'string'}, 'start_time': {'format': 'date-time', 'title': 'Start Time', 'type': 'string'}, 'max_duration': {'title': 'Max Duration', 'type': 'number'}}, 'required': ['id', 'category', 'meal', 'start_time', 'max_duration'], 'title': 'OrderState', 'type': 'object'}, 'PlayerState': {'properties': {'id': {'title': 'Id', 'type': 'string'}, 'pos': {'items': {'type': 'number'}, 'title': 'Pos', 'type': 'array'}, 'facing_direction': {'items': {'type': 'number'}, 'title': 'Facing Direction', 'type': 'array'}, 'holding': {'anyOf': [{'$ref': '#/$defs/ItemState'}, {'$ref': '#/$defs/CookingEquipmentState'}, {'type': 'null'}], 'title': 'Holding'}, 'current_nearest_counter_pos': {'anyOf': [{'items': {'type': 'number'}, 'type': 'array'}, {'type': 'null'}], 'title': 'Current Nearest Counter Pos'}, 'current_nearest_counter_id': {'anyOf': [{'type': 'string'}, {'type': 'null'}], 'title': 'Current Nearest Counter Id'}}, 'required': ['id', 'pos', 'facing_direction', 'holding', 'current_nearest_counter_pos', 'current_nearest_counter_id'], 'title': 'PlayerState', 'type': 'object'}}, 'properties': {'players': {'items': {'$ref': '#/$defs/PlayerState'}, 'title': 'Players', 'type': 'array'}, 'counters': {'items': {'$ref': '#/$defs/CounterState'}, 'title': 'Counters', 'type': 'array'}, 'kitchen': {'$ref': '#/$defs/KitchenInfo'}, 'score': {'anyOf': [{'type': 'number'}, {'type': 'integer'}], 'title': 'Score'}, 'orders': {'items': {'$ref': '#/$defs/OrderState'}, 'title': 'Orders', 'type': 'array'}, 'ended': {'title': 'Ended', 'type': 'boolean'}, 'env_time': {'format': 'date-time', 'title': 'Env Time', 'type': 'string'}, 'remaining_time': {'title': 'Remaining Time', 'type': 'number'}}, 'required': ['players', 'counters', 'kitchen', 'score', 'orders', 'ended', 'env_time', 'remaining_time'], 'title': 'StateRepresentation', 'type': 'object'}
+```
+
+The BaseModel and TypedDicts can be found in `overcooked_simulator.state_representation`. The
+`overcooked_simulator.state_representation.StateRepresentation` represents the json state that the `get_json_state`
+returns.
+
+## Generate images from JSON states
+You might have stored some json states and now you want to visualize them via the
+pygame-2d visualization. You can do that by running the `drawing.py` script and referencing a json file.
+```bash
+python3 overcooked_simulator/gui_2d_vis/drawing.py --state my_state.json
+```
+- You can specify a different visualization config with `-v` or `--visualization_config`.
+- You can specify the name of the output file with `-o` or `--output_file`. The default is `screenshot.jpg`.
+
+# Configuration
+
+The environment configuration is currently done with 3 config files + GUI configuration.
+
+## Item Config
+
+The item config defines which ingredients, cooking equipment and meals can exist and how meals and processed ingredients
+can be cooked/created.
+
+For example
+
+    CuttingBoard:
+      type: Equipment
+
+    Stove:
+      type: Equipment
+
+    Pot:
+      type: Equipment
+      equipment: Stove
+
+    Tomato:
+      type: Ingredient
+
+    ChoppedTomato:
+      type: Ingredient
+      needs: [ Tomato ]
+      seconds: 4.0
+      equipment: CuttingBoard
+
+    TomatoSoup:
+      type: Meal
+      needs: [ ChoppedTomato, ChoppedTomato, ChoppedTomato ]
+      seconds: 6.0
+      equipment: Pot
+
+
+## Layout Config
+
+You can define the layout of the kitchen via a layout file. The position of counters are based on a grid system, even
+when the players do not move grid steps but continuous steps. Each character defines a different type of counter. Which
+character is mapped to which counter is defined in the Environment config.
+
+For example
+
+```
+#QU#FO#TNLB#
+#__________M
+#__________K
+W__________I
+#__A_____A_D
+C__________E
+C__________G
+#__________#
+#P#S+#X##S+#
+```
+
+## Environment Config
+
+The environment config defines how a level/environment is defined. Here, the available plates, meals, order and player
+configuration is done.
+
+For example
+
+```yaml
+plates:
+  clean_plates: 1
+  dirty_plates: 2
+  plate_delay: [ 5, 10 ]
+  # range of seconds until the dirty plate arrives.
+
+game:
+  time_limit_seconds: 300
+
+meals:
+  all: true
+
+layout_chars:
+  _: Free
+  hash: Counter
+  A: Agent
+  P: PlateDispenser
+  C: CuttingBoard
+  X: Trashcan
+  W: ServingWindow
+  S: Sink
+  +: SinkAddon
+  U: Pot  # with Stove
+  T: Tomato
+
+orders:
+  ...
+
+player_config:
+  radius: 0.4
+  player_speed_units_per_seconds: 8
+  interaction_range: 1.6
+```
 
-## Direct integration into your code.
-Initialize an environment....
+## PyGame Visualization Config
 
-**TODO** JSON State description.
+Here the visualisation for all objects is defined. Reference the images or define a list of base shapes that represent
+the counters, ingredients, meals and players.
 
 
 # Citation
@@ -54,14 +317,17 @@ Initialize an environment....
 # Structure of the Documentation
 The API documentation follows the file and content structure in the repo.
 On the left you can find the navigation panel that brings you to the implementation of
+- the **counter factory** converts the characters in the layout file to counter instances,
 - the **counters**, including the kitchen utility objects like dispenser, cooking counter (stove, deep fryer, oven),
   sink, etc.,
 - the **game items**, the holdable ingredients, cooking equipment, composed ingredients, and meals,
-- in **main**, you find an example how to start a simulation,
+- the **game server**, which can manage several running environments and can communicates via FastAPI post requests and
+websockets,
 - the **orders**, how to sample incoming orders and their attributes,
 - the **environment**, handles the incoming actions and provides the state,
 - the **player**/agent, that interacts in the environment,
-- a **simulation runner**, that calls the step function of the environment for a real-time interaction, and
+- type hints are defined in **state representation** for the json state and **server results** for the data returned by
+the game server in post requests.
 - **util**ity code.
 
 
diff --git a/overcooked_simulator/__main__.py b/overcooked_simulator/__main__.py
index 7b055adeff6acd29a6912719a181632cb83cc400..436201f2ec085f299699bdbf641f6f01029270b8 100644
--- a/overcooked_simulator/__main__.py
+++ b/overcooked_simulator/__main__.py
@@ -23,6 +23,16 @@ def start_pygame_gui(cli_args):
 
 def main(cli_args=None):
     if cli_args is None:
+        parser = argparse.ArgumentParser(
+            prog="Overcooked Simulator",
+            description="Game Engine Server + PyGameGUI: Starts overcooked game engine server and a PyGame 2D Visualization window in two processes.",
+            epilog="For further information, see https://scs.pages.ub.uni-bielefeld.de/cocosy/overcooked-simulator/overcooked_simulator.html",
+        )
+
+        url_and_port_arguments(parser)
+        disable_websocket_logging_arguments(parser)
+        add_list_of_manager_ids_arguments(parser)
+
         cli_args = parser.parse_args()
     game_server = None
     pygame_gui = None
@@ -40,7 +50,7 @@ def main(cli_args=None):
         print("Received Keyboard interrupt")
     finally:
         if game_server is not None and game_server.is_alive():
-            print("Terminate game server")
+            print("Terminate gparserame server")
             game_server.terminate()
         if pygame_gui is not None and pygame_gui.is_alive():
             print("Terminate pygame gui")
@@ -49,15 +59,4 @@ def main(cli_args=None):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        prog="Overcooked Simulator",
-        description="Game Engine Server + PyGameGUI: Starts overcooked game engine server and a PyGame 2D Visualization window in two processes.",
-        epilog="For further information, see https://scs.pages.ub.uni-bielefeld.de/cocosy/overcooked-simulator/overcooked_simulator.html",
-    )
-
-    url_and_port_arguments(parser)
-    disable_websocket_logging_arguments(parser)
-    add_list_of_manager_ids_arguments(parser)
-
-    args = parser.parse_args()
-    main(args)
+    main()
diff --git a/overcooked_simulator/counter_factory.py b/overcooked_simulator/counter_factory.py
index 8f36b974384b8c78bbf92b004e830dc6769630d1..93193edad4afa17ad4e1c40e39f60a40c880ceaf 100644
--- a/overcooked_simulator/counter_factory.py
+++ b/overcooked_simulator/counter_factory.py
@@ -1,3 +1,36 @@
+"""
+The `CounterFactory` initializes the `Counter` classes from the characters in the layout config.
+The mapping depends on the definition in the `environment_config.yml` in the `layout_chars` section.
+
+```yaml
+layout_chars:
+  _: Free
+  hash: Counter
+  A: Agent
+  P: PlateDispenser
+  C: CuttingBoard
+  X: Trashcan
+  W: ServingWindow
+  S: Sink
+  +: SinkAddon
+  U: Pot  # with Stove
+  Q: Pan  # with Stove
+  O: Peel  # with Oven
+  F: Basket  # with DeepFryer
+  T: Tomato
+  N: Onion  # oNioN
+  L: Lettuce
+  K: Potato  # Kartoffel
+  I: Fish  # fIIIsh
+  D: Dough
+  E: Cheese  # chEEEse
+  G: Sausage  # sausaGe
+  B: Bun
+  M: Meat
+```
+
+# Code Documentation
+"""
 import inspect
 import sys
 from typing import Any, Type, TypeVar
@@ -27,12 +60,15 @@ T = TypeVar("T")
 def convert_words_to_chars(layout_chars_config: dict[str, str]) -> dict[str, str]:
     """Converts words in a given layout chars configuration dictionary to their corresponding characters.
 
+    This is useful for characters that can not be keys in a yaml file. For example, `#` produces a comment.
+    Therefore, you can specify `hash` as a key (`hash: Counter`). `word_refs` defines the conversions. *Click on `▶ View Source`.*
+
     Args:
-        layout_chars_config (dict[str, str]): A dictionary containing layout character configurations, where the keys are words
+        layout_chars_config: A dictionary containing layout character configurations, where the keys are words
             representing layout characters and the values are the corresponding character representations.
 
     Returns:
-        dict[str, str]: A dictionary where the keys are the layout characters and the values are their corresponding words.
+        A dictionary where the keys are the layout characters and the values are their corresponding words.
     """
     word_refs = {
         "hash": "#",
@@ -78,53 +114,57 @@ class CounterFactory:
     ) -> None:
         """Constructor for the `CounterFactory` class. Set up the attributes necessary to instantiate the counters.
 
-        Args:
-            layout_chars_config (dict[str, str]): A dictionary mapping layout characters to their corresponding names.
-            item_info (dict[str, ItemInfo]): A dictionary containing information about different items.
-            serving_window_additional_kwargs (dict[str, Any]): Additional keyword arguments for serving window configuration.
-            plate_config (PlateConfig): The configuration for plate usage.
-
-        Returns:
-            None
-
         Initializes the object with the provided parameters. It performs the following tasks:
         - Converts the layout character configuration from words to characters.
         - Sets the item information dictionary.
         - Sets the additional keyword arguments for serving window configuration.
         - Sets the plate configuration.
 
-        It also sets the following attributes:
-        - `no_counter_chars`: A set of characters that represent counters for agents or free spaces.
-        - `counter_classes`: A dictionary of counter classes imported from the 'overcooked_simulator.counters' module.
-        - `cooking_counter_equipments`: A dictionary mapping cooking counters to the list of equipment items associated with them.
+        Args:
+            layout_chars_config: A dictionary mapping layout characters to their corresponding names.
+            item_info: A dictionary containing information about different items.
+            serving_window_additional_kwargs: Additional keyword arguments for serving window configuration.
+            plate_config: The configuration for plate usage.
         """
-        self.layout_chars_config = convert_words_to_chars(layout_chars_config)
-        self.item_info = item_info
-        self.serving_window_additional_kwargs = serving_window_additional_kwargs
-        self.plate_config = plate_config
-        self.order_and_score = order_and_score
+        self.layout_chars_config: dict[str, str] = convert_words_to_chars(
+            layout_chars_config
+        )
+        """Layout chars to the counter names."""
+        self.item_info: dict[str, ItemInfo] = item_info
+        """All item infos from the `item_info` config."""
+        self.serving_window_additional_kwargs: dict[
+            str, Any
+        ] = serving_window_additional_kwargs
+        """The additional keyword arguments for the serving window."""
+        self.plate_config: PlateConfig = plate_config
+        """The plate config from the `environment_config`"""
+        self.order_and_score: OrderAndScoreManager = order_and_score
+        """The order and score manager to pass to `ServingWindow` and the `Tashcan` which can affect the scores."""
 
-        self.no_counter_chars = set(
+        self.no_counter_chars: set[str] = set(
             c
             for c, name in self.layout_chars_config.items()
             if name in ["Agent", "Free"]
         )
+        """A set of characters that represent counters for agents or free spaces."""
 
-        self.counter_classes = dict(
+        self.counter_classes: dict[str, Type] = dict(
             inspect.getmembers(
                 sys.modules["overcooked_simulator.counters"], inspect.isclass
             )
         )
+        """A dictionary of counter classes imported from the 'overcooked_simulator.counters' module."""
 
-        self.cooking_counter_equipments = {
-            cooking_counter: [
+        self.cooking_counter_equipments: dict[str, set[str]] = {
+            cooking_counter: {
                 equipment
                 for equipment, e_info in self.item_info.items()
                 if e_info.equipment and e_info.equipment.name == cooking_counter
-            ]
+            }
             for cooking_counter, info in self.item_info.items()
             if info.type == ItemType.Equipment and info.equipment is None
         }
+        """A dictionary mapping cooking counters to the list of equipment items associated with them."""
 
     def get_counter_object(self, c: str, pos: npt.NDArray[float]) -> Counter:
         """Create and returns a counter object based on the provided character and position."""
@@ -137,7 +177,9 @@ class CounterFactory:
                 if item_info.equipment.name not in self.counter_classes:
                     return CookingCounter(
                         name=item_info.equipment.name,
-                        cooking_counter_equipments=self.cooking_counter_equipments,
+                        equipments=self.cooking_counter_equipments[
+                            item_info.equipment.name
+                        ],
                         pos=pos,
                         occupied_by=CookingEquipment(
                             name=item_info.name,
diff --git a/overcooked_simulator/counters.py b/overcooked_simulator/counters.py
index 8e76013e7ce405e8215a249122eeea4f9b22da5c..68c44c3384bbad5a29ac6a19736e1f22df521ede 100644
--- a/overcooked_simulator/counters.py
+++ b/overcooked_simulator/counters.py
@@ -3,8 +3,8 @@ what should happen when the agent wants to pick something up from the counter. O
 the `Counter.drop_off` method receives the item what should be put on the counter. Before that the
 `Counter.can_drop_off` method checked if the item can be put on the counter. The progress on Counters or on objects
 on the counters are handled via the Counters. They have the task to delegate the progress call via the `progress`
-method, e.g., the `CuttingBoard.progress`. On which type of counter the progress method is called is currently
-defined in the environment class.
+method, e.g., the `CuttingBoard.progress`. The environment class detects which classes in this module have the
+`progress` method defined and on instances of these classes the progress will be delegated.
 
 Inside the item_info.yaml, equipment needs to be defined. It includes counters that are part of the
 interaction/requirements for the interaction.
@@ -18,6 +18,7 @@ interaction/requirements for the interaction.
     Stove:
       type: Equipment
 
+
 The defined counter classes are:
 - `Counter`
 - `CuttingBoard`
@@ -58,8 +59,10 @@ from overcooked_simulator.game_items import (
 
 
 log = logging.getLogger(__name__)
+"""The logger for this module."""
 
 COUNTER_CATEGORY = "Counter"
+"""The string for the `category` value in the json state representation for all counters."""
 
 
 class Counter:
@@ -81,13 +84,17 @@ class Counter:
             pos: Position of the counter in the environment. 2-element vector.
             occupied_by: The item on top of the counter.
         """
-        self.uuid = uuid.uuid4().hex if uid is None else None
+        self.uuid: str = uuid.uuid4().hex if uid is None else None
+        """A unique id for better tracking in GUIs with assets which instance moved or changed."""
         self.pos: npt.NDArray[float] = pos
+        """The position of the counter."""
         self.occupied_by: Optional[Item] = occupied_by
+        """What is on top of the counter, e.g., `Item`s."""
         self.orientation: npt.NDArray[float] = np.array([0, 1], dtype=float)
-
+        """In what direction the counter is facing."""
     @property
-    def occupied(self):
+    def occupied(self) -> bool:
+        """Is something on top of the counter."""
         return self.occupied_by is not None
 
     def set_orientation(self, orientation: npt.NDArray[float]) -> None:
@@ -158,6 +165,7 @@ class Counter:
         )
 
     def to_dict(self) -> dict:
+        """For the state representation. Only the relevant attributes are put into the dict."""
         return {
             "id": self.uuid,
             "category": COUNTER_CATEGORY,
@@ -175,24 +183,28 @@ class Counter:
 
 
 class CuttingBoard(Counter):
-    """Cutting ingredients on. The requirement in a new object could look like
+    """Cutting ingredients on. The requirement in a new object could look like.
+
+        ChoppedTomato:
+          type: Ingredient
+          needs: [ Tomato ]
+          seconds: 4.0
+          equipment: CuttingBoard
 
-    ```yaml
-    ChoppedTomato:
-      type: Ingredient
-      needs: [ Tomato ]
-      seconds: 4.0
-      equipment: CuttingBoard
-    ```
     The character `C` in the `layout` file represents the CuttingBoard.
     """
 
     def __init__(self, pos: np.ndarray, transitions: dict[str, ItemInfo], **kwargs):
-        self.progressing = False
-        self.transitions = transitions
-        self.inverted_transition_dict = {
+        self.progressing: bool = False
+        """Is a player progressing/cutting on the board."""
+        self.transitions: dict[str, ItemInfo] = transitions
+        """The allowed transitions to a new item. Keys are the resulting items and the `ItemInfo` (value) contains 
+        the needed items in the `need` attribute."""
+        self.inverted_transition_dict: dict[str, ItemInfo] = {
             info.needs[0]: info for name, info in self.transitions.items()
         }
+        """For faster accessing the needed item. Keys are the ingredients that the player can put and chop on the 
+        board."""
         super().__init__(pos=pos, **kwargs)
 
     def progress(self, passed_time: timedelta, now: datetime):
@@ -267,10 +279,16 @@ class ServingWindow(Counter):
         plate_dispenser: PlateDispenser = None,
         **kwargs,
     ):
-        self.order_and_score = order_and_score
-        self.plate_dispenser = plate_dispenser
-        self.meals = meals
-        self.env_time_func = env_time_func
+        self.order_and_score: OrderAndScoreManager = order_and_score
+        """Reference to the OrderAndScoreManager class. It determines which meals can be served and it manages the 
+        score."""
+        self.plate_dispenser: PlateDispenser = plate_dispenser
+        """Served meals are mentioned on the plate dispenser. So that the plate dispenser can add a dirty plate after 
+        some time."""
+        self.meals: set[str] = meals
+        """All allowed meals by the `environment_config.yml`."""
+        self.env_time_func: Callable[[], datetime] = env_time_func
+        """Reference to get the current env time by calling the `env_time_func`."""
         super().__init__(pos=pos, **kwargs)
 
     def drop_off(self, item) -> Item | None:
@@ -299,7 +317,7 @@ class Dispenser(Counter):
 
     At the moment all ingredients have an unlimited stock.
 
-    The character for each dispenser in the `layout` file is currently hard coded in the environment class:
+    The character for each dispenser in the `layout` file is defined in the `environment_config.yml`:
     ```yaml
     T: Tomato
     L: Lettuce
@@ -307,14 +325,14 @@ class Dispenser(Counter):
     B: Bun
     M: Meat
     ```
-    The plan is to put the info also in the config.
 
     In the implementation, an instance of the item to dispense is always on top of the dispenser.
     Which also is easier for the visualization of the dispenser.
     """
 
     def __init__(self, pos: npt.NDArray[float], dispensing: ItemInfo, **kwargs):
-        self.dispensing = dispensing
+        self.dispensing: ItemInfo = dispensing
+        """`ItemInfo` what the the Dispenser is dispensing. One ready object always is on top of the counter."""
         super().__init__(
             pos=pos,
             occupied_by=self.create_item(),
@@ -337,6 +355,7 @@ class Dispenser(Counter):
         return f"{self.dispensing.name}Dispenser"
 
     def create_item(self) -> Item:
+        """Create a new item to put on the dispenser after the previous one was picked up."""
         kwargs = {
             "name": self.dispensing.name,
             "item_info": self.dispensing,
@@ -386,12 +405,20 @@ class PlateDispenser(Counter):
         **kwargs,
     ) -> None:
         super().__init__(pos=pos, **kwargs)
-        self.dispensing = dispensing
-        self.occupied_by = deque()
-        self.out_of_kitchen_timer = []
-        self.plate_config = plate_config
-        self.next_plate_time = datetime.max
-        self.plate_transitions = plate_transitions
+        self.dispensing: ItemInfo = dispensing
+        """Plate ItemInfo."""
+        self.occupied_by: deque = deque()
+        """The queue of plates. New dirty ones are put at the end and therefore under the current plates."""
+        self.out_of_kitchen_timer: list[datetime] = []
+        """Internal timer for how many plates are out of kitchen and how long."""
+        self.plate_config: PlateConfig = plate_config
+        """The config how many plates are present in the kitchen at the beginning (and in total) and the config for 
+        the random "out of kitchen" timer."""
+        self.next_plate_time: datetime = datetime.max
+        """For efficient checking if dirty plates should be created, instead of looping through the 
+        `out_of_kitchen_timer` list every frame."""
+        self.plate_transitions: dict[str, ItemInfo] = plate_transitions
+        """Transitions for the plates. Relevant for the sink, because a plate can become a clean one there."""
         self.setup_plates()
 
     def pick_up(self, on_hands: bool = True) -> Item | None:
@@ -402,8 +429,6 @@ class PlateDispenser(Counter):
         return not self.occupied_by or self.occupied_by[-1].can_combine(item)
 
     def drop_off(self, item: Item) -> Item | None:
-        """At the moment items can be put on the top of the plate dispenser or the top plate if it is clean and can
-        be put on a plate."""
         if not self.occupied_by:
             self.occupied_by.append(item)
         elif self.occupied_by[-1].can_combine(item):
@@ -411,6 +436,7 @@ class PlateDispenser(Counter):
         return None
 
     def add_dirty_plate(self):
+        """Add a dirty plate after a timer is completed."""
         self.occupied_by.appendleft(self.create_item())
 
     def update_plate_out_of_kitchen(self, env_time: datetime):
@@ -464,6 +490,11 @@ class PlateDispenser(Counter):
         return "PlateReturn"
 
     def create_item(self, clean: bool = False) -> Plate:
+        """Create a plate.
+
+        Args:
+            clean: Whether the plate is clean or dirty.
+        """
         kwargs = {
             "clean": clean,
             "transitions": self.plate_transitions,
@@ -482,7 +513,8 @@ class Trashcan(Counter):
         self, order_and_score: OrderAndScoreManager, pos: npt.NDArray[float], **kwargs
     ):
         super().__init__(pos, **kwargs)
-        self.order_and_score = order_and_score
+        self.order_and_score: OrderAndScoreManager = order_and_score
+        """Reference to the `OrderAndScoreManager`, because unnecessary removed items can affect the score."""
 
     def pick_up(self, on_hands: bool = True) -> Item | None:
         pass
@@ -515,19 +547,18 @@ class CookingCounter(Counter):
     def __init__(
         self,
         name: str,
-        cooking_counter_equipments: dict[str, list[str]],
+        equipments: set[str],
         **kwargs,
     ):
-        self.name = name
-        self.cooking_counter_equipments = cooking_counter_equipments
+        self.name: str = name
+        """The type/name of the cooking counter, e.g., Stove, DeepFryer, Oven."""
+        self.equipments: set[str] = equipments
+        """The valid equipment for the cooking counter, e.g., for a Stove {'Pot','Pan'}."""
         super().__init__(**kwargs)
 
     def can_drop_off(self, item: Item) -> bool:
         if self.occupied_by is None:
-            return (
-                isinstance(item, CookingEquipment)
-                and item.name in self.cooking_counter_equipments[self.name]
-            )
+            return isinstance(item, CookingEquipment) and item.name in self.equipments
         else:
             return self.occupied_by.can_combine(item)
 
@@ -536,7 +567,7 @@ class CookingCounter(Counter):
         if (
             self.occupied_by
             and isinstance(self.occupied_by, CookingEquipment)
-            and self.occupied_by.name in self.cooking_counter_equipments[self.name]
+            and self.occupied_by.name in self.equipments
             and self.occupied_by.can_progress()
         ):
             self.occupied_by.progress(passed_time, now)
@@ -571,7 +602,8 @@ class Sink(Counter):
         **kwargs,
     ):
         super().__init__(pos=pos, **kwargs)
-        self.progressing = False
+        self.progressing: bool = False
+        """If a player currently cleans a plate."""
         self.sink_addon: SinkAddon = sink_addon
         """The connected sink addon which will receive the clean plates"""
         self.occupied_by: deque[Plate] = deque()
@@ -588,7 +620,8 @@ class Sink(Counter):
             self.transition_needs.update([info.needs[0]])
 
     @property
-    def occupied(self):
+    def occupied(self) -> bool:
+        """If there is a plate in the sink."""
         return len(self.occupied_by) != 0
 
     def progress(self, passed_time: timedelta, now: datetime):
@@ -631,7 +664,7 @@ class Sink(Counter):
     def can_drop_off(self, item: Item) -> bool:
         return isinstance(item, Plate) and not item.clean
 
-    def drop_off(self, item: Item) -> Item | None:
+    def drop_off(self, item: Plate) -> Item | None:
         self.occupied_by.appendleft(item)
         return None
 
@@ -639,6 +672,7 @@ class Sink(Counter):
         return None
 
     def set_addon(self, sink_addon: SinkAddon):
+        """Set the closest addon in post_setup."""
         self.sink_addon = sink_addon
 
     def to_dict(self) -> dict:
@@ -658,7 +692,8 @@ class SinkAddon(Counter):
     def __init__(self, pos: npt.NDArray[float], occupied_by=None):
         super().__init__(pos=pos)
         # maybe check if occupied by is already a list or deque?
-        self.occupied_by = deque([occupied_by]) if occupied_by else deque()
+        self.occupied_by: deque = deque([occupied_by]) if occupied_by else deque()
+        """The stack of clean plates."""
 
     def can_drop_off(self, item: Item) -> bool:
         return self.occupied_by and self.occupied_by[-1].can_combine(item)
@@ -667,6 +702,7 @@ class SinkAddon(Counter):
         return self.occupied_by[-1].combine(item)
 
     def add_clean_plate(self, plate: Plate):
+        """Called from the `Sink` after a plate is cleaned / the progress is complete."""
         self.occupied_by.appendleft(plate)
 
     def pick_up(self, on_hands: bool = True) -> Item | None:
diff --git a/overcooked_simulator/game_items.py b/overcooked_simulator/game_items.py
index a77f8d78159ee6e6542209230e7bc50a6c0e76df..9a406367ff71e802b5ba31e6d7bf2cd95304f062 100644
--- a/overcooked_simulator/game_items.py
+++ b/overcooked_simulator/game_items.py
@@ -29,15 +29,22 @@ from enum import Enum
 from typing import Optional, TypedDict
 
 log = logging.getLogger(__name__)
+"""The logger for this module."""
 
 ITEM_CATEGORY = "Item"
+"""The string for the `category` value in the json state representation for all normal items."""
+
 COOKING_EQUIPMENT_ITEM_CATEGORY = "ItemCookingEquipment"
+"""The string for the `category` value in the json state representation for all cooking equipments."""
 
 
 class ItemType(Enum):
     Ingredient = "Ingredient"
+    """All ingredients and process ingredients."""
     Meal = "Meal"
+    """All combined ingredients that can be served."""
     Equipment = "Equipment"
+    """All counters and cooking equipments."""
 
 
 @dataclasses.dataclass
@@ -105,15 +112,21 @@ class Item:
     """Base class for game items which can be held by a player."""
 
     item_category = ITEM_CATEGORY
+    """Class dependent category (is changed for the `CookingEquipment` class). """
 
     def __init__(
         self, name: str, item_info: ItemInfo, uid: str = None, *args, **kwargs
     ):
-        self.name = self.__class__.__name__ if name is None else name
-        self.item_info = item_info
-        self.progress_equipment = None
-        self.progress_percentage = 0.0
-        self.uuid = uuid.uuid4().hex if uid is None else uid
+        self.name: str = self.__class__.__name__ if name is None else name
+        """The name of the item, e.g., `Tomato` or `ChoppedTomato`"""
+        self.item_info: ItemInfo = item_info
+        """The information about the item from the `item_info.yml` config."""
+        self.progress_equipment: str | None = None
+        """The equipment with which the item was last progressed."""
+        self.progress_percentage: float = 0.0
+        """The current progress percentage of the item if it is progress-able."""
+        self.uuid: str = uuid.uuid4().hex if uid is None else uid
+        """A unique identifier for the item. Useful for GUIs that handles specific asset instances."""
 
     def __repr__(self):
         if self.progress_equipment is None:
@@ -126,12 +139,15 @@ class Item:
 
     @property
     def extra_repr(self):
+        """Additional string to add to the representation of the item (in __repr__)."""
         return ""
 
     def can_combine(self, other) -> bool:
+        """Check if the item can be combined with the other. After it returned True the `combine` method is called."""
         return False
 
     def combine(self, other) -> Item | None:
+        """Combine the item with another item based on possible transitions/needs."""
         pass
 
     def progress(self, equipment: str, percent: float):
@@ -148,10 +164,12 @@ class Item:
             )
 
     def reset(self):
+        """Reset the progress."""
         self.progress_equipment = None
         self.progress_percentage = 0.0
 
     def to_dict(self) -> dict:
+        """For the state representation. Only the relevant attributes are put into the dict."""
         return {
             "id": self.uuid,
             "category": self.item_category,
@@ -163,11 +181,13 @@ class Item:
 class CookingEquipment(Item):
     """Pot, Pan, ... that can hold items. It holds the progress of the content (e.g., the soup) in itself (
     progress_percentage) and not in the items in the content list."""
-    item_category = "Cooking Equipment"
+
+    item_category = COOKING_EQUIPMENT_ITEM_CATEGORY
 
     def __init__(self, transitions: dict[str, ItemInfo], *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.transitions = transitions
+        self.transitions: dict[str, ItemInfo] = transitions
+        """What is needed to cook a meal / create another ingredient."""
         self.active_transition: Optional[ActiveTransitionTypedDict] = None
         """The info how and when to convert the content_list to a new item."""
 
@@ -227,6 +247,7 @@ class CookingEquipment(Item):
         return return_value
 
     def can_progress(self) -> bool:
+        """Check if the cooking equipment can progress items at all."""
         return self.active_transition is not None
 
     def progress(self, passed_time: datetime.timedelta, now: datetime.datetime):
@@ -239,10 +260,12 @@ class CookingEquipment(Item):
         # todo set active transition for fire/burnt?
 
     def reset_content(self):
+        """Reset the content attributes after the content was picked up from the equipment."""
         self.content_list = []
         self.content_ready = None
 
     def release(self):
+        """Release the content when the player "picks up" the equipment with a plate in the hands"""
         content = self.content_list
         self.reset_content()
         return content
@@ -256,6 +279,7 @@ class CookingEquipment(Item):
         self.active_transition = None
 
     def get_potential_meal(self) -> Item | None:
+        """The meal that could be served depends on the attributes `content_ready` and `content_list`"""
         if self.content_ready:
             return self.content_ready
         if len(self.content_list) == 1:
@@ -281,8 +305,9 @@ class CookingEquipment(Item):
 class Plate(CookingEquipment):
     """The plate can have to states: clean and dirty. In the clean state it can hold content/other items."""
 
-    def __init__(self, transitions, clean, *args, **kwargs):
-        self.clean = clean
+    def __init__(self, transitions: dict[str, ItemInfo], clean: bool, *args, **kwargs):
+        self.clean: bool = clean
+        """If the plate is clean or dirty."""
         self.meals = set(transitions.keys())
         """All meals can be hold by a clean plate"""
         super().__init__(
@@ -292,14 +317,12 @@ class Plate(CookingEquipment):
             **kwargs,
         )
 
-    def finished_call(self):
-        self.clean = True
-        self.name = self.create_name()
-
     def progress(self, equipment: str, percent: float):
         Item.progress(self, equipment, percent)
 
     def create_name(self):
+        """The name depends on the clean or dirty state of the plate. Clean plates are `Plate`, otherwise
+        `DirtyPlate`."""
         return "Plate" if self.clean else "DirtyPlate"
 
     def can_combine(self, other):
diff --git a/overcooked_simulator/game_server.py b/overcooked_simulator/game_server.py
index 95e8e18fb6733f319d85fc58dfcef18471e66521..1848a5329552dc8e0f733fcc3b330fd78ada83dd 100644
--- a/overcooked_simulator/game_server.py
+++ b/overcooked_simulator/game_server.py
@@ -250,7 +250,7 @@ class EnvironmentHandler:
             manager_id in self.manager_envs
             and env_id in self.manager_envs[manager_id]
             and self.envs[env_id].status
-            not in [EnvironmentStatus.STOPPED, Environment.PAUSED]
+            not in [EnvironmentStatus.STOPPED, EnvironmentStatus.PAUSED]
         ):
             self.envs[env_id].status = EnvironmentStatus.PAUSED
 
@@ -266,7 +266,7 @@ class EnvironmentHandler:
             manager_id in self.manager_envs
             and env_id in self.manager_envs[manager_id]
             and self.envs[env_id].status
-            not in [EnvironmentStatus.STOPPED, Environment.PAUSED]
+            not in [EnvironmentStatus.STOPPED, EnvironmentStatus.PAUSED]
         ):
             self.envs[env_id].status = EnvironmentStatus.PAUSED
             self.envs[env_id].last_step_time = time.time_ns()
diff --git a/overcooked_simulator/gui_2d_vis/drawing.py b/overcooked_simulator/gui_2d_vis/drawing.py
index f136ef80b17acff11d21f357f2eb5a9aa73fb764..77f05c05bd5e3e42c837e1210aa21d0b996f0ad0 100644
--- a/overcooked_simulator/gui_2d_vis/drawing.py
+++ b/overcooked_simulator/gui_2d_vis/drawing.py
@@ -1,10 +1,14 @@
+import argparse
 import colorsys
+import json
+import math
 from datetime import datetime, timedelta
 from pathlib import Path
 
 import numpy as np
 import numpy.typing as npt
 import pygame
+import yaml
 from scipy.spatial import KDTree
 
 from overcooked_simulator import ROOT_DIR
@@ -71,10 +75,10 @@ class Visualizer:
         self,
         screen,
         state,
-        width,
-        height,
         grid_size,
     ):
+        width = int(np.ceil(state["kitchen"]["width"] * grid_size))
+        height = int(np.ceil(state["kitchen"]["height"] * grid_size))
         self.draw_background(
             surface=screen,
             width=width,
@@ -83,13 +87,13 @@ class Visualizer:
         )
         self.draw_counters(
             screen,
-            state,
+            state["counters"],
             grid_size,
         )
 
         self.draw_players(
             screen,
-            state,
+            state["players"],
             grid_size,
         )
 
@@ -135,16 +139,17 @@ class Visualizer:
     def draw_players(
         self,
         screen: pygame.Surface,
-        state_dict: dict,
+        players: dict,
         grid_size: float,
     ):
         """Visualizes the players as circles with a triangle for the facing direction.
         If the player holds something in their hands, it is displayed
         Args:            state: The game state returned by the environment.
         """
-        for p_idx, player_dict in enumerate(state_dict["players"]):
+        for p_idx, player_dict in enumerate(players):
             player_dict: PlayerState
             pos = np.array(player_dict["pos"]) * grid_size
+            pos += grid_size / 2  # correct for grid offset
 
             facing = np.array(player_dict["facing_direction"], dtype=float)
 
@@ -208,13 +213,15 @@ class Visualizer:
                 )
 
             if player_dict["current_nearest_counter_pos"]:
-                pos = player_dict["current_nearest_counter_pos"]
+                nearest_pos = (
+                    np.array(player_dict["current_nearest_counter_pos"]) * grid_size
+                )
+
                 pygame.draw.rect(
                     screen,
                     colors[self.player_colors[p_idx]],
                     rect=pygame.Rect(
-                        pos[0] * grid_size - (grid_size // 2),
-                        pos[1] * grid_size - (grid_size // 2),
+                        *nearest_pos,
                         grid_size,
                         grid_size,
                     ),
@@ -292,12 +299,7 @@ class Visualizer:
                     radius = part["radius"] * grid_size
                     color = colors[part["color"]]
 
-                    pygame.draw.circle(
-                        screen,
-                        color,
-                        draw_pos,
-                        radius,
-                    )
+                    pygame.draw.circle(screen, color, draw_pos, radius)
 
     def draw_item(
         self,
@@ -371,11 +373,13 @@ class Visualizer:
         grid_size: float,
     ):
         """Visualize progress of progressing item as a green bar under the item."""
+        pos -= grid_size / 2
+
         bar_height = grid_size * 0.2
         progress_width = percent * grid_size
         progress_bar = pygame.Rect(
-            pos[0] - (grid_size / 2),
-            pos[1] - (grid_size / 2) + grid_size - bar_height,
+            pos[0],
+            pos[1] + grid_size - bar_height,
             progress_width,
             bar_height,
         )
@@ -389,8 +393,11 @@ class Visualizer:
         different parts to be drawn.
         Args:            counter: The counter to visualize.
         """
-        pos = np.array(counter_dict["pos"]) * grid_size
+        pos = np.array(counter_dict["pos"], dtype=float) * grid_size
         counter_type = counter_dict["type"]
+
+        pos += grid_size // 2  # correct for grid offset
+
         self.draw_thing(
             screen,
             pos,
@@ -449,15 +456,15 @@ class Visualizer:
                 scale=item_scale,
             )
 
-    def draw_counters(self, screen: pygame, state, grid_size):
+    def draw_counters(self, screen: pygame, counters, grid_size):
         """Visualizes the counters in the environment.
 
         Args:            state: The game state returned by the environment.
         """
-        for counter in state["counters"]:
+        for counter in counters:
             self.draw_counter(screen, counter, grid_size)
 
-        for counter in state["counters"]:
+        for counter in counters:
             if counter["occupied_by"]:
                 item_pos = counter["pos"]
                 item_scale = 1.0
@@ -482,11 +489,17 @@ class Visualizer:
                     screen,
                     counter["occupied_by"],
                     grid_size,
+                    np.array(counter["pos"]) * grid_size + (grid_size / 2),
                     np.array(item_pos) * grid_size,
                     item_scale,
                 )
 
             if SHOW_COUNTER_CENTERS:
+                pygame.draw.circle(
+                    screen,
+                    colors["green1"],
+                    np.array(counter["pos"]) * grid_size + (grid_size / 2),
+                    3,
                 pos = np.array(counter["pos"]) * grid_size
                 pygame.draw.circle(screen, colors["green1"], pos, 3)
                 pygame.draw.circle(screen, colors["green1"], pos, 3)
@@ -536,17 +549,15 @@ class Visualizer:
                 ),
                 width=2,
             )
-            center = np.array(order_upper_left) + np.array(
-                [grid_size / 2, grid_size / 2]
-            )
+            center = np.array(order_upper_left)
             self.draw_thing(
-                pos=center,
+                pos=center + (grid_size / 2),
                 parts=config["Plate"]["parts"],
                 screen=order_screen,
                 grid_size=grid_size,
             )
             self.draw_item(
-                pos=center,
+                pos=center + (grid_size / 2),
                 item={"type": order["meal"]},
                 plate=True,
                 screen=order_screen,
@@ -562,7 +573,7 @@ class Visualizer:
 
             percentage = order_done_seconds / order["max_duration"]
             self.draw_progress_bar(
-                pos=center,
+                pos=center + (grid_size / 2),
                 percent=percentage,
                 screen=order_screen,
                 grid_size=grid_size,
@@ -574,3 +585,55 @@ class Visualizer:
             orders_height // 2,
         ]
         screen.blit(order_screen, orders_rect)
+
+    def save_state_image(
+        self, grid_size: int, state: dict, filename: str | Path
+    ) -> None:
+        width = int(np.ceil(state["kitchen"]["width"] * grid_size))
+        height = int(np.ceil(state["kitchen"]["height"] * grid_size))
+
+        flags = pygame.HIDDEN
+        screen = pygame.display.set_mode((width, height), flags=flags)
+
+        self.draw_gamescreen(screen, state, grid_size)
+        pygame.image.save(screen, filename)
+
+
+def save_screenshot(state: dict, config: dict, filename: str | Path) -> None:
+    viz = Visualizer(config)
+    viz.create_player_colors(len(state["players"]))
+    pygame.init()
+    pygame.font.init()
+    viz.save_state_image(grid_size=40, state=state, filename=filename)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="Overcooked Simulator Image Generation",
+        description="Generate images for a state in json.",
+        epilog="For further information, see https://scs.pages.ub.uni-bielefeld.de/cocosy/overcooked-simulator/overcooked_simulator.html",
+    )
+    parser.add_argument(
+        "-s",
+        "--state",
+        type=argparse.FileType("r", encoding="UTF-8"),
+        default=ROOT_DIR / "gui_2d_vis" / "sample_state.json",
+    )
+    parser.add_argument(
+        "-v",
+        "--visualization_config",
+        type=argparse.FileType("r", encoding="UTF-8"),
+        default=ROOT_DIR / "gui_2d_vis" / "visualization.yaml",
+    )
+    parser.add_argument(
+        "-o",
+        "--output_file",
+        type=str,
+        default="screenshot.jpg",
+    )
+    args = parser.parse_args()
+    with open(args.visualization_config, "r") as f:
+        viz_config = yaml.safe_load(f)
+    with open(args.state, "r") as f:
+        state = json.load(f)
+    save_screenshot(state, viz_config, args.output_file)
diff --git a/overcooked_simulator/gui_2d_vis/overcooked_gui.py b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
index 3a6d9ceaa4bb4658655d13f6d43d4b9c53272ff1..7922bc2fedf5de237a620e01ef0ad2496e942da9 100644
--- a/overcooked_simulator/gui_2d_vis/overcooked_gui.py
+++ b/overcooked_simulator/gui_2d_vis/overcooked_gui.py
@@ -95,7 +95,6 @@ class PyGameGUI:
         self.request_url = f"http://{url}:{port}"
         self.manager_id = random.choice(manager_ids)
 
-        # TODO cache loaded images?
         with open(ROOT_DIR / "gui_2d_vis" / "visualization.yaml", "r") as file:
             self.visualization_config = yaml.safe_load(file)
 
@@ -128,9 +127,8 @@ class PyGameGUI:
         self.vis.create_player_colors(len(self.player_names))
 
     def get_window_sizes(self, state: dict):
-        counter_positions = np.array([c["pos"] for c in state["counters"]])
-        kitchen_width = counter_positions[:, 0].max() + 0.5
-        kitchen_height = counter_positions[:, 1].max() + 0.5
+        kitchen_width = state["kitchen"]["width"]
+        kitchen_height = state["kitchen"]["height"]
         if self.visualization_config["GameWindow"]["WhatIsFixed"] == "window_width":
             game_width = self.visualization_config["GameWindow"]["size"]
             kitchen_aspect_ratio = kitchen_height / kitchen_width
@@ -349,8 +347,6 @@ class PyGameGUI:
         self.vis.draw_gamescreen(
             self.game_screen,
             state,
-            self.game_width,
-            self.game_height,
             self.grid_size,
         )
 
diff --git a/overcooked_simulator/gui_2d_vis/sample_state.json b/overcooked_simulator/gui_2d_vis/sample_state.json
new file mode 100644
index 0000000000000000000000000000000000000000..78d930781479e3d121f06abd8b5713dbc64edbde
--- /dev/null
+++ b/overcooked_simulator/gui_2d_vis/sample_state.json
@@ -0,0 +1,545 @@
+{
+  "players": [
+    {
+      "id": "0",
+      "pos": [
+        3.0,
+        4.0
+      ],
+      "facing_direction": [
+        0,
+        1
+      ],
+      "holding": null,
+      "current_nearest_counter_pos": null,
+      "current_nearest_counter_id": null
+    },
+    {
+      "id": "1",
+      "pos": [
+        9.0,
+        4.0
+      ],
+      "facing_direction": [
+        0,
+        1
+      ],
+      "holding": null,
+      "current_nearest_counter_pos": null,
+      "current_nearest_counter_id": null
+    }
+  ],
+  "counters": [
+    {
+      "id": "b032a0fcf1c54828a7bf3afc189547a6",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        0.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "7b46f1d45f1e463494940d19f47648d2",
+      "category": "Counter",
+      "type": "Stove",
+      "pos": [
+        1.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "240eb5d4e4dd4215a0ab9fbaa93f02f5",
+        "category": "ItemCookingEquipment",
+        "type": "Pan",
+        "progress_percentage": 0.0,
+        "content_list": [],
+        "content_ready": null
+      }
+    },
+    {
+      "id": "9675528d610c4a90b8cdea2f52203151",
+      "category": "Counter",
+      "type": "Stove",
+      "pos": [
+        2.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "54124ff47f9a4b2abb6e2df96a9db919",
+        "category": "ItemCookingEquipment",
+        "type": "Pot",
+        "progress_percentage": 0.0,
+        "content_list": [],
+        "content_ready": null
+      }
+    },
+    {
+      "id": "08d334733913488789fa7d195405c931",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        3.0,
+        0.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "abace5244d92463693c5d57cd70ab9da",
+      "category": "Counter",
+      "type": "DeepFryer",
+      "pos": [
+        4.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "987ea1da780c44cabc06883f5c721414",
+        "category": "ItemCookingEquipment",
+        "type": "Basket",
+        "progress_percentage": 0.0,
+        "content_list": [],
+        "content_ready": null
+      }
+    },
+    {
+      "id": "6fd0f625ee2b43609d10b9163c6fbfcf",
+      "category": "Counter",
+      "type": "Oven",
+      "pos": [
+        5.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "a8609b417c0a4c1f8a0d10ef042f2df6",
+        "category": "ItemCookingEquipment",
+        "type": "Peel",
+        "progress_percentage": 0.0,
+        "content_list": [],
+        "content_ready": null
+      }
+    },
+    {
+      "id": "d5e5f842e2fb490f9ac7b837cca583f2",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        6.0,
+        0.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "3b054bbd2d5f49a39fa7d1bf10801575",
+      "category": "Counter",
+      "type": "TomatoDispenser",
+      "pos": [
+        7.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "8d431317da8a439ab25a69f592d7e30d",
+        "category": "Item",
+        "type": "Tomato",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "0b88cdca1637439db60152762f9c535c",
+      "category": "Counter",
+      "type": "OnionDispenser",
+      "pos": [
+        8.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "c2143424f8424e33aee383396a00f50b",
+        "category": "Item",
+        "type": "Onion",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "7608477a75cd4e4db83add343f721b10",
+      "category": "Counter",
+      "type": "LettuceDispenser",
+      "pos": [
+        9.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "0301dd77d05044ab80d7bc45808546e9",
+        "category": "Item",
+        "type": "Lettuce",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "7510d5c4af424d3d988de020a80a3c63",
+      "category": "Counter",
+      "type": "BunDispenser",
+      "pos": [
+        10.0,
+        0.0
+      ],
+      "occupied_by": {
+        "id": "8c2e04c5828447088f8553f52c849132",
+        "category": "Item",
+        "type": "Bun",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "b4a03b6cdc8a4a088f46ceb68eaa0eff",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        11.0,
+        0.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "9396b59b16834e4ca594bcd7fe91f74f",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        1.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "1bf2e112845442e09ce33f2807981c3a",
+      "category": "Counter",
+      "type": "MeatDispenser",
+      "pos": [
+        11.0,
+        1.0
+      ],
+      "occupied_by": {
+        "id": "d723a509ef6740538c42d0f491c2cdac",
+        "category": "Item",
+        "type": "Meat",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "fa6e820423ca41a5874f12e71419cf2a",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        2.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "9525e23ec15445bb86962ead0b6f1076",
+      "category": "Counter",
+      "type": "PotatoDispenser",
+      "pos": [
+        11.0,
+        2.0
+      ],
+      "occupied_by": {
+        "id": "5f52f039c4f549418ee0be66ca98bee7",
+        "category": "Item",
+        "type": "Potato",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "638e272aeaa14bf8b6287a8a05a27ab8",
+      "category": "Counter",
+      "type": "ServingWindow",
+      "pos": [
+        0.0,
+        3.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "5df4d8b89e01407facfb309c253d9b61",
+      "category": "Counter",
+      "type": "FishDispenser",
+      "pos": [
+        11.0,
+        3.0
+      ],
+      "occupied_by": {
+        "id": "f5f4ac928e4845b7a1b31d84861d1f49",
+        "category": "Item",
+        "type": "Fish",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "b1ab927ccbe140988bca5fe51694c299",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        4.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "e3592afac20e4247b15641eae1818366",
+      "category": "Counter",
+      "type": "DoughDispenser",
+      "pos": [
+        11.0,
+        4.0
+      ],
+      "occupied_by": {
+        "id": "b7d7d308359849e6bbd2fe7baa6d87a8",
+        "category": "Item",
+        "type": "Dough",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "9d4a6e4726094f49873888bc74faf530",
+      "category": "Counter",
+      "type": "CuttingBoard",
+      "pos": [
+        0.0,
+        5.0
+      ],
+      "occupied_by": null,
+      "progressing": false
+    },
+    {
+      "id": "c436371aeab84e4ab47454b89a3f4cb2",
+      "category": "Counter",
+      "type": "CheeseDispenser",
+      "pos": [
+        11.0,
+        5.0
+      ],
+      "occupied_by": {
+        "id": "e2fd82e3fab84dc59ab8b6c6c7963379",
+        "category": "Item",
+        "type": "Cheese",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "2ca7c251080c404787ef46e08e574290",
+      "category": "Counter",
+      "type": "CuttingBoard",
+      "pos": [
+        0.0,
+        6.0
+      ],
+      "occupied_by": null,
+      "progressing": false
+    },
+    {
+      "id": "9c824c5009814894918c17ea951e34b5",
+      "category": "Counter",
+      "type": "SausageDispenser",
+      "pos": [
+        11.0,
+        6.0
+      ],
+      "occupied_by": {
+        "id": "d6d343968b9b4175a4e621f1febaaf1f",
+        "category": "Item",
+        "type": "Sausage",
+        "progress_percentage": 0.0
+      }
+    },
+    {
+      "id": "2c13a76d110749b39ff62866c658ae3e",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        7.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "14b547181dd346daac28467b07277a3d",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        11.0,
+        7.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "188db4c4275f4ba7835989333b252b26",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        0.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "f0a4790c290144cba2942705c3240a0f",
+      "category": "Counter",
+      "type": "PlateDispenser",
+      "pos": [
+        1.0,
+        8.0
+      ],
+      "occupied_by": [
+        {
+          "id": "db41a69e880d4d64b535b265199d4148",
+          "category": "ItemCookingEquipment",
+          "type": "DirtyPlate",
+          "progress_percentage": 0.0,
+          "content_list": [],
+          "content_ready": null
+        },
+        {
+          "id": "13fd4df8399d48c490091bb9ed3e1302",
+          "category": "ItemCookingEquipment",
+          "type": "DirtyPlate",
+          "progress_percentage": 0.0,
+          "content_list": [],
+          "content_ready": null
+        },
+        {
+          "id": "5075019e20ed45c1aa7aa01ca5198c50",
+          "category": "ItemCookingEquipment",
+          "type": "Plate",
+          "progress_percentage": 0.0,
+          "content_list": [],
+          "content_ready": null
+        }
+      ]
+    },
+    {
+      "id": "439b6269458c42f5b8be157da1bc0e22",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        2.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "6a92b071a1db46bf80d82a413ea400d4",
+      "category": "Counter",
+      "type": "Sink",
+      "pos": [
+        3.0,
+        8.0
+      ],
+      "occupied_by": [],
+      "progressing": false
+    },
+    {
+      "id": "b3852e97d88f42d4b724d4fb9b0cf8d8",
+      "category": "Counter",
+      "type": "SinkAddon",
+      "pos": [
+        4.0,
+        8.0
+      ],
+      "occupied_by": []
+    },
+    {
+      "id": "9f13c6729dd647f39575d1654d67a094",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        5.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "0dc1e47f4a6448cb9e457c0d998f25f6",
+      "category": "Counter",
+      "type": "Trashcan",
+      "pos": [
+        6.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "87e13494ebaf44318f3ea05ef9d137a6",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        7.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "6b83337a91414633a0793dfa4086f442",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        8.0,
+        8.0
+      ],
+      "occupied_by": null
+    },
+    {
+      "id": "355025b2ec144d008264009761caaee7",
+      "category": "Counter",
+      "type": "Sink",
+      "pos": [
+        9.0,
+        8.0
+      ],
+      "occupied_by": [],
+      "progressing": false
+    },
+    {
+      "id": "3a12e9b62b8c4473b7b84f307cf3de4b",
+      "category": "Counter",
+      "type": "SinkAddon",
+      "pos": [
+        10.0,
+        8.0
+      ],
+      "occupied_by": []
+    },
+    {
+      "id": "b3ce6a85df3d40c0b376a61c3346fd09",
+      "category": "Counter",
+      "type": "Counter",
+      "pos": [
+        11.0,
+        8.0
+      ],
+      "occupied_by": null
+    }
+  ],
+  "kitchen": {
+    "width": 12.0,
+    "height": 9.0
+  },
+  "score": 0,
+  "orders": [
+    {
+      "id": "823ab5e3211840778261c405ff7df91b",
+      "category": "Order",
+      "meal": "FriedFish",
+      "start_time": "2000-01-01T00:00:00",
+      "max_duration": 47.371306
+    },
+    {
+      "id": "2cb614440155426cae04b212f8f2729f",
+      "category": "Order",
+      "meal": "FriedFish",
+      "start_time": "2000-01-01T00:00:00",
+      "max_duration": 53.237394
+    }
+  ],
+  "ended": false,
+  "env_time": "2000-01-01T00:00:00.675141",
+  "remaining_time": 299.324859
+}
diff --git a/overcooked_simulator/gui_2d_vis/visualization.yaml b/overcooked_simulator/gui_2d_vis/visualization.yaml
index a4fea62b941bc7dd8e3328967cb3ec3f7f1f4217..ce510e721a0b0e32a515e26cc8acf6a90ea1fabf 100644
--- a/overcooked_simulator/gui_2d_vis/visualization.yaml
+++ b/overcooked_simulator/gui_2d_vis/visualization.yaml
@@ -2,7 +2,7 @@
 
 GameWindow:
   WhatIsFixed: grid  # grid or window_width or window_height
-  size: 60
+  size: 50
   screen_margin: 100
   min_width: 700
   min_height: 600
diff --git a/overcooked_simulator/order.py b/overcooked_simulator/order.py
index b44d014d24a1df78cefcc1b5cf55b536f166b465..b61c9218f7adce6b059933eaf27d66c9bb9e8ee1 100644
--- a/overcooked_simulator/order.py
+++ b/overcooked_simulator/order.py
@@ -36,6 +36,9 @@ For an easier usage of the random orders, also some classes for type hints and d
 - `ScoreCalcGenFuncType`
 - `ExpiredPenaltyFuncType`
 
+For the scoring of using the trashcan the `penalty_for_each_item` example function is defined. You can set/replace it
+in the `environment_config`.
+
 
 ## Code Documentation
 """
@@ -53,8 +56,10 @@ from typing import Callable, Tuple, Any, Deque, Protocol, TypedDict
 from overcooked_simulator.game_items import Item, Plate, ItemInfo
 
 log = logging.getLogger(__name__)
+"""The logger for this module."""
 
 ORDER_CATEGORY = "Order"
+"""The string for the `category` value in the json state representation for all orders."""
 
 
 @dataclasses.dataclass
@@ -76,7 +81,7 @@ class Order:
     expired_penalty: float
     """The penalty to the score if the order expires"""
     uuid: str = dataclasses.field(default_factory=lambda: uuid.uuid4().hex)
-
+    """The unique identifier for the order."""
     finished_info: dict[str, Any] = dataclasses.field(default_factory=dict)
     """Is set after the order is completed."""
     _timed_penalties: list[Tuple[datetime, float]] = dataclasses.field(
@@ -84,10 +89,9 @@ class Order:
     )
     """Converted penalties the env is working with from the `timed_penalties`"""
 
-    def order_time(self, env_time: datetime) -> timedelta:
-        return self.start_time - env_time
-
     def create_penalties(self, env_time: datetime):
+        """Create the general timed penalties list to check for penalties after some time the order is still not
+        fulfilled."""
         for penalty_info in self.timed_penalties:
             match penalty_info:
                 case (offset, penalty):
@@ -116,6 +120,7 @@ class OrderGeneration:
 
     def __init__(self, available_meals: dict[str, ItemInfo], **kwargs):
         self.available_meals: list[ItemInfo] = list(available_meals.values())
+        """Available meals restricted through the `environment_config.yml`."""
 
     @abstractmethod
     def init_orders(self, now) -> list[Order]:
@@ -138,10 +143,12 @@ class OrderAndScoreManager:
     """The Order and Score Manager that is called from the serving window."""
 
     def __init__(self, order_config, available_meals: dict[str, ItemInfo]):
-        self.score = 0
+        self.score: float = 0.0
+        """The current score of the environment."""
         self.order_gen: OrderGeneration = order_config["order_gen_class"](
             available_meals=available_meals, kwargs=order_config["order_gen_kwargs"]
         )
+        """The order generation."""
         self.serving_not_ordered_meals: Callable[
             [Item], Tuple[bool, float]
         ] = order_config["serving_not_ordered_meals"]
@@ -170,6 +177,7 @@ class OrderAndScoreManager:
         """Cache last expired orders for `OrderGeneration.get_orders` call."""
 
     def update_next_relevant_time(self):
+        """For more efficient checking when to do something in the progress call."""
         next_relevant_time = datetime.max
         for order in self.open_orders:
             next_relevant_time = min(
@@ -180,6 +188,8 @@ class OrderAndScoreManager:
         self.next_relevant_time = next_relevant_time
 
     def serve_meal(self, item: Item, env_time: datetime) -> bool:
+        """Is called by the ServingWindow to serve a meal. Returns True if the meal can be served and should be
+        "deleted" from the hands of the player."""
         if isinstance(item, Plate):
             meal = item.get_potential_meal()
             if meal is not None:
@@ -220,6 +230,7 @@ class OrderAndScoreManager:
         return False
 
     def increment_score(self, score: int | float):
+        """Add a value to the current score and log it."""
         self.score += score
         log.debug(f"Score: {self.score}")
 
@@ -271,6 +282,8 @@ class OrderAndScoreManager:
             self.update_next_relevant_time()
 
     def find_order_for_meal(self, meal) -> Tuple[Order, int] | None:
+        """Get the order that will be fulfilled for a meal. At the moment the oldest order in the list that has the
+        same meal (name)."""
         for index, order in enumerate(self.open_orders):
             if order.meal.name == meal.name:
                 return order, index
@@ -282,6 +295,7 @@ class OrderAndScoreManager:
             order.create_penalties(env_time)
 
     def order_state(self) -> list[dict]:
+        """Similar to the `to_dict` in `Item` and `Counter`. Relevant for the state of the environment"""
         return [
             {
                 "id": order.uuid,
@@ -294,6 +308,7 @@ class OrderAndScoreManager:
         ]
 
     def apply_penalty_for_using_trash(self, remove: Item | list[Item]):
+        """Is called if a item is put into the trashcan."""
         self.increment_score(self.penalty_for_trash(remove))
 
 
@@ -463,7 +478,7 @@ class RandomOrderGeneration(OrderGeneration):
         super().__init__(available_meals, **kwargs)
         self.kwargs: RandomOrderKwarg = RandomOrderKwarg(**kwargs["kwargs"])
         self.next_order_time: datetime | None = datetime.max
-        self.number_cur_orders = 0
+        self.number_cur_orders: int = 0
         self.needed_orders: int = 0
         """For the sample on dur but when it was restricted due to max order number."""
 
diff --git a/overcooked_simulator/overcooked_environment.py b/overcooked_simulator/overcooked_environment.py
index 77b9e625cff1bae071ebaa85c4b802a65788acc7..44b8ad43757b34715231b7ff37575deac0849b7d 100644
--- a/overcooked_simulator/overcooked_environment.py
+++ b/overcooked_simulator/overcooked_environment.py
@@ -162,6 +162,9 @@ class Environment:
             self.free_positions,
         ) = self.parse_layout_file()
 
+        self.world_borders_x = [-0.5, self.kitchen_width - 0.5]
+        self.world_borders_y = [-0.5, self.kitchen_height - 0.5]
+
         progress_counter_classes = list(
             filter(
                 lambda cl: hasattr(cl, "progress"),
@@ -281,25 +284,24 @@ class Environment:
         [counter_size/2, counter_size/2], counters are directly next to each other (of no empty space is specified
         in layout).
         """
-        current_y: float = 0.5
+
+        starting_at: float = 0.0
+        current_y: float = starting_at
         counters: list[Counter] = []
         designated_player_positions: list[npt.NDArray] = []
         free_positions: list[npt.NDArray] = []
 
-        self.kitchen_width = 0
-
         if self.as_files:
             with open(self.layout_config, "r") as layout_file:
                 lines = layout_file.readlines()
         else:
             lines = self.layout_config.split("\n")
-        self.kitchen_height = len(lines)
 
         grid = []
 
         for line in lines:
             line = line.replace("\n", "").replace(" ", "")  # remove newline char
-            current_x = 0.5
+            current_x: float = starting_at
             grid_line = []
             for character in line:
                 character = character.capitalize()
@@ -322,13 +324,12 @@ class Environment:
                             free_positions.append(np.array([current_x, current_y]))
 
                 current_x += 1
-                if current_x > self.kitchen_width:
-                    self.kitchen_width = current_x
 
             grid.append(grid_line)
             current_y += 1
 
-        self.kitchen_width -= 0.5
+        self.kitchen_width: float = len(lines[0]) + starting_at
+        self.kitchen_height = len(lines) + starting_at
 
         self.determine_counter_orientations(
             counters, grid, np.array([self.kitchen_width / 2, self.kitchen_height / 2])
@@ -634,9 +635,13 @@ class Environment:
 
         Returns: True if the player touches the world bounds, False if not.
         """
-        collisions_lower = any((player.pos - (player.radius)) < 0)
+        collisions_lower = any(
+            (player.pos - (player.radius))
+            < [self.world_borders_x[0], self.world_borders_y[0]]
+        )
         collisions_upper = any(
-            (player.pos + (player.radius)) > [self.kitchen_width, self.kitchen_height]
+            (player.pos + (player.radius))
+            > [self.world_borders_x[1], self.world_borders_y[1]]
         )
         return collisions_lower or collisions_upper
 
@@ -675,6 +680,7 @@ class Environment:
         state = {
             "players": [p.to_dict() for p in self.players.values()],
             "counters": [c.to_dict() for c in self.counters],
+            "kitchen": {"width": self.kitchen_width, "height": self.kitchen_height},
             "score": self.order_and_score.score,
             "orders": self.order_and_score.order_state(),
             "ended": self.game_ended,
diff --git a/overcooked_simulator/player.py b/overcooked_simulator/player.py
index 4c7257ddb13a1d7e2ca4f28ca8358659ce10643b..5f87b87f7032e54f12cace9fecda1d882aac5c54 100644
--- a/overcooked_simulator/player.py
+++ b/overcooked_simulator/player.py
@@ -20,6 +20,7 @@ from overcooked_simulator.game_items import Item, Plate
 from overcooked_simulator.state_representation import PlayerState
 
 log = logging.getLogger(__name__)
+"""The logger for this module."""
 
 
 @dataclasses.dataclass
@@ -79,9 +80,13 @@ class Player:
         calculated with."""
 
         self.current_movement: npt.NDArray[float] = np.zeros(2, float)
+        """The movement vector that will be used to calculate the movement in the next step call."""
         self.movement_until: datetime.datetime = datetime.datetime.min
+        """The env time until the player wants to move."""
 
     def set_movement(self, move_vector, move_until):
+        """Called by the `perform_action` method. Movements will be performed (pos will be updated) in the `step`
+        function of the environment"""
         self.current_movement = move_vector
         self.movement_until = move_until
 
@@ -127,7 +132,8 @@ class Player:
         Args:
             counter: The counter, can the player reach it?
 
-        Returns: True if the counter is in range of the player, False if not.
+        Returns:
+            True if the counter is in range of the player, False if not.
 
         """
         return np.linalg.norm(counter.pos - self.facing_point) <= self.interaction_range
@@ -182,6 +188,7 @@ class Player:
         return f"Player(name:{self.name},pos:{str(self.pos)},holds:{self.holding})"
 
     def to_dict(self) -> PlayerState:
+        """For the state representation. Only the relevant attributes are put into the dict."""
         # TODO add color to player class for vis independent player color
         return {
             "id": self.name,
diff --git a/overcooked_simulator/state_representation.py b/overcooked_simulator/state_representation.py
index 1912be8a9b530ac50b41c9648c9f50e16fd3fe31..c57c30c0812c003d8008b9347a03a2c96b00b39c 100644
--- a/overcooked_simulator/state_representation.py
+++ b/overcooked_simulator/state_representation.py
@@ -1,3 +1,6 @@
+"""
+Type hint classes for the representation of the json state.
+"""
 from datetime import datetime
 
 from pydantic import BaseModel
@@ -51,9 +54,19 @@ class PlayerState(TypedDict):
     current_nearest_counter_id: str | None
 
 
+class KitchenInfo(BaseModel):
+    """Basic information of the kitchen."""
+
+    width: float
+    height: float
+
+
 class StateRepresentation(BaseModel):
+    """The format of the returned state representation."""
+
     players: list[PlayerState]
     counters: list[CounterState]
+    kitchen: KitchenInfo
     score: float | int
     orders: list[OrderState]
     ended: bool
diff --git a/overcooked_simulator/utils.py b/overcooked_simulator/utils.py
index edf11d72e8b4a5cbaa8672d9da727d6e4d3c312f..2754f08912d31ff513ed891646d29571800e8960 100644
--- a/overcooked_simulator/utils.py
+++ b/overcooked_simulator/utils.py
@@ -1,3 +1,7 @@
+"""
+Some utility functions.
+"""
+
 import logging
 import os
 import sys
@@ -36,6 +40,8 @@ def get_closest(point: npt.NDArray[float], counters: list[Counter]):
 
 
 def custom_asdict_factory(data):
+    """Convert enums to their value."""
+
     def convert_value(obj):
         if isinstance(obj, Enum):
             return obj.value
diff --git a/tests/test_start.py b/tests/test_start.py
index 3bb5f3a4e09ced5e05dd764bc74d6d4b6878e760..1dbdafdafd48a79f29969cc686bc87665562799e 100644
--- a/tests/test_start.py
+++ b/tests/test_start.py
@@ -146,7 +146,7 @@ def test_collision_detection(env_config, layout_config, item_info):
     player2.move_abs(counter_pos)
     assert env.detect_player_collision(player1), "Players at same pos. Not detected."
 
-    player1.move_abs(np.array([0, 0]))
+    player1.move_abs(np.array([-1, -1]))
     assert env.detect_collision_world_bounds(
         player1
     ), "Player collides with world bounds."
@@ -234,8 +234,8 @@ def test_processing(env_config, layout_config, item_info):
                 equipment=ItemInfo(name="CuttingBoard", type=ItemType.Equipment),
                 type=ItemType.Ingredient,
                 needs=["Tomato"],
-                )
-            },
+            )
+        },
     )
     env.counters.append(counter)
     env.progressing_counters.append(counter)