diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index fecca331b7e8253d5196f9bfd881a40b4ab672d6..9dc397a0162fb46687f84b30516fc5d613703314 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -396,6 +396,7 @@ class TreeObsForRailEnv(ObservationBuilder):
             cell_transitions = self.env.rail.get_transitions(*position, direction)
             total_transitions = bin(self.env.rail.get_full_transitions(*position)).count("1")
             num_transitions = np.count_nonzero(cell_transitions)
+
             exploring = False
 
             # Detect Switches that can only be used by other agents.
@@ -404,14 +405,10 @@ class TreeObsForRailEnv(ObservationBuilder):
 
             if num_transitions == 1:
                 # Check if dead-end, or if we can go forward along direction
-                nbits = 0
-                tmp = self.env.rail.get_full_transitions(*position)
-                while tmp > 0:
-                    nbits += (tmp & 1)
-                    tmp = tmp >> 1
+                nbits = total_transitions
                 if nbits == 1:
                     # Dead-end!
-                    last_is_dead_end = True
+                    last_is_dead_end = False
 
                 if not last_is_dead_end:
                     # Keep walking through the tree along `direction'
@@ -435,8 +432,6 @@ class TreeObsForRailEnv(ObservationBuilder):
 
         # `position' is either a terminal node or a switch
 
-        observation = []
-
         # #############################
         # #############################
         # Modify here to append new / different features for each visited cell!
@@ -464,6 +459,7 @@ class TreeObsForRailEnv(ObservationBuilder):
                            other_agent_same_direction,
                            other_agent_opposite_direction
                            ]
+
         else:
             observation = [own_target_encountered,
                            other_target_encountered,