diff --git a/run.py b/run.py
index 143f44ad102adf0b5f4791136fe44a81870abdc5..048dcdee5b3a5ebe3cb1fb01cbd5f07aea3f5480 100644
--- a/run.py
+++ b/run.py
@@ -27,7 +27,7 @@ VERBOSE = True
 
 # Checkpoint to use (remember to push it!)
 # checkpoint = "./checkpoints/201112143850-5400.pth" # 21.220418678677177 DEPTH=2 AGENTS=10
-checkpoint = "./checkpoints/201113211844-6200.pth" # 19.690047767961005 DEPTH=2 AGENTS=20
+checkpoint = "./checkpoints/201113211844-6700.pth" # 19.690047767961005 DEPTH=2 AGENTS=20
 
 
 # Use last action cache
diff --git a/utils/fast_tree_obs.py b/utils/fast_tree_obs.py
index 0666ef4224eb43a09a96510648dc21c43cadaf75..c388d2a9cdeda91c8835975cc9c366d6bb2ff01a 100755
--- a/utils/fast_tree_obs.py
+++ b/utils/fast_tree_obs.py
@@ -1,3 +1,5 @@
+from typing import List, Optional
+
 import numpy as np
 from flatland.core.env_observation_builder import ObservationBuilder
 from flatland.core.grid.grid4_utils import get_new_position
@@ -232,6 +234,12 @@ class FastTreeObs(ObservationBuilder):
 
         return has_opp_agent, has_same_agent, has_target, visited
 
+    def get_many(self, handles: Optional[List[int]] = None):
+        self.dead_lock_avoidance_agent.start_step()
+        observations = super().get_many(handles)
+        self.dead_lock_avoidance_agent.end_step()
+        return observations
+
     def get(self, handle):
         # all values are [0,1]
         # observation[0]  : 1 path towards target (direction 0) / otherwise 0 -> path is longer or there is no path
@@ -262,9 +270,6 @@ class FastTreeObs(ObservationBuilder):
         # observation[25] : If there is a switch on the path which agent can not use -> 1
         # observation[26] : If there the dead-lock avoidance agent predicts a deadlock -> 1
 
-        if handle == 0:
-            self.dead_lock_avoidance_agent.start_step()
-
         observation = np.zeros(self.observation_dim)
         visited = []
         agent = self.env.agents[handle]