diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index b4b1818ff3d992fe215aef037de023465b1d444a..8e746c4869334e461a809e10cbdc012435945757 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -842,19 +842,19 @@ class File(h5py.File):
 
     def plot(
         self,
-        ax=None,
-        lw_distances=False,
-        lw=2,
-        ms=32,
-        figsize=None,
-        step_size=4,
-        c=None,
-        cmap="Set1",
+        ax: mpl.axes = None,
+        lw_distances: bool = False,
+        lw: int = 2,
+        ms: int = 32,
+        figsize: Tuple[int] = None,
+        step_size: int = 4,
+        c: List = None,
+        cmap: mpl.colors.Colormap = "Set1",
         skip_timesteps=0,
         max_timesteps=None,
         show=False,
         legend=True,
-    ):
+    ) -> mpl.axes:
         """Plot the file using matplotlib.pyplot
 
         The tracks in the file are plotted using matplotlib.plot().
@@ -862,9 +862,11 @@ class File(h5py.File):
         Args:
             ax (matplotlib.axes, optional): An axes object to plot in. If None is given, a new figure is created.
             lw_distances (bool, optional):  Flag to show the distances between individuals through line width.
+            lw(int, optional): Line width of the tracks. Defaults to 2.
+            ms(int, optional): Marker size of the tracks. Defaults to 32.
             figsize (Tuple[int], optional): Size of a newly created figure.
             step_size (int, optional): when using lw_distances, the track is split into sections which have a common line width. This parameter defines the length of the sections.
-            c (Array[color_representation], optional): An array of colors. Each item has to be matplotlib.colors.is_color_like(item).
+            c (List, optional): An array of colors. Each item has to be matplotlib.colors.is_color_like(item).
             cmap (matplotlib.colors.Colormap, optional): The colormap to use
             skip_timesteps (int, optional): Skip timesteps in the begining of the file
             max_timesteps (int, optional): Cut of timesteps in the end of the file.
@@ -1076,10 +1078,16 @@ class File(h5py.File):
             plt.plot([], [], linestyle="dotted", alpha=0.5, color="k", zorder=0)[0],
         ]
         categories = [entity.attrs.get("category", None) for entity in self.entities]
+
         entity_polygons = [
-            patches.Polygon(shape_vertices(options["entity_scale"]), facecolor=color)
-            for color in [
-                "gray" if category == "robot" else "k" for category in categories
+            patches.Polygon(
+                shape_vertices(options["entity_scale"]),
+                edgecolor=edgecolor,
+                facecolor=color,
+            )
+            for edgecolor, color in [
+                ("k", "white") if category == "robot" else ("k", "k")
+                for category in categories
             ]
         ]