diff --git a/cooperative_cuisine/configs/agents/random_agent.py b/cooperative_cuisine/configs/agents/random_agent.py
index 4aa85d8a22d37217ca410b78d7000320bf8ed23b..99f8a2384625318f6293b4ab21a0c706615b45a0 100644
--- a/cooperative_cuisine/configs/agents/random_agent.py
+++ b/cooperative_cuisine/configs/agents/random_agent.py
@@ -23,7 +23,7 @@ from cooperative_cuisine.utils import custom_asdict_factory
 TIME_TO_STOP_ACTION = 3.0
 
 ADD_RANDOM_MOVEMENTS = False
-DIAGONAL_MOVEMENTS = False
+DIAGONAL_MOVEMENTS = True
 AVOID_OTHER_PLAYERS = True
 
 
@@ -307,9 +307,7 @@ async def agent():
                         ...
 
             if not task_type:
-                # task_type = random.choice(["GOTO"])
-                task_type = random.choice(["GOTO", "PUT"])
-                # task_type = random.choice(["GOTO", "PUT", "INTERACT"])
+                task_type = random.choice(["GOTO", "PUT", "INTERACT"])
                 threshold = datetime.now() + timedelta(seconds=TIME_TO_STOP_ACTION)
                 if task_type == "GOTO":
                     # counter_type = random.choice(list(counters.keys()))
diff --git a/cooperative_cuisine/state_representation.py b/cooperative_cuisine/state_representation.py
index fdf122182d9ce50abb6db751bb0c75fca7afba3f..7887d0f8a7d542bded6fb91ddaf1c7632843c8ee 100644
--- a/cooperative_cuisine/state_representation.py
+++ b/cooperative_cuisine/state_representation.py
@@ -191,19 +191,19 @@ class StateRepresentation(BaseModel):
 
 
 def astar_heuristic(x, y):
-    """Heuristic distance function used in astart algorithm."""
+    """Heuristic distance function used in astar algorithm."""
     return np.linalg.norm(np.array(x) - np.array(y))
 
 
 def create_movement_graph(state: StateRepresentation, diagonal=True) -> Graph:
     """
-    Creates a graph which represents the connections of empty kitchen tiles.
+    Creates a graph which represents the connections of empty kitchen tiles and such
+    possible coarse movements of an agent.
     Args:
         state: State representation to determine the graph to.
-        diagonal: if True use 8 way connection, ie diagonal connections between the spaces.
+        diagonal: if True use 8 way connection, i.e. diagonal connections between the spaces.
 
     Returns: Graph representing the connections between empty kitchen tiles.
-
     """
     width, height = state["kitchen"]["width"], state["kitchen"]["height"]
     free_space = np.ones((width, height), dtype=bool)
@@ -257,7 +257,7 @@ def restrict_movement_graph(
     graph: Graph,
     player_positions: list[tuple[float, float] | list[float]] | npt.NDArray[float],
 ) -> Graph:
-    """Modifies a given connection graph. Removed the nodes of spaces on which a players stand.
+    """Modifies a given movement graph. Removed the nodes of spaces on which players stand.
 
     Args:
         graph: The graph to modify.
@@ -279,48 +279,4 @@ 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[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()))