training input and target about faster-trajectory-transformer

[Howuhh]

Hi! As this is the code in getitem of the dataset, it only operates on [seq_len, ...] shapes, i.e. unbatched sequences

[wj210]

Yes, but however the slicing is done after flattening the seq_len together with the token dims so effectively we only slide it by 1 token? Instead of 1 entire timestep worth of tokens

[Howuhh]

@wj210 But that's the point! We discretize each domain such as states, actions, reward, then concatenate it to the one sequence and then flatten it (see architecture figure in the original paper). Thus, our task is to predict next token (not entire next timestamp!).
For example, we will predict from token 0 of trajectory 0 ($s_0^0$) next token 1 of trajectory 0 ($s_0^1$), which is just next state dimension. Actual sequence length for the transformer is seq_len * (state_dim + action_dim + 1 for return to go)

[wj210]

Right, yes i guess i was just confused about that... Thanks again!