Paper: Selective Visual Representations Improve Convergence and Generalization for Embodied AI (ICLR-2024)

Reviewer comment (Experiments):
Given that the proposed approach is about the codebook module, it would have been good to see if this approach could be applied to other pretrained visual encoders, for eg ViT based models.

Author response:
**W1) Applying the codebook to other pretrained visual encoders**


Thanks for the great feedback! We added section **A.2** in the appendix. In order to show the applicability of the codebook to other visual encoders, we used pretrained **DINOv2**[1] visual features and applied the codebook to bottleneck the new goal-conditioned representations. We use the frozen **DINOv2 ViT-S/14** model to encode RGB images into a 384x7x7 tensor. We fuse this tensor with a 32-dimensional goal embedding and the previous action embedding and flatten the result to obtain a 1574-dimensional goal-conditioned observation embedding. We employed a codebook with similar dimensions, K = 256 and Dc = 10, to bottleneck the goal-conditioned representations. As shown in **Table 5** (shown in below as well), our approach outperforms the DINOv2 baseline models across a variety of Object Navigation metrics in various benchmarks. This experiment underscores the capability of our codebook module in effectively bottlenecking other visual features for embodied-AI tasks.

| Benchmark                   | Model              | SR(%)↑ | EL↓   | Curvature↓ | SPL↑  | SEL↑  |
|-----------------------------|--------------------|-------:|------:|-----------:|------:|------:|
| **ProcTHOR-10k (test)**     | DINOv2             | 74.25  | 151.00| 0.24       | 49.53 | 43.20 |
|                             | +Codebook (Ours)   | **76.31** | **129.00**| **0.12**   | **50.26** | **44.70** |
| **ArchitecTHOR (0-shot)**   | DINOv2             | 57.25  | 218.00| 0.25       | 36.83 | 29.00 |
|                             | +Codebook (Ours)   | **59.75** | **194.00**| **0.11**   | **36.00** | **31.70** |
| **AI2-iTHOR (0-shot)**      | DINOv2             | 74.67  | 97.00 | 0.19       | 59.45 | 26.50 |
|                             | +Codebook (Ours)   | **76.93** | **68.00** | **0.07**   | **60.14** | **28.30** |
| **RoboTHOR (0-shot)**       | DINOv2             | 60.54  | -     | -          | **29.36** | -     |
|                             | +Codebook (Ours)   | **61.03** | -     | -          | 28.01 | -     |


[1] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, et al. Dinov2: Learning robust visual features without supervision. arXiv preprint arXiv:2304.07193, 2023

Reviewer comment (Experiments):
It's not clear how their approach compares against a full-scale visual encoder fine tuning baseline? Finetuning of the visual encoder should also allow it to forget information that is not relevant to the task. While finetuning is computationally expensive, it will still be interesting to see how EmbCLIP-codebook performs with respect to it.

Author response:
**W2) Full-scale visual encoder fine-tuning as a baseline**


Thank you for the suggestion. To explore the impact of full-scale fine-tuning, we fine-tuned the entire policy, including the visual encoder, for both EmbCLIP and EmbCLIP-Codebook models. We used the best checkpoints reported in **Table 1** and fine-tuned the entire policy, including the ResNet-50 visual encoder, for an additional 30 million steps. Details of the implementation, including multi-node training, large batch size, and small learning rate, can be found in **Appendix 6**. The results are presented in the following table (also included in the updated paper as **Table 7**).

| Benchmark                   | Model                                    | SR(%)↑ | EL↓ | SPL↑ |
|-----------------------------|------------------------------------------|-------:|----:|-----:|
| **ProcTHOR-10k (test)**     | EmbCLIP                                  | 67.70  | 182.00 | 49.00 |
|                             | EmbCLIP-Codebook                         | 73.72  | 136.00 | 48.37 |
|                             | EmbCLIP + CLIP ResNet fine-tuning        | 74.37  | 138.00 | 52.46 |
|                             | EmbCLIP-Codebook + CLIP ResNet fine-tuning | **79.38** | **120.00** | **52.67** |
| **ArchitecTHOR (0-shot)**   | EmbCLIP                                  | 55.80  | 222.00 | 38.30 |
|                             | EmbCLIP-Codebook                         | 58.33  | 174.00 | 35.57 |
|                             | EmbCLIP + CLIP ResNet fine-tuning        | 59.00  | 182.00 | **38.92** |
|                             | EmbCLIP-Codebook + CLIP ResNet fine-tuning | **62.58** | **168.00** | 37.10 |
| **AI2-iTHOR (0-shot)**      | EmbCLIP                                  | 70.00  | 121.00 | 57.10 |
|                             | EmbCLIP-Codebook                         | **78.40** | **86.00** | 54.39 |
|                             | EmbCLIP + CLIP ResNet fine-tuning        | 72.53  | 88.00 | 59.28 |
|                             | EmbCLIP-Codebook + CLIP ResNet fine-tuning | 77.75 | 99.00 | **59.33** |

As the table shows, while there is a substantial improvement from fine-tuning the entire policy model, including the visual encoder, compared to using the frozen one, a noticeable gap remains between the fine-tuned EmbCLIP and fine-tuned EmbCLIP-Codebook. Notably, with further end-to-end fine-tuning, EmbCLIP-Codebook achieves significantly superior results compared to all other models.

