diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index f345648533c7ab554403b7031245139e39e51b75..e7f736f5c3ad2b9c2a6ff8ad99f5f5706224bfde 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -414,7 +414,7 @@ class RailEnv(Environment):
         # Agents not done (arrival_time is None)
         else:
             # CANCELLED check (never departed)
-            if (agent.state == TrainState.READY_TO_DEPART):
+            if (agent.state.is_off_map_state()):
                 reward = -1 * self.cancellation_factor * \
                     (agent.get_travel_time_on_shortest_path(self.distance_map) + self.cancellation_time_buffer)