Understanding attention maps

[LLusvarghi]

Hi @SimonGer

Wonderful work!
I am trying to understand the attention maps provided by the model (here) to get some insights on the learned driving policy. In this regard, I have a couple of questions:

1. Attention maps provide one attention score for each embedding token, right?
2. How does the mapping between the input batch data and the embedding tokens work? I would like to know at which index of the embedding array I can find the tokenized version of the different input data elements (e.g., ego speed, surrounding objects, etc.), so that I can later extract the corresponding attention score.

Ultimately, my goal is to extract the attention score assigned by the model to each piece of input data.
Any help from your side on this matter will be truly appreciated!

[SimonGer]

Hi Luca,

thanks! The attention score visualization is a very cool part of PlanT 1.0 that I unfortunately didn't get to during my work, so I don't think I can be a big help there :(

I think you are correct for point 1, but I'm not 100% sure.

You can look at the original PlanT repo for some inspiration about the visualization stuff. This commit could be useful.
I think you can also find out a lot by looking at how the batches are built in the agent and what is being processed in the model.
Basically you can see that there are certain fixed tokens that get added to the front of the list of embeddings and then we have the objects, which have the same order as they were given in by the agent. Maybe its already enough to find out what each of the tokens before is in your case and then zip the remaining ones with the boxes, but I'm just guessing.

Best of Luck! If you get anything working I would love to hear about it that would be really cool!!