diff --git a/spinn_gym/games/breakout/breakout_machine_vertex.py b/spinn_gym/games/breakout/breakout_machine_vertex.py
index d6c5163..3fe9366 100644
--- a/spinn_gym/games/breakout/breakout_machine_vertex.py
+++ b/spinn_gym/games/breakout/breakout_machine_vertex.py
@@ -38,6 +38,7 @@
 
 from spynnaker.pyNN.data import SpynnakerDataView
 from spynnaker.pyNN.models.common import PopulationApplicationVertex
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -149,8 +150,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             BreakoutMachineVertex._BREAKOUT_REGIONS.BREAKOUT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
         if self.app_vertex.source_vertex is None:
             raise ValueError(
                 "The breakout vertex doesn't have a source vertex!")
diff --git a/spinn_gym/games/double_inverted_pendulum/double_pendulum_machine_vertex.py b/spinn_gym/games/double_inverted_pendulum/double_pendulum_machine_vertex.py
index c3bdce1..2bb740b 100644
--- a/spinn_gym/games/double_inverted_pendulum/double_pendulum_machine_vertex.py
+++ b/spinn_gym/games/double_inverted_pendulum/double_pendulum_machine_vertex.py
@@ -33,6 +33,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -166,8 +167,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._DOUBLE_PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting double pendulum recording region:\n")
diff --git a/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py b/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
index b2d4a48..c78f92b 100644
--- a/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
+++ b/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
@@ -36,6 +36,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -164,8 +165,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting pendulum recording region:\n")
diff --git a/spinn_gym/games/logic/logic_machine_vertex.py b/spinn_gym/games/logic/logic_machine_vertex.py
index 3f94612..467a10e 100644
--- a/spinn_gym/games/logic/logic_machine_vertex.py
+++ b/spinn_gym/games/logic/logic_machine_vertex.py
@@ -37,6 +37,8 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
+
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
 
@@ -148,8 +150,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._LOGIC_REGIONS.LOGIC.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting logic recording region:\n")
diff --git a/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py b/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
index 9cabf08..9b985e4 100644
--- a/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
+++ b/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
@@ -37,6 +37,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -159,8 +160,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._BANDIT_REGIONS.BANDIT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting bandit recording region:\n")
diff --git a/spinn_gym/games/store_recall/store_recall_machine_vertex.py b/spinn_gym/games/store_recall/store_recall_machine_vertex.py
index 9f49932..d237bed 100644
--- a/spinn_gym/games/store_recall/store_recall_machine_vertex.py
+++ b/spinn_gym/games/store_recall/store_recall_machine_vertex.py
@@ -36,6 +36,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -153,8 +154,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._RECALL_REGIONS.RECALL.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
-            vertex))
+        spec.write_value(routing_info.get_safe_first_key_from_pre_vertex(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting recall recording region:\n")