diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
new file mode 100644
index 0000000000000000000000000000000000000000..e39e367142546a872aef9841461e601b1584c415
--- /dev/null
+++ b/.gitlab-ci.yml
@@ -0,0 +1,33 @@
+image: themattrix/tox
+
+##########################################
+##########################################
+## We have to set the following env vars
+## in the admin interface :
+## - AWS_DEFAULT_REGION
+## - BUCKET_NAME
+## - AWS_ACCESS_KEY_ID
+## - AWS_SECRET_ACCESS_KEY
+
+stages:
+    - tests
+    - benchmarks_and_profiling
+    - deploy_docs
+cache:
+    paths:
+        - .tox
+
+before_script:
+    - echo "Setting Up...."
+
+tests:
+    stage: tests
+    script:
+        - apt update
+        - apt install -y libgl1-mesa-glx xvfb graphviz xdg-utils libcairo2-dev libjpeg-dev libgif-dev
+        - pip install tox
+        - xvfb-run tox -v --recreate
+
+
+
+
diff --git a/requirements_torch_training.txt b/requirements_torch_training.txt
index 2bce630587233b4c771ef7a43bc3aaf7f78fbb07..8033798ff5c83ae4eee1f6969a8a9f4c6187c1b6 100644
--- a/requirements_torch_training.txt
+++ b/requirements_torch_training.txt
@@ -1 +1,2 @@
-torch==1.1.0
\ No newline at end of file
+git+http://gitlab.aicrowd.com/flatland/flatland.git@master
+torch>=1.1.0
\ No newline at end of file
diff --git a/tox.ini b/tox.ini
index 3c22b56780ffa59d41f64cad3f9698c3f62a204d..f55e297488a729a9bb95b370ff7c4980801e1909 100644
--- a/tox.ini
+++ b/tox.ini
@@ -21,7 +21,7 @@ passenv =
 deps =
     -r{toxinidir}/requirements_torch_training.txt
 commands =
-    python torch_training/training_navigation.py
+    python torch_training/multi_agent_training.py
 
 [flake8]
 max-line-length = 120
diff --git a/utils/observation_utils.py b/utils/observation_utils.py
index 0c97b186a9331f185cf1a1d3f99685581cb551f7..5e01121fe84c24d8d5a46d92bb30578e1dcda2b0 100644
--- a/utils/observation_utils.py
+++ b/utils/observation_utils.py
@@ -1,5 +1,7 @@
 import numpy as np
 
+from flatland.envs.observations import TreeObsForRailEnv
+
 
 def max_lt(seq, val):
     """
@@ -48,7 +50,7 @@ def norm_obs_clip(obs, clip_min=-1, clip_max=1):
     return np.clip((np.array(obs) - min_obs) / norm, clip_min, clip_max)
 
 
-def split_tree(tree, num_features_per_node=9, current_depth=0):
+def split_tree(tree, current_depth=0):
     """
     Splits the tree observation into different sub groups that need the same normalization.
     This is necessary because the tree observation includes two different distance:
@@ -64,6 +66,7 @@ def split_tree(tree, num_features_per_node=9, current_depth=0):
     :param current_depth: Keeping track of the current depth in the tree
     :return: Returns the three different groups of distance and binary values.
     """
+    num_features_per_node = TreeObsForRailEnv.observation_dim
 
     if len(tree) < num_features_per_node:
         return [], [], []
@@ -88,7 +91,6 @@ def split_tree(tree, num_features_per_node=9, current_depth=0):
         child_tree = tree[(num_features_per_node + children * child_size):
                           (num_features_per_node + (children + 1) * child_size)]
         tmp_tree_data, tmp_distance_data, tmp_agent_data = split_tree(child_tree,
-                                                                      num_features_per_node,
                                                                       current_depth=current_depth + 1)
         if len(tmp_tree_data) > 0:
             tree_data.extend(tmp_tree_data)