Selective Visual Representations Improve Convergence and Generalization for Embodied AI

We thank you for your valuable feedback. In the following, we use W for Weaknesses and Q for Questions.

W1) Applying the codebook to other pretrained visual encoders

Thanks a lot for your great comment! We included section A.2 in the appendix which shows the codebook module is representation-agnostic and can be applicable to other pretrained visual encoders as well. We used pretrained DINOv2[1] visual features and used 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.

[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

Q1) Codebook as a bridge between modalities?

Yes, our model encodes goal related information beyond just the appearance of the goal (location, context, places, etc). Our model bridges between task-related information and visual information. Our tasks are represented as an embedding and our codebook connects the two. Figure 8 shows the encoded information in different latent codes. While some codes are directly responsible for goal-related information such as goal visibility and proximity to the agent, others encode other information required for the object navigation task such as walkable areas, corners, etc.

Q2) What are the failure cases of the EmbCLIP-Codebook?

We included section A.5 which analyzes the failure cases for EmbCLIP-Codebook. Fig. 9 illustrates examples of two modes of failure in our EmbCLIP-Codebook agent: perception and exploration. Although the codebook module enables the agent to effectively filter out distractions and concentrate on the target object, the agent’s performance remains constrained by the perceptual capabilities of the pretrained visual encoder. So there exists instances where the codebook fails to identify target objects. The top row of Fig. 9 showcases examples of failures related to perception. These instances predominantly occur when the goal object is either too small or challenging to identify (the baseball bat on the table in the top left example). In these scenarios, although the agent traverses past the object, it fails to accurately locate the target. The second row of the figure presents additional instances of failure, wherein the agent fails to explore specific areas of the environment where the target object is located.

Q3) Are there cases where two objects of the same class are present in the scene?

Yes, that is indeed possible. Take, for example, a scenario where the target object is a 'chair.' In a typical scene, multiple chairs are often found around dining tables. The objective of Object Goal Navigation is to find any one instance of the target object type – in this case, any chair in the scene. Therefore, the task requires the agent to find any single instance of the specified object type.

We thank you for your valuable feedback. In the following, we use W for Weaknesses and Q for Questions.

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.

[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

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).

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.

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.

Q1) What is meant by “Samplers” mentioned in experiment details?

The concept of a Sampler is defined in AllenAct [1]. The primary function of a Sampler is to sample an episode (or a Task) for the agent. Each episode specifies the action space, success criteria, and the rewards that the environment returns to the agent. After the agent completes an episode, the Sampler selects a new episode for the agent. Consequently, a greater number of Samplers indicates that a higher number of episodes are being processed by the agent in parallel. In other words, more Samples imply a larger batch size for training a policy model.

[1] Luca Weihs and Jordi Salvador and Klemen Kotar and Unnat Jain and Kuo-Hao Zeng and Roozbeh Mottaghi and Aniruddha Kembhavi, “AllenAct: A Framework for Embodied AI Research”, Arxiv 2020, url: https://allenact.org/.

Q2) Why is RoboTHOR missing some metrics?

RoboTHOR test results are only available through submissions to the leaderboard (link). We also included the link to the submission in a footnote on page 6. The leaderboard metrics only include Success Rate and SPL.

We thank you for your valuable feedback. In the following, we use W for Weaknesses and Q for Questions.

W1) Incorporating text prompts as a human prior into the models

That’s a very interesting suggestion! It can definitely be applied as a future work of our approach. In our current proposed method, we anticipate the automatic extraction of such common-sense knowledge and human priors from the observations through the bottleneck. Integrating such priors in the form of language could indeed be a compelling direction for future work.

W2) Codebook collapse ablations

Thanks for the insightful comment. We included an ablation study in section A.7. Codebook collapse is a prevalent issue defined as the assignment of high probabilities to a specific subset of latent codes in the codebook, causing the majority of codes to be underutilized. While it’s important to note that our codebook module doesn’t experience a complete collapse, incorporating regularization into the codebook enhances overall performance by promoting a more balanced utilization of all the latent codes. While various solutions have been suggested, we have determined that applying a simple dropout on top of the codebook probabilities is the most effective and straightforward approach. In Fig. 10, we compare the average codebook probability distributions between two models—one trained with a 0.1 dropout probability and the other without any dropout. The figure shows that this regularization leads to a more uniform average usage of the latent codes. Table 8 presents the corresponding evaluation results. Clearly, introducing such regularization to the codebook module results in an improved performance. We additionally tried the Linde-Buzo-Gray[1] splitting algorithm as an alternative approach to balance the codebook usage. More specifically, during the training process, if there is a code vector that is underutilized, we perform a split on the most frequently used latent code, creating two equal embeddings and replacing the unused one. The corresponding training curve is depicted in Fig. 11, which shows a severe imbalance and eventual collapse during the training.

[1] Yoseph Linde, Andres Buzo, and Robert Gray. An algorithm for vector quantizer design. In IEEE Transactions on communications, 1980.

Q1) What’s the effect of adding a skip connection?

We observed that incorporating a skip connection into our architecture leads to poorer performance, likely due to overfitting. There is a possibility that our existing hyperparameters may need adjustment to better suit the architecture with the skip connection, but despite our efforts, we were unable to identify the optimal hyperparameters, given the time constraints of the discussion period. We believe this requires further investigation. Nonetheless, based on our current findings, it appears that employing a full bottleneck yields better results than introducing a skip connection. Table below shows the results for 300 million steps of training.

We thank you for the valuable feedback. In the following, we use W for Weaknesses and Q for Questions.

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.

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.

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.

Q2) Nearest-neighbor analysis using learned codes as queries

Thanks for the interesting suggestion! We added this analysis in section A.3 in the appendix. Using a set of 10k ProcTHOR frames each accompanied by a corresponding goal specification that may or may not be visible in the frame, we input these frames and their associated goals into our pretrained EmbCLIP-Codebook model. We save their corresponding 10-dim hidden compact representations h which are convex combinations of the latent codes. To visually assess the information encoded in each individual learnable code, we perform nearest neighbor retrieval on the dataset using each 10-dim codebook entries as queries. Fig. 8 illustrates the 8 nearest neighbors for 9 different codes. The figure reveals that certain codes primarily contribute to goal visibility, determined by their proximity to the agent, while others capture various geometric and semantic details of scene layouts, such as walkable areas, corners, doors, etc.