diff --git a/cooperative_cuisine/configs/agents/random_agent.py b/cooperative_cuisine/configs/agents/random_agent.py
index e72bd4b67a39a83087c7f9ddc496d93c49c596e1..4aa85d8a22d37217ca410b78d7000320bf8ed23b 100644
--- a/cooperative_cuisine/configs/agents/random_agent.py
+++ b/cooperative_cuisine/configs/agents/random_agent.py
@@ -34,7 +34,7 @@ def get_free_neighbours(
     free_space = np.ones((width, height), dtype=bool)
     for counter in state["counters"]:
         grid_idx = np.array(counter["pos"]).astype(int)
-        free_space[*grid_idx] = False
+        free_space[grid_idx[0], grid_idx[1]] = False
     i, j = np.array(counter_pos).astype(int)
     free = []
 
diff --git a/cooperative_cuisine/state_representation.py b/cooperative_cuisine/state_representation.py
index cbc70c51642cfa9994cbcbffe6cd00f50e5f36d2..fdf122182d9ce50abb6db751bb0c75fca7afba3f 100644
--- a/cooperative_cuisine/state_representation.py
+++ b/cooperative_cuisine/state_representation.py
@@ -13,8 +13,6 @@ from networkx import Graph
 from pydantic import BaseModel
 from typing_extensions import Literal, TypedDict
 
-from cooperative_cuisine import ROOT_DIR
-
 
 class OrderState(TypedDict):
     """Format of the state representation of an order."""
@@ -281,48 +279,48 @@ def create_json_schema() -> dict[str, Any]:
 
 
 if __name__ == "__main__":
-    sample_state_path = ROOT_DIR / "pygame_2d_vis" / "sample_state.json"
-    with open(sample_state_path, "r") as f:
-        state = json.load(f)
-
-    graph = create_movement_graph(state, diagonal=False)
-    width, height = state["kitchen"]["width"], state["kitchen"]["height"]
-    free_space = np.ones((width, height), dtype=bool)
-    for counter in state["counters"]:
-        grid_idx = np.array(counter["pos"]).round().astype(int)
-        free_space[*grid_idx] = False
-
-    other_players = [[2, 2], [3, 3]]
-    restricted = restrict_movement_graph(graph, other_players)
-
-    print(graph.nodes)
-
-    source = (1, 1)
-    target = (10, 7)
-
-    try:
-        path = networkx.astar_path(
-            G=restricted, source=source, target=target, heuristic=astar_heuristic
-        )
-    except networkx.exception.NetworkXNoPath:
-        print("NO PATH FOUND")
-        path = []
-
-    width, height = free_space.shape
-    for i in range(width):
-        for j in range(height):
-            if (i, j) == source:
-                print(" S ", end="")
-            elif (i, j) == target:
-                print(" T ", end="")
-            elif (i, j) in path:
-                print(" x ", end="")
-            elif [i, j] in other_players:
-                print(" O ", end="")
-            elif free_space[i, j]:
-                print("   ", end="")
-            else:
-                print(" # ", end="")
-        print()
-
-    # print(json.dumps(create_json_schema()))
+    # sample_state_path = ROOT_DIR / "pygame_2d_vis" / "sample_state.json"
+    # with open(sample_state_path, "r") as f:
+    #     state = json.load(f)
+    #
+    # graph = create_movement_graph(state, diagonal=False)
+    # width, height = state["kitchen"]["width"], state["kitchen"]["height"]
+    # free_space = np.ones((width, height), dtype=bool)
+    # for counter in state["counters"]:
+    #     grid_idx = np.array(counter["pos"]).round().astype(int)
+    #     free_space[grid_idx[0], grid_idx[1]] = False
+    #
+    # other_players = [[2, 2], [3, 3]]
+    # restricted = restrict_movement_graph(graph, other_players)
+    #
+    # print(graph.nodes)
+    #
+    # source = (1, 1)
+    # target = (10, 7)
+    #
+    # try:
+    #     path = networkx.astar_path(
+    #         G=restricted, source=source, target=target, heuristic=astar_heuristic
+    #     )
+    # except networkx.exception.NetworkXNoPath:
+    #     print("NO PATH FOUND")
+    #     path = []
+    #
+    # width, height = free_space.shape
+    # for i in range(width):
+    #     for j in range(height):
+    #         if (i, j) == source:
+    #             print(" S ", end="")
+    #         elif (i, j) == target:
+    #             print(" T ", end="")
+    #         elif (i, j) in path:
+    #             print(" x ", end="")
+    #         elif [i, j] in other_players:
+    #             print(" O ", end="")
+    #         elif free_space[i, j]:
+    #             print("   ", end="")
+    #         else:
+    #             print(" # ", end="")
+    #     print()
+
+    print(json.dumps(create_json_schema()))