Reviewer comment (Experiments):
It is also not clear how important is the use of the codebook vectors compared to just the introduction of a bottleneck in the architecture? To test this, I recommend the authors try using their proposed bottleneck architecture without the codebook vectors.

Author response:
**W3) Bottlenecked architecture without the codebook module as a baseline**


Thanks for the insightful comment. We included a new section, **A.4**, in the appendix. **Table 6** compares our proposed bottleneck architecture (EmbCLIP-Codebook) with 2 other information bottleneck baselines on 3 Object Navigation benchmarks. **EmbCLIP-AE** utilizes an auto-encoder on the goal-conditioned embedding E. This autoencoder comprises a series of linear layers, mapping the representation to $\mathcal{P} \in \mathcal{R}^{256}$ and $h \in \mathcal{R}^{10}$ before resizing back to $\hat{E} \in \mathcal{R}^{1574}$. The results in Table 6 demonstrate the superior performance of our method compared to EmbCLIP-AE across all 3 benchmarks which underscores the importance of the codebook module in the proposed architecture.

Reviewer comment (Experiments):
Learning codebook as information bottleneck helps discriminative mapping of observations to task outputs. Meanwhile, throwing away information means that there's little knowledge shared among different task. It'd be interesting to see if the performance comparison on zero-shot/few-shot learning of novel tasks. My guess is that EmbCLIP would adapt to the novel task faster and better than the proposed codebook method.

Author response:
**W1) Throwing away information using the codebook might result in worse performance in zero-shot/few-shot learning of novel tasks**

Bottlenecking and discarding information does not necessarily lead to poorer generalization in novel scenarios. An illustrative example of this is found in regularization techniques such as dropout, wherein they effectively prevent overfitting, resulting in simpler and more generalizable models. Furthermore, our codebook module is specifically designed to be **task-specific**. Our main goal is to bottleneck the information to best suit the task in hand. Therefore, performing a new task requires training another task-specific codebook module. It is however expected of our proposed approach to be more generalizable when applied to the **same task in new domains**. 
As demonstrated in **Section 4.2** of the paper, our method significantly outperforms non-bottlenecked architectures in adapting to new visual domains through the lightweight fine-tuning of a small adaptation module.

Reviewer comment (Experiments):
The paper does not provide comparison to other bottleneck-based baselines. For example, one can learn a self-attention modules atop CLIP feature maps. Also, learning an auto-encoder, where the bottleneck is low-dimensional latent feature, seems to be another reasonable baseline. Both self-attention modules and the auto-encoder will condition on the goal task description and previous action.

Author response:
**W2, Q3) Comparison with other information bottleneck baselines**

Thanks a lot for this insightful feedback. We included a new section, **A.4**, in the appendix. **Table 6** compares our proposed bottleneck architecture (EmbCLIP-Codebook) with 2 other information bottleneck baselines on 3 Object Navigation benchmarks. **EmbCLIP-AE** utilizes an auto-encoder on the goal-conditioned embedding E, while **EmbCLIP-SelfAttn** applies self-attention to the goal-conditioned CLIP feature maps. The results in Table 6 demonstrate the superior performance of our method compared to the other information bottleneck baselines across all 3 benchmarks.

Reviewer comment (Experiments):
What's the performance of fine-tuning entire model on Habitat and adaptation module on HMD semantics (table 2)? Also, what's the before-finetuning results of EmbCLIP+codebook on both benchmark? Right now the comparison is conducted over different dataset, not apple-to-apple.

Author response:
**Q1) Habitat zero-shot results and fine-tuning results on the same benchmark**

Sorry for the confusion. All results are originally evaluated on the same benchmark (HM3D Semantics dataset) used in Habitat 2022 ObjectNav Challenge. We have updated **Table 2** to clarify our fine-tuning results. 
We’re also reporting the zero-shot results for both models here. The models are evaluated at 220M and 420M training steps. We’re not including the results in the paper as the 0-shot results are very low compared to the fine-tuned versions.




| Benchmark              | Model            | Training Step | SR(%)↑ | EL↓ | SPL↑      |
| ---------------------- | ---------------- | ------------- | ------------ | -------------- | -------- |
|**Habitat Challenge 2022** (zero-shot) | EmbCLIP          | 220M          | **11.20**    | 313            | **6.50** |
| **Habitat Challenge 2022** (zero-shot) | EmbCLIP-Codebook | 220M          | 8.15         | **253**        | 4.71     |
|**Habitat Challenge 2022** (zero-shot) | EmbCLIP          | 420M          | 8.50         | 362            | 4.85     |
|**Habitat Challenge 2022** (zero-shot) | EmbCLIP-Codebook | 420M          | **9.85**     | **345**        | **5.23** |

Reviewer comment (Experiments):
Perhaps the authors should compare with other information-bottleneck baselines, e.g. self-attention / autoencoder. Such results could clarify if codebook / information bottleneck is the key.

Author response:
**W2, Q3) Comparison with other information bottleneck baselines**

Thanks a lot for this insightful feedback. We included a new section, **A.4**, in the appendix. **Table 6** compares our proposed bottleneck architecture (EmbCLIP-Codebook) with 2 other information bottleneck baselines on 3 Object Navigation benchmarks. **EmbCLIP-AE** utilizes an auto-encoder on the goal-conditioned embedding E, while **EmbCLIP-SelfAttn** applies self-attention to the goal-conditioned CLIP feature maps. The results in Table 6 demonstrate the superior performance of our method compared to the other information bottleneck baselines across all 3 benchmarks.