Paper: Decision ConvFormer: Local Filtering in MetaFormer is Sufficient for Decision Making (ICLR-2024)

Reviewer comment (Reproducibility):
How to compute the embeddings of a float number (reward) in the subsection 3.1? Some explanations might be helpful.

Author response:
To compute the RTG embeddings, we adopted the approach used in Ref. [C.1], which involves using a single linear layer to transform a one-dimensional float RTG input into a vector of a predefined hidden size.

Reviewer comment (Theory):
The reasons why the propose method is effective are needed to explained. It seems that the self-attention operation is more expressive then the proposed block (three causal convolution filters). Is the proposed DC only suitable for some settings, e.g. the setting with less data?

Author response:
The reason why DC is effective can be understood from an inductive bias perspective.  Consider the use of MLPs and convolution layers in the computer vision domain. While MLPs are more general and expressive than convolution layers, the local-interaction inductive bias of a convolution layer makes it particularly efficient for image processing. Similarly, although transformers are more expressive than DC, when solving offline RL problems with Markov properties, injecting this Markovianness inductive bias can significantly enhance learning efficiency and performance. Our DC model successfully incorporates the Markovianness inductive bias by employing convolution filters that feature three local windows, thus making it more effective compared with DT. We think that DC is effective for most offline RL tasks with strong Markov property.

Reviewer comment (Evaluation):
Why the ODC is worse than DC on some tasks?

Author response:
As mentioned in [C.2], the ODT/ODC architecture, which incorporates elements of stochasticity for online adaptation, is more complex than the DT/DC architecture. This added complexity might pose challenges in the Antmaze domain which is characterized by its sparse reward nature, potentially complicating the optimization process and leading to lower performance.

Reviewer comment (Writing):
The author could remove the mention of the MetaFormer framework. As someone who has never heard it before, I first thought MetaFormer is a new transformer variant, but then I realized it's just a framework, which is a bit confusing to me.

Author response:
We thank the reviewer for pointing out the potential confusion regarding the use of the terms "MetaFormer" and "token mixer block" in our paper. We understand that these terms, particularly for readers unfamiliar with them, might lead to some misunderstanding.

Although we acknowledge the reviewer's suggestion to remove mentioning the MetaFormer framework, the development of DC in our paper is essentially based on the MetaFormer framework. We choose to keep it. To prevent any confusion, however,  we will provide a clear explanation in the early part of the introduction. Specifically, we will emphasize two key points:

1) The term "MetaFormer" is used to refer to a general framework rather than a new variant of the transformer model. We will elaborate on its role and significance in the context of our research to offer a clearer understanding for readers.

2) We will clarify that the "token mixer block" in our paper does not refer to the token mixing concept found in the MLP-Mixer. We will provide a precise definition and contextualize its usage within our proposed framework.

Reviewer comment (Writing):
The name "token mixer block" should be avoided, since it reminds me of the token mixing in the MLP-Mixer, which made me confused in the beginning.

Author response:
We thank the reviewer for pointing out the potential confusion regarding the use of the terms "MetaFormer" and "token mixer block" in our paper. We understand that these terms, particularly for readers unfamiliar with them, might lead to some misunderstanding.

Although we acknowledge the reviewer's suggestion to remove mentioning the MetaFormer framework, the development of DC in our paper is essentially based on the MetaFormer framework. We choose to keep it. To prevent any confusion, however,  we will provide a clear explanation in the early part of the introduction. Specifically, we will emphasize two key points:

1) The term "MetaFormer" is used to refer to a general framework rather than a new variant of the transformer model. We will elaborate on its role and significance in the context of our research to offer a clearer understanding for readers.

2) We will clarify that the "token mixer block" in our paper does not refer to the token mixing concept found in the MLP-Mixer. We will provide a precise definition and contextualize its usage within our proposed framework.

Reviewer comment (Experiments):
Can you further describe details of the motivating examples? Do you only learn the attention directly of that one layer or all layers?

Author response:
In our motivating examples, rather than using an attention module, we trained $3K \times 3K$ parameters directly in all layers. Here, $K$ denotes the context length, and the multiplication by 3 accounts for the input modalities: RTG, state, and action. We apply causal masking, similar to the masked attention approach, to prevent referencing future timestep information in predictions. Therefore, the learned parameters in each layer have the shape of a lower-triangular matrix with dimensions of $3K \times 3K$.

Reviewer comment (Experiments):
Can you also test the OOD generalization on novel tasks with multi-task learning?

Author response:
Thank you for your insightful suggestion regarding the exploration of out-of-distribution (OOD) generalization in novel tasks with multi-task learning.
We acknowledge that this is a meaningful direction for future research, but we feel that this is beyond the scope of the current paper.    At present, the field of offline multi-task/meta-RL, especially concerning OOD generalization, is still in its early stages of development, as indicated by recent literature [B.1,B.2,B.3]. The problem becomes more difficult when it is combined with return-conditioned approaches. These approaches face significant hurdles in multi-task settings, primarily due to the diverse scales and structures of rewards encountered across various complex tasks.

