Revisiting the Design Choices in Max-Return Sequence Modeling

What dataset characteristics specifically correlate with optimal context lengths? The paper mentions variations in context length without thoroughly linking these findings to specific dataset properties. It would strengthen the findings to identify specific dataset characteristics that correlate with optimal context lengths. Again, properties such as trajectory length, state/action space dimensionality, and reward sparsity could be analyzed to determine how they influence context length choices. Exploring these aspects could reveal generalizable patterns for determining optimal context lengths in other datasets and applications, potentially providing valuable tuning guidelines.

Do you anticipate that the observed effects of architecture and context length on performance will hold for other reinforcement learning tasks or domains not included in the study? Any insight into the potential for generalizing these findings would strengthen the conclusions. Or do you have any marginal evidence that might suggest this?

Will you provide code, model checkpoints, or more detailed instructions for setting up experiments? This would help others replicate and validate your findings and would enhance the practical impact of your contributions.

1. Given that one of the potential advantages of this method over traditional RL is reduced complexity, are the computational resources used reduced as well? More details on this would help understand the method's practical benefits.
2. In Section 5.1, the quality of the datasets, which claims to be the reason affecting the performance against IQL, can be further elaborated.
3. Table 4 suggests a linear relationship between the performance and context length $K$. Could you discuss the reason for choosing to model these relationships linearly? And, if I was not mistaken, in this regression analysis, $K$ values are limited to the discrete support of $\{2,5,10,20\}$, which may not accurately reflect the relationship by the interpolation. Presenting these plots visually assessing this relationship could provide clearer evidence of linearity or indicate the need for a different modeling approach, such as polynomial regression or categorically treating each $K$ value. Alternatively, more detailed justification for the chosen method would help clarify the rationale behind this modeling choice.

• In the future work, you mention it would be interesting to explore how to better integrate classical RL with sequence modelling to harness the strength of both, do you have any concrete ideas here? I feel like including something like it in this current paper could make it much stronger
• You mention that positional embeddings might hinder trajectory stitching. Could you elaborate on this insight? What types of positional embeddings did you use? Do you think different tasks / datasets might benefit from different types of positional embeddings? (e.g. absolute vs cosine etc.)

• In 5.2.1, the authors use a "0" token, but presumably the model has never (or rarely?) seen this. Why not just compute the Jacobian or some similar quantity to see the dependence of the prediction on past inputs? Or am I wrong an the models actually do see mask tokens?
• "Thus, the Reinformer pays more attention to global information", this is an interesting observation, but it doesn't translate to supporting a claim the architecture impacts performance. If different architectures "weigh" temporal information differently but have the exact same policy logits, are they really different? There needs to be a more causal/mechanistic claim here.
• There's a fair amount of speculation (especially about environment properties) that feels hidden in the text, for example, the authors claim that "the Reinformer’s focus on global historical trajectory information is particularly adept at considering and utilizing waypoint-related information effectively." This is speculation from (a) Figure 2 (b) and the authors' understanding of maze-large. A better, and less speculative, experiment would directly intervene on the agent to test a hypothesis; for example in this case, the agents could be made to "cheat" by being given waypoint information, and if this intervention removes the Reinformer's superiority, then we've learned that the other models are probably worse at modeling this specifically.
• Is the x axis inverted on Figures 3 and 4 and 6b? Table 2 (and the text) suggests e.g. that Reinformer is superior to Reimba on Maze2d-large.
• "Correspondingly, a smaller context length aligns more closely with the Markov Decision Process (MDP) framework, and thus performs better." I'm not sure this conclusion naturally follows. If I understand correctly, the argument is that offline RL on a worse & diverse replay buffer means more effective capacity is spent modeling unimportant states, which means that we should show less of them at a time to be closer to a 1-Markov setting. To me this is not obviously true because even with a lower K, the model still sees all the states as part of its training data, and is still spending effective capacity on them regardless of the sequence in which they appear. To me a more plausible reason is that credit assignment is harder if the model has to consider many possible explanations (i.e. K times more) for the output; or in other words with longer inputs the variance is larger. But this could be demonstrated empirically rather than through speculation.