diff --git a/spinn_gym/games/breakout/breakout_machine_vertex.py b/spinn_gym/games/breakout/breakout_machine_vertex.py
index e1ba1a2..d0b2061 100644
--- a/spinn_gym/games/breakout/breakout_machine_vertex.py
+++ b/spinn_gym/games/breakout/breakout_machine_vertex.py
@@ -37,8 +37,8 @@
     front_end_common_constants
 
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
 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
@@ -150,10 +150,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             BreakoutMachineVertex._BREAKOUT_REGIONS.BREAKOUT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        data = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert data is not None
-        spec.write_value(data)
+        spec.write_value(routing_info.get_key_from(
+            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 8620fff..ba3936a 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,7 +33,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -167,10 +167,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._DOUBLE_PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        data = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert data is not None
-        spec.write_value(data)
+        spec.write_value(routing_info.get_key_from(
+            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 ee52931..378ea90 100644
--- a/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
+++ b/spinn_gym/games/inverted_pendulum/inverted_pendulum_machine_vertex.py
@@ -36,7 +36,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -165,10 +165,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._PENDULUM_REGIONS.PENDULUM.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        data = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert data is not None
-        spec.write_value(data)
+        spec.write_value(routing_info.get_key_from(
+            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 7c2f667..80d73f6 100644
--- a/spinn_gym/games/logic/logic_machine_vertex.py
+++ b/spinn_gym/games/logic/logic_machine_vertex.py
@@ -37,7 +37,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -150,10 +150,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._LOGIC_REGIONS.LOGIC.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        first_key = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert first_key is not None
-        spec.write_value(first_key)
+        spec.write_value(routing_info.get_key_from(
+            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 34e4c22..676f090 100644
--- a/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
+++ b/spinn_gym/games/multi_arm_bandit/bandit_machine_vertex.py
@@ -37,7 +37,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -160,10 +160,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._BANDIT_REGIONS.BANDIT.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        first_key = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert first_key is not None
-        spec.write_value(first_key)
+        spec.write_value(routing_info.get_key_from(
+            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 c02c67c..1b4163c 100644
--- a/spinn_gym/games/store_recall/store_recall_machine_vertex.py
+++ b/spinn_gym/games/store_recall/store_recall_machine_vertex.py
@@ -36,7 +36,7 @@
 
 # sPyNNaker imports
 from spynnaker.pyNN.data import SpynnakerDataView
-from spynnaker.pyNN.utilities import constants
+from spynnaker.pyNN.utilities.constants import SPIKE_PARTITION_ID
 
 # spinn_gym imports
 from spinn_gym.games import SpinnGymMachineVertex
@@ -154,10 +154,8 @@ def generate_data_specification(
         spec.switch_write_focus(
             self._RECALL_REGIONS.RECALL.value)
         routing_info = SpynnakerDataView.get_routing_infos()
-        data = routing_info.get_first_key_from_pre_vertex(
-            vertex, constants.SPIKE_PARTITION_ID)
-        assert data is not None
-        spec.write_value(data)
+        spec.write_value(routing_info.get_key_from(
+            vertex, SPIKE_PARTITION_ID))
 
         # Write recording region for score
         spec.comment("\nWriting recall recording region:\n")