diff --git a/spinn_gym/games/breakout/breakout_machine_vertex.py b/spinn_gym/games/breakout/breakout_machine_vertex.py
index 9521208..1de6bba 100644
--- a/spinn_gym/games/breakout/breakout_machine_vertex.py
+++ b/spinn_gym/games/breakout/breakout_machine_vertex.py
@@ -135,8 +135,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             BreakoutMachineVertex._BREAKOUT_REGIONS.BREAKOUT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
         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 12bc043..242b7ee 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
@@ -162,8 +162,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             self._DOUBLE_PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
 
         # 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 3501fe9..954a684 100644
--- a/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
+++ b/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
@@ -160,8 +160,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             self._PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
 
         # 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 34d1775..b251df3 100644
--- a/spinn_gym/games/logic/logic_machine_vertex.py
+++ b/spinn_gym/games/logic/logic_machine_vertex.py
@@ -145,8 +145,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             self._LOGIC_REGIONS.LOGIC.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
 
         # 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 f116819..aba4908 100644
--- a/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
+++ b/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
@@ -155,8 +155,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             self._BANDIT_REGIONS.BANDIT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
 
         # 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 a1bcea1..9d6dc4c 100644
--- a/spinn_gym/games/store_recall/store_recall_machine_vertex.py
+++ b/spinn_gym/games/store_recall/store_recall_machine_vertex.py
@@ -149,8 +149,8 @@ def generate_data_specification(self, spec, placement):
         spec.switch_write_focus(
             self._RECALL_REGIONS.RECALL.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        spec.write_value(routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID))
+        spec.write_value(routing_info.get_single_first_key_from_pre_vertex(
+            vertex))
 
         # Write recording region for score
         spec.comment("\nWriting recall recording region:\n")