Fix visualisation bug from unnormalising ego-vehicle out of feature b…

…ounds see #79
eleurent · Mar 4, 2023 · 02e0303 · 02e0303
1 parent 35216ba
commit 02e0303
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rl_agents/agents/deep_q_network/graphics.py b/rl_agents/agents/deep_q_network/graphics.py
@@ -106,6 +106,7 @@ def compute_vehicles_attention(cls, agent, state):
             if not obs_type.absolute and v_index > 0:
                 v_position += agent.env.unwrapped.vehicle.position
             vehicle = min(agent.env.road.vehicles, key=lambda v: np.linalg.norm(v.position - v_position))
+            vehicle = agent.env.unwrapped.vehicle if v_index == 0 else vehicle
             v_attention[vehicle] = attention[:, v_index]
         return v_attention