Clarification on Visual State Interaction

[Gigi-G]

Dear @idejie

Thank you for sharing the code for your impressive work — it's been a pleasure to explore it.

I had a question regarding the Visual State Interaction. In the paper, it appears that the interaction is described as occurring specifically between the visual embeddings of the start and end frames:

However, in the code https://github.com/idejie/PlanLLM/blob/58b2a2e03ec84c375b9b2fe2d531d5878954312e/models/schema/state_encoder.py#L79C9-L79C56, the variable state_feat seems to include not only $x_{0}^{v}$ (the embedding of the start) and $x_{T}^{v}$ (the embedding of the end), but also features from intermediate frames (as indicated in the comment on this line https://github.com/idejie/PlanLLM/blob/58b2a2e03ec84c375b9b2fe2d531d5878954312e/models/schema/state_encoder.py#L74C11-L74C40).

Am I missing something? Could you please clarify whether this broader inclusion is intentional in the context of the visual interaction module?

Thanks in advance for your time and insights!

Best regards,
Luigi

[idejie]

Hi, thanks for your attention on our work.

We add a self-attention on tje context of the visual interaction from start to end.
And we slice the start and end state for late operationsmodels/schema/state_encoder.py#L82-L85.

Similar implementation of previous method SCHMEA.

[Gigi-G]

Hi @idejie, thanks for getting back to me.

From what I can tell in SCHEMA’s code, the overall design is similar. However, their state_encoder (which appears in your repository as well here) processes the start and end frames independently of the intermediate frames because the linear layers act node-wise. In your version, the self-attention links every frame to every other frame, which unintentionally lets information from the intermediate frames leak into start/end predictions.

To prevent that, I modified the block here like this:

# Encode state features: [start0, start1, …, endT]
state_feat = self.state_encoder(self.dropout(state_feat))

# Separate start, end, and intermediate features
start_state_feat = state_feat[:, 0:1, :]      # [batch, 1, embed_dim]
end_state_feat   = state_feat[:, -1:, :]      # [batch, 1, embed_dim]
inter_state_feat = state_feat[:, 1:-1, :]     # [batch, T-1, embed_dim]

# Let self-attention see only start and end
start_end_feat = self.state_interaction(
    torch.cat([start_state_feat, end_state_feat], dim=1)  # [batch, 2, embed_dim]
)

start_state_feat = start_end_feat[:, 0, :]  # [batch, embed_dim]
end_state_feat   = start_end_feat[:, 1, :]  # [batch, embed_dim]

With this change, self-attention focuses solely on the start and end signals, preventing leakage from intermediate frames. After applying it, the results drop to about the same level as the original SCHEMA rather than the higher numbers reported in the paper.

Am I overlooking something?

Best regards,
Luigi