diff --git a/arena_bringup/launch/start_training.launch b/arena_bringup/launch/start_training.launch
index 6804fc9f..8a69b9e5 100644
--- a/arena_bringup/launch/start_training.launch
+++ b/arena_bringup/launch/start_training.launch
@@ -24,6 +24,9 @@
   <param name="robot_yaml_path" value="$(find arena-simulation-setup)/robot/$(arg model)/$(arg model).model.yaml" />
   <param name="tmp_model_path" value="/tmp" />
 
+  <!-- Path to a json file to load custom neural network -->
+  <param name="custom_network_path" value="../network.json"/>
+
   <rosparam command="load" file="$(find arena_bringup)/params/flatland.yaml" />
   <rosparam command="load" file="$(find arena_bringup)/params/drl.yaml" />
 
diff --git a/training/scripts/train_agent.py b/training/scripts/train_agent.py
old mode 100755
new mode 100644
diff --git a/training/tools/argsparser.py b/training/tools/argsparser.py
index 46645e7e..127dc53f 100644
--- a/training/tools/argsparser.py
+++ b/training/tools/argsparser.py
@@ -23,6 +23,7 @@ def training_args(parser):
     import rosnav.model.custom_policy
     import rosnav.model.custom_sb3_policy
     from rosnav.model.agent_factory import AgentFactory
+    import rosnav.model.custom_policy_from_json
 
     group.add_argument(
         "--agent",