While our current work with the DC algorithm does not directly address these advanced aspects of OOD generalization in multi-task learning, we believe that the insights and methodologies developed here could be useful for future investigation in this area.

Reviewer comment (Experiments):
Can you visualize the attention of that hybrid model in some Atari games?

Author response:
We plotted the attention maps of hybrid DC for Atari Breakout and Assault in Section 3 of our anonymous link [https://sites.google.com/view/decision-convformer-iclr2024](https://sites.google.com/view/decision-convformer-iclr2024).
 Since the attention module is used only at the last layer in hybrid DC, we plotted the map of the final layer's attention module. It is seen that the attention map of the last attention layer is quite spread out in the considered  Atari games where strong Markov property does not seem to hold.
As discussed in our response to Q2, in hybrid DC, since local associations are captured by the previous convolution layers, considering a global perspective in the last layer seems to further enhance the performance of decision-making on domains with weak Markov property.

Reviewer comment (Experiments):
My major concern is the seemingly incomplete justification of the proposed architectures. In my understanding, just DT with a small K (i.e., short context length) could be sufficient and show comparable to or even better than DC. Additional comparisons (in terms of performance and computational complexity) on DC and DT with different choices of K would be helpful. Otherwise, it is unclear whether the gain of DC (or hybrid DC) is mainly from the good combinations of hyperparameters (including the embedding dimension, GELU, K, ...), or indeed the convolution filtering.

Author response:
Thank you for the detailed feedback and valuable suggestions regarding our paper. We have incorporated additional analysis and experimental results as recommended, to enhance the clarity and significance of our study. We believe the revision significantly strengthens the manuscript and better shows its significance.

### **W1. Origin of DC's Gain.**
We wish to direct the reviewer's attention to Appendices G.2 and G.3 of our original manuscript, which investigates DC's gain regarding the context length. Specifically, Appendix G.2 elaborates on how different context lengths $K$ and filter sizes $L$ influence DC, while Appendix G.3 provides the performance of DT with respect to context length $K$. For a clearer understanding, we will note in the main text of our manuscript that the content about the context length and window size of DC and DT is in the appendix.

Here, we briefly summarize the results in Appendices G.2 and G.3.  Table R1, which is equivalent to Table 18 in our manuscript, shows the performance across different context length $K$ and filter size $L$  for DC on MuJoCo hopper-medium. Since RTG, state, and action form the input at a particular time step, $L$ is a multiple of 3 for DC. Here, $L=3$ corresponds to the case that only the current time step is considered, and $L=6$ corresponds to the case that the convolution filter covers two time steps. The result in Table R1 shows that having a longer context length yields a slight performance enhancement although the gain is not substantial.

**Table R1.** Expert-normalized returns averaged across five seeds in hopper-medium for different combinations of $K$ and $L$.
|  K  \ L | 3    | 6    | 30   |
|-------------------|------|------|------|
| 8                 | 83.5 | 92.5 | -    |
| 20                | 90.1 | 94.2 | 93.5 |

Regarding DT with small $K$ focusing more on local information, as suggested by the reviewer, we already performed experiments and the results are in Table 19 in Appendix G.3 in our manuscript, which is shown again in Table R2 below.
As seen in Table R2,  the performance of DT deteriorates as the context length $K$ decreases.

**Table R2.** Performance of DT across context lengths K= 20, 8, and 2.
| Dataset                                  | DT (K=20) | DT (K=8) | DT (K=2) | DC (K=8) |
|------------------------------------------|-----------|----------|----------|----------|
| hopper medium & medium-replay            | 77.0      | 74.8     | 72.5     | **93.4** |
| walker2d medium & medium-replay          | 73.4      | 72.5     | 71.6     | **77.9** |
| halfcheetah medium & medium-replay       | 39.8      | 39.3     | 39.4     | **42.2** |

The results in Tables R1 and R2 can be explained from two perspectives.

1) **Input length** (historical information): it's important to note that, as indicated in [A.1], a transformer like DT, which operates with a global viewpoint, benefits from extensive historical data. The comprehensive historical context allows the transformer to better understand the policy that influenced past actions, thus improving learning efficiency and potentially optimizing training dynamics.

