diff --git a/src/robofish/io/file.py b/src/robofish/io/file.py
index 1ce2f1ee4e3119ae178919e0f7620b80f41ec2db..1ea891f30c7e33a029633024a5c64cfb888481e6 100644
--- a/src/robofish/io/file.py
+++ b/src/robofish/io/file.py
@@ -928,7 +928,7 @@ class File(h5py.File):
         else:
             step_size = poses.shape[1]
 
-        cmap = matplotlib.cm.get_cmap(cmap)
+        cmap = plt.get_cmap(cmap)
 
         x_world, y_world = self.world_size
         if figsize is None:
@@ -951,6 +951,10 @@ class File(h5py.File):
                 this_c = c
 
             timesteps = poses.shape[1] - 1
+            colours = plt.get_cmap(cmap)(np.linspace(0, 1, 1 + timesteps // step_size))
+
+            if not lw_distances:
+                colours[:] = this_c
             for t in range(0, timesteps, step_size):
                 if lw_distances:
                     lw = np.mean(line_width[t : t + step_size + 1])
@@ -958,7 +962,7 @@ class File(h5py.File):
                 ax.plot(
                     poses[fish_id, t : t + step_size + 1, 0],
                     poses[fish_id, t : t + step_size + 1, 1],
-                    c=this_c,
+                    c=colours[t//step_size],
                     lw=lw,
                 )
             # Plotting outside of the figure to have the label