From 6008e4cb28980cf9bec691d1c63da6938bac7a14 Mon Sep 17 00:00:00 2001
From: Erik Nygren <erik.nygren@sbb.ch>
Date: Thu, 16 May 2019 17:55:35 +0200
Subject: [PATCH] fixed bug in rendering of tree state observation

---
 flatland/envs/observations.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/flatland/envs/observations.py b/flatland/envs/observations.py
index bde43bf8..dba5faf9 100644
--- a/flatland/envs/observations.py
+++ b/flatland/envs/observations.py
@@ -385,7 +385,7 @@ class TreeObsForRailEnv(ObservationBuilder):
                                                                           depth + 1)
                 observation = observation + branch_observation
                 if len(branch_visited) != 0:
-                    visited.union(branch_visited)
+                    visited = visited.union(branch_visited)
             elif last_isSwitch and possible_transitions[branch_direction]:
                 new_cell = self._new_position(position, branch_direction)
                 branch_observation, branch_visited = self._explore_branch(handle,
@@ -395,10 +395,7 @@ class TreeObsForRailEnv(ObservationBuilder):
                                                                           depth + 1)
                 observation = observation + branch_observation
                 if len(branch_visited) != 0:
-                    visited.union(branch_visited)
-                    print(visited)
-                    print(branch_visited)
-                    print("-------")
+                    visited = visited.union(branch_visited)
             else:
                 num_cells_to_fill_in = 0
                 pow4 = 1
-- 
GitLab