2) **Output length** (joint action optimization): Reducing the context length $K$ can be detrimental because it loses the benefit from the joint prediction of multiple sequential actions. Note that  the objective function of DT and DC is given by $$ \mathcal{L}_{\text{DT or DC}} := E\_{\tau \sim D} [ \frac{1}{K} \sum\_{t=1}^{K}  ( a_t  - \pi\_{\theta}(\hat{R}\_{t-K+1:t}, s\_{t-K+1:t}, a\_{t-K+1:t-1}) )^2 ].$$

    So, DT or DC are trained not only to predict the last action $a_t$ but also to predict the previous actions $a_{t-K+1}$ to $a_{t-1}$ in a causal manner. So, if we set $K$ as a small number, say, $K=2$, then the training guidance comes from only two terms. On the other hand, if we set $K$ large, e.g., $K=20$, we can exploit the strong training guidance from 20 terms that are intertwined. It seems that this joint optimization yields benefits in training.

Reviewer comment (Experiments):
Can you provide evaluations of the hybrid DC for the benchmark in Table 1 (environments with Markov property)? This would help choose architectures when the prior knowledge of the degree of Markov property is limited. If the hybrid DC is comparable to or better than DC and computational cost is not important, then one may simply consider the hybrid DC for such cases.

Author response:
### **Q1. Hybrid DC on MuJoCo and Antmaze.**

Below are the results of hybrid DC on MuJoCo and Antmaze.

**Table R5.** Offline performance results of DT, DC and DC$^\text{hybrid}$ on Mujoco and Antmaze.
|                         | DT    | DC    | DC$^\text{hybrid}$ |
|-------------------------|-------|-------|--------------|
| halfcheetah-medium      | 42.6  | **43.0** | **43.0**     |
| hopper-medium           | 68.4  | **92.5** | 86.8         |
| walker2d-medium         | 75.5  | **79.2** | 78.5         |
| halfcheetah-medium-replay | 37.0  | **41.3** | 40.1         |
| hopper-medium-replay    | 85.6  | **94.2** | 93.7         |
| walker2d-medium-replay  | 71.2  | **76.6** | 75.1         |
| halfcheetah-medium-expert | 88.8  | **93.0** | 92.8         |
| hopper-medium-expert    | 109.6 | 110.4 | **110.6**    |
| walker2d-medium-expert  | 109.3 | **109.6** | 108.9       |
| antmaze-umaze           | 69.4  | **85.0** | 81.2         |
| antmaze-umaze-diverse   | 62.2  | **78.5** | 72.3         |
| **mean**                | 76.4  | **82.2** | **80.3**     |


As shown in Table R5, hybrid DC does not outperform DC in these domains but still shows relatively superior performance compared to DT. Therefore, for tasks with strong Markov property, convolution layers alone are sufficient. However, for tasks where strong Markov property does not hold, hybrid DC is a good alternative.

Reviewer comment (Evaluation):
Can you report the computational complexity of the hybrid DC as you did for DT and DC in Table 14, 15, 16?

Author response:
### **Q2. Computational Complexity of the Hybrid DC.**

The complexity of hybrid DC can be seen in Tables R6, R7, and R8.

**Table R6.** The resource usage for training DT, DC, DC$^\text{hybrid}$ on MuJoCo and Antmaze.
| Complexity              | DT     | DC     | DC$^\text{hybrid}$ |
|-------------------------|--------|--------|--------------|
| Training time (s)       | 426    | 396    | 412          |
| GPU memory usage        | 0.7GiB | 0.6GiB | 0.7GiB       |
| All params #            | 1.1M   | 0.8M   | 0.9M         |
| Token mixer params #    | 198K   | 8K     | 71K          |

**Table R7.** The resource usage for training ODT, ODC, ODC$^\text{hybrid}$ on MuJoCo and Antmaze.
| Complexity              | ODT    | ODC    | ODC$^\text{hybrid}$ |
|-------------------------|--------|--------|---------------|
| Training time (s)       | 2147   | 854    | 1842          |
| GPU memory usage        | 4GiB   | 1.4GiB | 3GiB          |
| All params #            | 13.4M  | 8.1M   | 9.5M          |
| Token mixer params #    | 4202K  | 43K    | 1083K         |

**Table R8.** The resource usage for training DT, DC, DC$^\text{hybrid}$ on Atari.
| Complexity              | DT     | DC     | DC$^\text{hybrid}$  |
|-------------------------|--------|--------|---------------|
| Training time (s)       | 764    | 193    | 725           |
| GPU memory usage        | 3.7GB  | 1.8GB  | 3.1GB         |
| All params #            | 2.1M   | 1.7M   | 1.8M          |
| Token mixer params #    | 396K   | 16K    | 79K           |

Please be aware that we used the PyTorch grouped conv1D on our implementation, which is known for its inefficiencies within the PyTorch framework, as discussed in [https://github.com/pytorch/pytorch/issues/73764](https://github.com/pytorch/pytorch/issues/73764). We expect that the training time could be reduced by either employing a different CUDA kernel or opting for an alternative library like JAX.