LAQ: Latent action reshape error, pretrained laq cannot reconstruct images from ssv2

[Woodenonez]

Hi,

It seems like this repo is not active anymore, so I'm marking this problem for the future, and I might fix it if I have time.

1. In the LAQ forward method:

if math.sqrt(self.code_seq_len) % 1 == 0: # "code_seq_len should be square number"
    action_h = int(math.sqrt(self.code_seq_len))
    action_w = int(math.sqrt(self.code_seq_len))
elif self.code_seq_len == 2:
    action_h = 2
    action_w = 1
else:
    raise ValueError("code_seq_len should be square number or defined as 2")

This code is used to reshape the latent action token for decoding. This works for the hyperparameters given in the repo, but it seems like a coincidence rather than a general fix. If the hyperparameters are changed, there will be an error in the decode method

tokens = rearrange(tokens, 'b t h w d -> (b t) (h w) d')
actions = rearrange(actions, 'b t h w d -> (b t) (h w) d') # BUG?: Check, this is the original code, but it seems wrong.
# actions = rearrange(actions, "b t h w d -> b (t h w) d")
attn_bias = self.spatial_rel_pos_bias(h, w, device = tokens.device)
tokens = self.dec_spatial_transformer(tokens, attn_bias = attn_bias, video_shape = video_shape, context=actions)
tokens = rearrange(tokens, '(b t) (h w) d -> b t h w d', b = b, h = h , w = w)

A more general method is to determine the action h and w from the CNN encoder in the VQ.

2. I also tried the given pretrained LAQ from HF, and it cannot reconstruct figures in SSv2, and it outputs basically noise. I personally believe the proposed method is solid, but I doubt the given code is complete or the newest.

If anyone else encounters the same problem and fixes it, it is much appreciated. I will also try to fix it and do some pretraining.

Cheers~

[Woodenonez]

An easy fix is to define action h and w as:

if math.sqrt(tokens.shape[1]) % 1 == 0:
    action_h = action_w = int(math.sqrt(tokens.shape[1]))
else:
    raise ValueError(f"Cannot reshape tokens of shape {tokens.shape} into a square for action tokens.")