A question about the visualization of the attention matrix

[Twolige]

Hi @eleurent, thanks for your work!

Nowadays I'm training an agent with the _sb3_highway_ppotransformer.py and just wonder,

1. How to visualize the attention matrix with different color and width like the way that you've showed in your paper Social Attention for Autonomous Decision-Making in Dense Traffic?
2. Can I use the _sb3_highway_ppotransformer.py to get similar visualization image?

Thanks for replying in advance.

[eleurent]

Hi @Twolige,

The function that I use to visualise the attention matrix is here:

https://github.com/eleurent/rl-agents/blob/9ed10b2e8407afb606857d332b51babb52d2d90c/rl_agents/agents/deep_q_network/graphics.py#L55

If you want to use sb3_highway_ppo_transformer.py, you will have to copy and call this function from the script, but it will also require some adaptation. In particular, the code for getting the attention matrix from the SB3 model must be changed (it was done there with my own agent implementation:

https://github.com/eleurent/rl-agents/blob/9ed10b2e8407afb606857d332b51babb52d2d90c/rl_agents/agents/deep_q_network/graphics.py#L87

once you have the attention matrix, you have to map its rows to the actual Vehicle objects in the scene as done here:

https://github.com/eleurent/rl-agents/blob/9ed10b2e8407afb606857d332b51babb52d2d90c/rl_agents/agents/deep_q_network/graphics.py#L108

and then use display_vehicles_attention to draw the lines on the simulation surface (env.unwrapped.viewer.sim_surface)

[eleurent]

Okay, maybe it is a bit tricky, I'll try to do it when I have a few spare minutes