Selective Visual Representations Improve Convergence and Generalization for Embodied AI
Authors: Ainaz Eftekhar
Kuo-Hao Zeng
Jiafei Duan
Ali Farhadi
Aniruddha Kembhavi
Ranjay Krishna
Venue: ICLR 2024 spotlight (2024)
Keywords: Embodied-AI
Task-conditioned Representations
Visual Navigation
Reinforcement Learning
Primary area: reinforcement learning

Abstract: Embodied AI models often employ off the shelf vision backbones like CLIP to encode their visual observations. Although such general purpose representations encode rich syntactic and semantic information about the scene, much of this information is often irrelevant to the specific task at hand. This introduces noise within the learning process and distracts the agent's focus from task-relevant visual cues.
Inspired by selective attention in humans—the process through which people filter their perception based on their experiences, knowledge, and the task at hand—we introduce a parameter-efficient approach to filter visual stimuli for embodied AI.
Our approach induces a task-conditioned bottleneck using a small learnable codebook module. This codebook is trained jointly to optimize task reward and acts as a task-conditioned selective filter over the visual observation.
Our experiments showcase state-of-the-art performance for object goal navigation and object displacement across $5$ benchmarks, ProcTHOR, ArchitecTHOR, RoboTHOR, AI2-iTHOR, and ManipulaTHOR. The filtered representations produced by the codebook are also able generalize better and converge faster when adapted to other simulation environments such as Habitat. Our qualitative analyses show that agents explore their environments more effectively and their representations retain task-relevant information like target object recognition while ignoring superfluous information about other objects.

Published as a conference paper at ICLR 2024
SELECTIVE VISUAL REPRESENTATIONS IMPROVE CON- VERGENCE AND GENERALIZATION FOR EMBODIED AI
Ainaz Eftekhar" 1,2 Ali Farhadi
Kuo-Hao Zeng x2 Ani Kembhavi-.2
Jiafei Duan 1.2
Ranjay Krishnal.2
University of Washington, Allen Institute for Artificial Intelligence { (ainaze, khzeng, jiafeid, alif, anik, ranjayk/ @allenai. org
http:://embodiied-oddebooooo github.io
ABSTRACT
Embodied AI models often employ off the shelf vision backbones like CLIP encode their visual observations. Although such general purpose representations encode rich syntactic and semantic information about the scene, much of this information is often irrelevant to the specific task at hand. This introduces noise within the learning process and distracts the agent's focus from task-relevant vi- sual cues. Inspired by selective attention in humans- the process through which people filter their perception based on their experiences. knowledge, and the task at hand-ww we introduce parameter-effiiien approach to filter visual stimuli for embodied AI. Our approach induces a cundicondttooned bottleneck using a small learnable codebook module. This codebook is trained jointly to optimize task reward and acts as a cask-conditioned selective filter over the visual observation. Our experiments showcase state-of-ther tfrt--art- performance for object goal naviga- tion and object displacement across benchmarks, ProcTHOR, ArchitecTHOR, RobOTHOR, AI2-ITHOR, and ManipulaTHOR, The filtered representations pro- duced by the codebook are also able generalize better and converge faster when adapted to other simulation environments such as Habitat. Our qualitative analyses show that agents explore their environments more effectively and their representa- tions retain task-relevant information like target object recognition while ignoring superfluous information about other objects. Code is available on the project page.
INTRODUCTION
Human visual perception is far from passive reception of all available visual stimuli; is an actively tuned mechanism that operates selectively, allocating attention and processing stimuli that are deemed relevant the current task (Bugelski & Alampay, 1961). An illustrative example of this phenomenon is the common experience of misplacing one's keys; we become subsequently oblivious to most visual cues in our environment, except for those directly related to the search for the lost keys. this case, we become particularly attentive to surfaces where we usually place our keys and navigate our environment by similarly processing the walkable areas around us (Figure 1).
The subfield of embodied artificial intelligence (AI) studies Al agents tasked with very analogous situations (Duan et al., 2022). Embodied AI tasks such as navigation (Batra et al., 2020b; Krantz et al., 2023), instruction following (Anderson et al., 2018; Krantz et al., 2020; Shridhar et al., 2020), manipulation (Ehsani et al., 2021; Xiang et al., 2020), and rearrangement (Batra et al., 2020a; Weihs et al., 2021a). necessitate dirl-eirected behaviors, such navigating specific goal or relocating target objects. Conventional frameworks usually employ general-purposs visual backbones (Khandelwal et al., 2022; Yadav al., 2023b) to extract representations from visual input. These representations are then fused with goal embeddings (e.g., object type. images, or natural language instructions) construct a conl-condiiioned representation E R where often has dimensions as large as a 1568-dimm dise-dimensional Hilbert space. E captures an abundance of details from the visual input, of which policy determines which action to take next. Given E's general-purposs nature, it often contains a significant amount of task- iareliirelevant information. For example to find specific object house, the agent doesn't need to know about other distractor objects in the agent's
"Equal contribution

---

Published as a conference paper at ICLR 2024
Bottom-Up Processing
Top-Down Selective Attention
Input Frame
Inpul Frame
saliing
TV Comen gaan
Chandeler
p windows
Block
Encoder with Codebook bottleneck
authiane nighht
Enpoder
Seuton solas
Coffee thtll with
Browm box
might
Wooden
contenn
Goal: Key
Walkable path
foor
Textiped
Goal: Key
the flat surllsse
Figure| 1: Selective Attention. Imagine an agent is tasked to locate a key an environment. Standard visual encoders such as CLIP encoder capture general purpose scene information which include details not relevant to the task, such as the color of the sofa or texture of the floor. This mirrors the concept of bottom-up up processing. where perception influenced by external stimuli in the environment. To address this, we equip the encoder with a codebook bottleneck that only retains the most task-relevant information such identifying flat surfaces likely to hold the key and the walkable paths to these surfaces. This represents top-down selective processing where the perception is guided by internal goals and expectations.
view, about their colors, materials, attributes, etc. These distractions introduce unnecessary noise into the learning process, distracting the agent's is focus away from more pertinent visual cues.
In this paper, we draw from the substantial body of research in cognitive psychology and neuroscience to induce selective task-specific representations that filter irrelevant sensory input and only retain the necessary stimuli. In Psychology, inattentional blindness posits that people become "blind" to objects and details irrelevant to the task at hand (as depicted by the famous invisible gorilla video on YouTube') (Simons & Chabris, 1999). Similar studies in human visual search (Wolfe, 1994; Treisman & Gelade, 1980) find that people become oblivious to the distracting items when searching through a cluttered environment. The principle of top-down selective attention (Desimone & Duncan, 1995; Buschman & Miller, 2007) also projects a similar theory: that our goals bottleneck our visual processing guided by internal goals (in contrast to bottom-up processing, which is how vision encoders are usually trained).
To apply this principle of selective attention to embodied Al agents, we propose a parameter-efficient approach to bottleneck visual representations conditioned on the task. We insert a simple codebook module into our agent, which consists of a collection of K = 256 learnable latent codes, each with a dimension of D. = 10, where D. < D. This codebook module accepts E as input, attends over the K latent codes, to produce a codebook representation É. Σ is weighted convex combination of the K latent codes, weighted by the attention estimates (see Figure 2). Given the bottleneck induced by needing to choose amongst K = 256 possible D. = 10-dimensional representations, the codebook module enforces selective filtering, encoding only the essential cues necessary for the task.
We demonstrate the effectiveness of our approach by achieving zero-shot, state-of-the-ar perfor- mances on 2 Embodied-A AI tasks: object goal (avigation(ObiNNa)) (Deitke et al., 2020) and object displacement (ObjDis) (Ehsani et al., 2022) across 5 benchmarks (PrecTHOR(Deitke et al., 2022b), ArchitecTHOR, RoboTHOR (Deitke et al., 2020), A12-TTHOR(Kolve et al., 2017a), and Manip- ulaTHOR (Ehsani et al., 2021)). Across all benchmarks, our approach yields significant absolute improvements in success rate and reductions in episode length. Moreover, we further show the adapt- ability of our codebook representations new visual domains by lightweight finetuning in Habitat environments (Savva et al., 2019). Finally, we conduct comprehensive analysis, verifying that the codebook encodes better information about the task, about the distance to the goal for navigation, etc. while its ability to identify individual objects in its view diminishes significantly. Surprisingly, we observe noticeable improvements in the agent's behavior in the form of smoother trajectories and more efficient exploration strategies.
RELATED WORK
We situate our work amongst methods that learn performant representations for Embodied AI tasks. Learning representation through proxy objectives. A common approach to learning useful representations is defining proxy objectives for the visual encoder. For example, a paradigm outlines a set of invariant transformations and learns representations with a contrastive learning objective while simultaneously optimizing task reward (Du et al., 2021; Singh et al., 2023; Guo et al., 2018; 8;; Laskin
I Link to invisible gorilla video depicting selective attention
2

---

Published as a conference paper at ICLR 2024
et al., 2020; Majumdar et al., 2022). Others use auxiliary tasks, such as predicting depth (Mirowski
et al., 2017; Gordon et al., 2019), reward (Jaderberg et al., 2017), forward dynamics (Gregor et al., 2019; Zeng et al., 2023; 2021; Guo et al., 2020; Kotar et al., 2023), inverse dynamics (Pathak et al., 2017; Ye et al., 202 b,:), scene graphs (Gadre et al., 2022; Du et al., 2020; Savinov et al., 2018; Wu et
or 2D maps (Chaplot et al., 2020b;a). In contrast, our proposed codebook module is jointly trained with the overall policy to optimize the task rewards only.
Learning representations through pretraining. Alternately, many proposed methods self-supervise representations by sampling image frames from the environment during pre-training stage ( (Yaday et al., 2023b;a; Kotar et al., 2023; Mezghan et al., 2022). Some utilize 3D constraints induce 3D- awareness into the visual representations (Wallingford et al., 2023; Marza et al., 2023). There are even been complicated games designed to encourage the emergence of useful representations (Weihs et al., 2021b). Of particular interest is the EmbCLIP (Khandelwal et al., 2022) model, which serves as our main baseline. This method leverages the large-scale, pretrained representations by CLIP (Radford et al., 2021) the visual encoder. Our model builds upon EmbCLIP, incorporating the proposed codebook module to filter out redundant stimuli.
Bottleneck and regularize representation. Bottlenecked representations are common in machine learning (Hinton et al., 2011; 1; Kingma& & Welling, 2014; 4; Alemi et al., 2017; Koh et al., 2020; Federici al., 2020; Mairal al., 2008), computer vision (He et al., 2016; Srinivas al., 2021; Riquelme et al., 1021; Van Den Oord et al., 2017; Chen et al., 2023; Kolesnikov et al., 2022: Hsu et al., 2023), natural language processing (Shazeer et al., 2017: Fedus et al., 2022; Mahabadi et al., 2021). and even reinforcement learning (Serban et al., 2020; Fan & Li, 2022: Pacelli & Majumdar, 2020). The current most popular bottlenecked representation is VQ-GAN (Esser et al., 2021). Nonetheless, relatively few have implemented these techniques in Embodied AL. Most applications of information bottlenecks in this space, as far as we are aware, have been applied to full-observable MiniGrid (Chevalier- Boisvert et al., 2023) or Atari (Bellemare et al., 2013) environments (Goyal et al., 2019; Igl et al., 2019; Bai et al., 2021; Lu et al., 2020). By contrast, we situate our work within the partially-observed physics- enabled, and visually rich Embodied AI environments. The closest related work is recent proposal for learning a codebook for embodied AI using a NERF-based rendering objective (Wallingford et al., 2023); this method pre-trains the codebook and later uses it for downstream tasks. By contrast, we learn our codebook representations while simultaneously learning the policy to act.
3 METHODOLOGY
We first provide a brief background on how embodied AI agents are designed today. Next, we introduce our contributions to make the model. We lay out the details of our task-conditioned codebook and also describe our training strategy.
Background. The modern design for embodied agents share a common core architecture (Batra et al., 2020b; Khandelwal et al., 2022). This architecture contains three main components. First, three encoders encode the necessary information required to take an action. A visual encoder transforms visual observations (e.g., RGB or depth) into a visual representation w; goal encoder transforms the task's 's objective (e.g., GPS location, target object type, or natural language instructions) into a goal embedding g: a previous-action encoder to embed the most recently executed action into an action embedding a. For the visual encoder v, we mainly use the RGB sensors as visual stimuli. The only exception is the object displacement task, for which we follow prior work to also include depth sensor, the segmentation mask sensor, and the end-effector sensor (Ehsani et al., 2022). The second component of the architecture fuses U, g. a together. The three representations are flattened and then concatenated to form the task-conditioned representation F. Finally, the third component keeps track of the history of the agent's 's trajectory and proposes the next action. This component includes recurrent state encoder to encode and remember past steps. So, r = Eglt = 1,... I' encodes all the past representations into a single state representation. This is injected by an actor-critic head, which predicts a distribution over action space and also produces a value estimation of the current state.
The need for tast-bottlenecked visual representations. EmbCLII (Khandelwal et al., 2022) is today's ' -fate-of-the-art model for the tasks that we consider which encodes the three representations in E' E RD, where = 1568. This embedding contains CLIP features, which were trained for general-purpose vision tasks and guided through language self-supervision (Radford et al., 202 I). Therefore, when presented with the appropriate input image, this representation can identify large number of object categories, their attributes. their spatial relationships, etc. It can even encode the materials and textures of objects, and outline the actions that people appear to be taking in the image. However, for tasks in Embodied AI, where the agent is asked to locate pair of, say, lost "keys",
3

---

Published as a conference paper at ICLR 2024
Codebook Module
Input Frame
Encoder
Visual
E
Scoring Function
Canvex Combination
Upsampling Leyer
State Encoder
Actor Critic
Action
CNN
Linear
Softmax
Linear
Embeddler
Goal: Chair
256 x
Codebook
10 x
E (0....111)
g
x
x574
Embedder
Previous
Action
a
C E R 256x10
The 3 representations corresponding to the input frame, the goal, and the previous action get concatenated form F R. R1565 The codebook module takes E and generates a probability simplex P E R526 over the latent codes. The hidden compact representation E R2 is a convex combination of the codes weighted by P. The final task bakkbbotteeeekedd codebook representation É is derived by upsampling h which is subsequently passed to the recurrent state encoder and the policy to produce an action.
Figure 2:
An overview
knowing the texture of the couch in the living room or the color of the flower vase on the dining table seems irrelevant. These additional pieces of information, therefore, are also sent over to the policy, which is often designed as recurrent neural network (RNN) followed by a small actor-critic module. These few parameters, therefore, must serve two purposes: they must identify what information useful for the task at hand and what action to take given that information.
The codebook module. We introduce a module that decouples the two objectives across different parameters in the network. The input encoders and the codebook focus on extracting essential information for the task from the visual input, whereas the policy (RNN and actor-critic heads) can focus on taking actions conditioned this filtered information.
We design the codebook to be a parameter- module to transform the general- purpose representation E into compact lusk-bottlenecked one. The module inputs E and generates compact E (see Figure 2). This module contains codes, defined
set of latent vectors C [1, c,.... cx] E RxxD, where K denotes codebook's 's size and D. is the dimensionality of each latent code Cg. To create a strong bottleneck, we set D, 10 and K 256. These codes initialized randomly via a normal distribution and optimized along with the overall training algorithm.
To extract É from E, the module contains scoring function o(.) and an upsampling layer. We first use p o(E) to generate a probability simplex over the K latent codes, where P [p,] such that Σ P; = The scoring function & is implemented as a single-layer MLP followed by softmax function. This configuration forces the agent to select which latent code(s) are more useful for representing the current frame. Next, the hidden compact representation h is derived by taking convex combination of the learnable codes {ci]/] weighted by their corresponding pi: C = Σi=1 Pt.c, Further, the upsampling layer, implemented as a linear layer, upsamples the hidden embedding h to the task-bottlenecde codebook representation É.
By design, we place the codebook module immediately before the RNN and after the E is fused. In this way, the module focuses on single-step processing, prompting the codebook to compile information at each individual step rather than relying on past steps. By positioning the codebook module right after E, it fuses the visual representation with the goal g. We compare our design choice with the codebook module positioned right after the visual encoder (without coal-conditioning in Sec. A.9. In essence, this module design is reminiscent of memory networks, which keep track of information (Han et al., 2020), except without the write functionality; not requiring a write operation makes sense because make the assumption that everything the agent needs to act is already encoded in the CLIP embeddings.
Training algorithm. We employ Proximal Policy Optimization (PPO) (Schulman al., 2017)2 to train the agent. PPO is an on-policy reinforcement learning algorithm, which optimizes all three components of the architecture, including the codebook module. One critical challenge we
We use DD-PPO to optimize the policy for the object displacement task (Fhsani et al., 2022).
4

---

Published as a conference paper at ICLR 2024
Table I: Our method with the codebook outperforms the baselines on 2 different tasks, including zero-shot evaluation Object Goal Navigation benchmarks and results Object Displacement benchmark. We use and 4 denote if larger or smaller values are preferred.
Object navigation Curvature,
Beachmark
Model
SRCY)
EL,
SPLT
SELT
PmcTHOR-IIn 11k (validation)
EmbCLIP
67.70
182,00
0.58
+codebook
73.72
136.00
0.23
48.37
43.69
ARCHmTtTHOR (0-shot)
FmhCTTPP
55.80
222.00
0.49
38.30
20:57
+Codebook
58.33
174.00
0.20
35.57
28.31
EmbCLIP
51.32
24.24
RoboTHOR (0-shot)
+Codeboak
55.00
23.65
A12-TTHOR (0-shot)
+Codcbook
EmbCLIP
78,40
70.00
121.00 86.00
0.16
0.29
54,39
57.10
21.45
26.76
Object displacement
PU y
SR(S5]
ManipulaTHHO
m-VOLE
-Codebook
86.00
65.10
encounter when training codebooks is codebook collapse (Van Den Oord et al., 2017; Riquelme et al., 2021; Shazeer et al., 2017), where only a handful of codes are used by the agent, limiting the model's full capacity. While several solutions? have been proposed to this problem, using dropout (Srivastava et al., 2014) 4) is the straightforward yet effective approach to overcome this challenge. Instead of directly employing the probability scores P for the convex combination, first apply dropout to it with rate of 0.1. This produces new probability scores P. which we then use to combine the latent codes: P'C. Intuitively, this equates randomly setting 10% of the latent codes to zero before integrating them through convex combination, preventing the agent from relying on only a handful of codes.
find that
4 EXPERIMENTS AND ANALYSIS
In our experiments, demonstrate that the goal-conditioned codebook embedding F.r results in significant improvements over the non-bottlenecked embedding E' across a variety of Embodied-Ar benchmarks (Sec. 4.1). We further show that our bottlenecked embeddings are more generalizable and can be applied across different domains without exhaustive finetuning (Sec. 4.2). We study (both qualitatively and quantitatively) the visual cues encoded in the codebook and how they relate the target task (Sec. 4.3). Finally, we evaluate whether our approach is suitable for other pretrained visual encoders. We substitute the CLIP representations with DINOv2 (Oquab et al., 2023) visual features and replicate our experiments. The results confirm that our proposed codebook module is representation-agee -agnostic and can effectively bottleneck various pretrained visual representations.
4.1 CODEBOOK BAASED REPRESENTATION IMPROVE PERFORMANCE IN EMBODIED-AL
Tasks. We present our results on several navigation benchmarks and a manipulation benchmark to demonstrate the benefits of compressing the task-conditioned embedding using the codebook. We consider ObjectNav (navigate to find a specific object category in a scene) in PROCTHOR (Deitke et al., 2022a), ARCHIEECTHOR, ROBOTHOR (Deitke et al., 2020), A12-iTHOR (Kolve et al., 2017b). We also consider Object Displacement (bringing a source object to a destination object using a robotic arm) (Ehsani et al., 2022) in ManipulaTHoO (Ehsani et al., 2021) as our manipulation task.
4.1.1 OBJECT NAVIGATION
Model. We adopt the same core architecture used in EmbCLIP (Khandelwal et al., 2022) for Object Navigation. RGB images are encoded using a frozen CLIP ResNet-50 model to a 2048 x 7 x 7 tensor and then compressed by a a 2-layer CNN to another 32 x 7 x tensor v. This tensor is concatenated with the 32-dim goal embedding and a previous action embedding o, passed through another 2-layer CNN and flattened form D 1568-dim goal-conditioned observation embedding F. While EmbCLIP directly passes this 1568-dim representation, E, to to a 1-layer recurrent state encoder with 512 hidden units, use codebook module with codebook's size of K 256 and D, 10 as a bottleneck to compress this goal-conditioned embedding and filter out the irrelevant information. We further pass the resulting goal- botteeotteeeckkd codebook embedding É to the following recurrent state encoder. We term our model as EmbCLIP- Coodebook in the rest of the paper.
We tried Linde-Buzo-Gray Buzo-Gray (Linde et al., 1980) splitting algorithm and Gumbel Softmax (Jang et al., 2017) with careful temperature tunning, but have found that the dropout is the most effective approach.
5

---

Published as a conference paper at ICLR 2024
Table 2: Results for models trained on ProcTHOR and evaluated with fine- fine-tuning on Habitat Object Goal Navigation benchmarks. Our method with the codebook can adapt Habitat environment through the lightweight finetuning. We use and 4 denote larger or smaller values are preferred.
Object goal navigation
Benchmark
Fine-tuning parts
Model
SR(%)
SPL.1
Invalid Actions($1)
Curvature
EmbCLIP
36.45
18.18
28,10
0.53
Adaptation Module
Habitat challenge 2022
+Codebook
50.25
25.76
21.50
0.26
(HMBD Semantics)
EmbCLIP
58.00
30.97
15.80
0.52
Entire Model
+Codebook
55.00
29.21
15.40
0.38
Experiment Details. All models are trained using the AllenAct framework. We follow (Deitke et al., 2022b) to pretrain the Emb- CLIP baseline and CembCLIP-Codebone
EmbOLIP C-CCoeeboo
on the the PROCTHOR-IE 10k houses with 96 samplers for 435M steps. The models are evaluated zero-shot on the downstream benchmarks, including PROCTHOR- 10k val scenes, ARCHITECTHOR val scenes, ROBOTHOR test scenes", and A12-iTHOR val scenes. Each model is evaluated for 5 checkpoints between 415M to 435M training steps and we select the best model results from the ARCHITECTHÖR val scenes. Please refer to the suppiementary Sec. A. for more training details.
Goal
Starting Pose
EmbCLIP
Result. Table 1 shows EmbCLLP-Coddboob achieves best Success Rate (SR) across all 4 benchmarks. Furthermore, Episode Length (EL) is significantly lower compared to EmbCLIP suggesting that our agent is able to find the goal in much fewer steps. We further present Curvature as a key evaluation metric comparing the smoothness of the trajectories. Smoother trajectories are generally safer and more energy-effie energy-efficient. Excessive rotations and sudden changes in direction can lead to increased energy consumption and increase the chances of collisions with other objects. The curvature value k is defined as dexddy-dyx for each coordinate (z, y) on the trajectory and
Goal
Starting
Figure 3: Sample Trajectory. EmbCLIP agent takes many re- dundant rotations, resulting in a high average curvature, whereas ours navigates more smoothly.
(da2++dy2)
it's averaged across all time steps in the agent's 's path. A higher value of average Curvature signifies more rotations and changes in direction. Conversely, a value closer to zero suggests predominantly straight and smooth trajectory. Our agent travels in much smoother paths across all benchmarks. Figure: shows an illustrative example. Please find more examples in Supplementary Sec. A.7.
SPL Limitations and Introducing a New Metric SEL. We further report the Success Weighted by Path Length (SPL) as one of the common metrics used in the field for Object Navigation task. SPL is defined as 1 Em= St "max(...p)) where 4 is the shortest possible path (traveled distance) to the N
target object, p; is the taken path by the agent and S: is a binary indicator of success for episode i. As shown in the table, EmbCLIP-Codebood steps (smaller EL). However, it lags behind in the SPL metric. This discrepancy arises because SPL is evaluated based on the distance traveled rather than the actual number of steps taken. This highlights a limitation in this metric since the efficiency of path should consider factors like time and energy consumption. Actions such as rotations and look-ups/downs, which also consume time and energy, are not accounted for in this metric. In light of this observation, we further report Success Weighted by Episode Length (SEL): N Σee "ma][ where w; is the shortest possible episode length to the target object, and e; is the episode length produced by the agent. We utilize the privileged information from the environment to develop an expert agent to collect w; for each episode. Shown in Table I, CmmCCLIP-Codebow outperforms EmbCLIP by a significant margin in SEL.
outperforms in Success Rate and solves the task with fewer
4.1.2 OBJECT DISPLACEMENT
Model. We adopt the same model architecture (OLE w. GT mask) used in m- VOLE (Ehsani et al., 2022). The model consists of 3-layers CNN to encode the RGB-D and a ground truth segmentation mask, frozen R8 to encode the query images of target classes, including the source object (i.c., an apple) as the object has not picked up yet or the destination object after the source object (i.e., a table) has been picked up, and a distance embedder to encode the relative distance computed by the Object Location Estimator Module, which measures the distance between
"I Link to
leaderboard

---

Published as a conference paper at ICLR 2024
the end-effector and the source/destination object. The resulting visual features of current observation and query image are then concatenated with the distance feature to construct the observation embedding E with D 1536. While m- VOLE directly passes the embedding E' to the following recurrent state encoder, we add our codebook module with codebook's ' size of K = 256 and D. = 10 to compress the embedding E' to the goal-conditioned codebook embedding E. Finally, the recurrent state encoder processes the embedding E and the actor-critic after the recurrent state encoder further predicts the action probability distribution as well as the state value estimation. We call our model as m-
in the rest of the paper.
Experiment Details. We, again, use the AllenAct framework to train our m- Vole-Codebood with the default reward shaping provided by (Ehsani et al., 2022). Following (Ehsani et al., 2022), we train our agent by DD-PPO (Wijmans et al., 2019) with 80 samplers for 20M steps on APND dataset (Ehsani et al., 2021) I) in 30 kitchen scenes, where we split them into 20 training scenes, 5 validation scenes, and 5 testing scenes. During the training stage, the agent is required to navigate to the source object, pick it up, and bring it to the destination object. Please refer to Sec. A.1 for more training details.
Results. To quantify the models' performance, we report PickUp Success Rate (PU) and Success Rate (SR) as our evaluation metrics. As shown in Table I, we can find that our m- --VVV---ooebbooo outperforms m- VOLE baseline on both PU and SR. It demonstrates that our codebook module is applicable to the Embodied Al tasks involving interactions between the agent and environment.
4.2 CODEBOOK EMBEDDING IS EASIER TO TRANSFER TO NEW VISUAL DOMAINS
Model. This section investigates the capability of the codebook embedding to transfer across new vi- sual domains without exhaustive finetuning. Efficient transfer of Embodied-Al agents between distinct vi- sual domains remains a important challenge for many downstream applications. For example, there might be a desire to retain the high-level decision- making policy established in the recurrent state encoder and the actor-critic during large-scale pretraining, Alter- natively, resource constraints may prevent end-to-end finetuning of the entire model in the target domain. To address this, we propose finetuning only the CNN layers that follow the frozen CLIP ResNet backbone, the goal encoder, and the previous action encoder. This way ensures adjustments only to the visual rep- resentation and the task embedding, thereby adapting the agent to the target domain without changing the policy's 's behaviors completely. We call these fine- tuned components as the Adaptation Module, shown in Figure 4. Meanwhile, other components, including the CLIP ResNet, codebook, recurrent state encoder, and actor-critic, remain frozen.
Adaptation Module
Imp4 4144 Frame
Risheet
CLIP
CNN
CNN
8 E (O....11)
Goal Chair
Enbaaddr
Cadebook
Previuus Action
Embostter
Figure 4: Lightweight Finetuning of the Adap- tation Module. We only finetune a few CNN lay- ers, action/goal embedders, and the codebook scor- ing function when moving to new visual domain.
Experiment Details. We conduct lightweight finetuning of the Adaptation Module in both EmbCLIP- Codebook and the baseline for Object Navigation on Habitat 2022 ObjectNav challenge. Since AI2THOR and Habitat simulators exhibit variations in visual characteristics, lighting, textures and other environmental factors, they are appropriate platforms for studying the lightweight transfer capabilities of Embodied-Al agents. Following (Deitke et al., 2022b), we use the checkpoint at 200M steps to perform the finetuning in Habitat for another 200M steps with 40 samplers. Since the target object types in Habitat are subset of the ones in PROCTHOR, we discard the parameters in the goal embedder corresponding to the unavailable object types (more details in Supplementary Sec. A.
Results. Table 2 shows that EmbbCLIP-Codebood achieves superior performance during the lightweight fine-tuning of the Adaptation Module across all metrics. To elaborate, we attain 13.8% higher success rate and a 7.58% improvement in SPL, all while executing fewer invalid actions and achieving smoother trajectories. Our codebook bottleneck effectively decouples the process of learning salient visual information for the task from the process of decision-makin based on this filtered information. Consequently, when faced with similar task in for adaptation is significantly reduced. In this scenario, only the modules responsible for extracting essential visual cues in the new domain require fine-tuning, while the decision-making modules
visual domain, the need
7

---

Published as a conference paper at ICLR 2024
Goal:
Clock
View
Laptep
House
Broiceboll
Bowl
DibCLIP
(Cadebeok
Figure 5: GradCAM Attention Visualization. While EmbCLIP is distracted by different objects and other visual cues even though the target object is visible in the frame, EibCLLP-Coodeboo& able effectively ignore such distractions and only focus on the object goal.
Table 4: Our method with the codebook consistently outperforms the baselines using DINOv2 (Oquab et al., 2023) visual features in zero-shot evaluation 4 Object Goal Navigation benchmarks.
Table 3: Linear probing shows code- book embeddings encode more task-
relevant cues
Objectt novigation
Tusk
Model
Accurncy Fl Score"
Denchmark
Model
SRCS7
ELJ Curvauuee SPLT SELT
Obyect Presenze
EmbCLIP
0.10
0.55
(all 125 entegories)
#Codbbook
0.05
0.29
ProcTHOR 10k (validazion)
DINO-2 (Oqubb 2023)
-Codebook (Ours
74.25
129.00
0.34 49.53
Object Presence
FInCCIPP
goal categories)
CCoebboo
0.11
ARchITECTHiOR (0-shu))
DINO
-Cadhhook (Qurs)
59,75 194.00
$7.25 218.00
0.25 36.83 29.00
0,11 36,00 31.70
Goal Visibility
EmbCLIP
0.87
(binary)
+Codehook
0.92
A1--THRR (0-shor)
-Codebook (Ours)
DINOV2
76.93
74.67 97.00
68.00
0.19 59.45 26.50
0.07 60.14 28.30
Discance to Goal
EmbCLIP
0.29
(5-way clossification
4Codcbook
RoboTTIOR (0-shoo)
DINO<2
-Codebook (Ours)
29.36
can remain fixed. In contrast, EmbCLIP, which does not decouple this skill learning, necessitates adaptation of both visual cue extraction and decision-making modules when transitioning to new domain. This the primary reason for the substantial performance gap observed in the lightweight fine-tuning setting. These results are upper-hounded by end-to-end finetuning of the entire model with access to more training resources. Here, EmbCLIP-Codete performs nearly on par with the EmbCLIP while still maintaining smoother trajectories, as evidenced by the curvature metric.
4.3 CODEBOOK ENCODES ONLY THE MOST IMPORTANT INFORMATION TO THE TASK
We conduct an analysis (both qualitatively and quantitatively) to explore the information encapsulated within our bottlenecked representations after training for Object Navigation on PROCTHOR- 10k.
Linear Probing. We use linear probing to predict a list of primitive tasks from the goal-conditioned visual embedding E in EmbCLIP and the bottlenecked version F. in ManDCLiP-CCodbeee generated for 10k frames in ARCHITECTHOR scenes. The selected tasks include Object Presence (identifying which of the 125 object categories are present in the image), Goal Presence (identifying which the 16 goal categories are present in the image), Goal Visibility (determining if the object goal is visible in the frame and within 1.0m distance from the agent), and Distance to Goal (predicting the distance to the target object only when it's closer than 5m the agent). We train simple linear classifiers for each of these tasks to predict the desired outcome from the embeddings. The results are summarized in Table 3. Our bottleeeeeeee representations effectively exclude information related object categories other than the specified goal. Consequently, although our performance in predicting the presence of all object categories may be suboptimal, this information usually irrelevant to the object navigation task. Instead our method improves accuracy in predicting goal visibility. Furthermore, our representations better encode information about the distance to the goal (when it's close enough to the agent) resulting in better navigation towards the goal in the final stages of the task.
Grad-Cam Visualization. We utilize Grad-Cann (Selvaraju al., 2017) to visualize the attention map in the visual observation. Given an action predicted by the policy model at a single step, we treat it as a classification target and calculate the gradients to minimize the classification objective. With the gradients, further apply XGradCAM to visualize the final CNN layer producing coal-conditioned embeddings E. As shown in Figure 5, we observe that EmbCLIP often focuses on many different objects, including their appearance, texture, other distracting visual cues, though the target object has been visible in view. However, visual cues and concentrate on the target object.
is able to ignore many distracting
Nearest Neighbor Visualization To visually assess the information captured in our embeddings, we employ nearest neighbor retrieval using sample of 10k frames from Procthor, each accompanied
8

---

Published as a conference paper at ICLR 2024
(a)
Goal:
Rankerial
HoueePlont
Taleation
Red
(b)
Goall:
Cartageeaan
CortageCCon
Red
Oual Vise
E
Quel Tokk:
Goal:
Bool
Mig
Gool:
Goal:
Sufle
AlsmCllekk
Tdoeiaion
MoussParrt
(d)
Gral:
Apple
Anea
AlamnCoock
Visa
EMAKCLI -
EnSCLI *Coohiioo
Goal:
Vigg
Apple
Apple
14.g
Gool::
Apple
Apple
E16CLIP
Query and Goad
Goaa
Figure 6: Nearest-Neighbor Retrieval in the Coal-Conditioned Embedding Space (E and E), The 4 examples show that EmbCLIP-Coddeooo prioritizes task semanticy while EmbCLIP focusses on semantics. In (a) and (b), our nearest neighbors are based on goal visibility and goal proximity to the agent whercas EmbCLIP ones based the overall semantics the (tables in the left, toilets far away). In (c) and (d), ours favor the overall scene layout whereas EmbCL.IP mostly focuses on colors and appearances.
by a corresponding goal specification, which may may not be visible in the frame. We feed these frames and goals into our pretrained Objectnav models, generating their respective goal-conditioned embeddings (denoted as E EmbCLIP and É in neighbors in both the E and Ê embedding spaces, as illustrated in Figure (a) (d).
We then identify nearest
In (a), the codebook-bate servations with the goal prominently visible, including goals that are not Vase. Similarly in (b), the nearest neighbors are observations with the goal visible, but now at further distance from the agent, akin to the query image. In contrast, the EmbCLIP neighbors for (a) include wooden table the left, like the query image, and in (b) include free space with toilet in the distance. These examples indicate that LmbCLLP-CCodebbr tends focus on EmbCLIP -ww.l.o-oooookkkk neighbors show similarity in the overall scene layout, where as EmbCLIP neighbors tend prioritize the appearance of the floor and walls.
nearest neighbors for the Vase query image are other instances of ob-
prioritizes task semantics over EmbCLIP which semantics. In (c) and (d), the query does not contain the goal object. Here,
4.4 CODEBOOK MODULE IS TEPRESSENTATIOONAGGOSS
To assess the applicability of the codebook other visual encoders and to quantify the extent of the improvements irrespective of the underlying representation, we replaced the CLIP representations with DINOv2 (Oquab al., 2023) visual features. DINOv2 relies on discriminative self-supervised pre-training (which is different from the image-text pretraining approach utilized in CLIP) and is proven to be effective across a wide range of tasks, including Embodied-Al We use the frozen DINOv2 ViT-S/14 model to encode RGB images into 384 7 x 7 tensor which is compressed to tensor v using 2-layer CNN. Similar to EmbCLIP, v is concatenated with a 32-dimensional goal embedding and the previous action embedding and the result is flattened to obtain 1574-dimensional goal-conditioned observation embedding, denoted as We employed a codebook with similar dimensions, K 256 and D. 10, bottleneck the goal-conditioned representations. The results are presented in Table 4. Consistently, our approach outperforms the DINOv2 baseline models across variety of Object Navigation metrics in various benchmarks. This experiment underscores the effectiveness of our codebook module in bottlenecking other visual features for embodied- AI tasks.
CONCLUSION
Inspired by selective attention in humans, we proposed a compact learnable codebook module for Embodied-Al that decouples identifying the salient visual information useful for the task from the process of decision-making based on that filtered information. It acts as a task-conditioned bottleneck that filters out unnecessary information, allowing the agent to focus more task-related visual cues. It significantly outperforms state-of-the baselines in Object goal navigation and Object displacement tasks across 5 benchmarks. This results in significantly faster adaptation to new visual domains. Our qualitative and quantitative analyses show that captures more task-relevant information, resulting in more effective exploration strategies.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank members from RAIVN Lab the University of Washington and PRIOR team at Allen Institute for AI for valuable feedbacks on this project. This work is in part supported by NSF IIS 1652052, IIS 17303 166, DARPA N66001-19-2-4031, DARPA 1911515--10054 and gifts from Allen Institute for Al.
REFERENCES
Alexander A Alemi, Ian Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. In ICLR, 2017. 3
Peter Anderson, Qi Wu, Damien Teney, Jake Bruce, Mark Johnson, Niko Sünderhauf, Ian Reid, Stephen Gould, and Anton Van Den Hengel. Iision-and-langunee navigation: Interpreting visually- grounded navigation instructions in real environments. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 3674-3683, 2018.
Chenjia Bai, Lingxiao Wang, Lei Han, Animesh Garg, Jianye Hao, Peng Liu, and Zhaoran Wang. Dynamic bottleneck for robust self-supervised exploration. In NewrlPS, 2021.
Dhruv Batra, Angel X Chang, Sonia Chernova, Andrew Davison, Jia Deng, Vladlen Koltun, Sergey Levine, Jitendra Malik, Igor Mordatch, Roozbeh Mottaghi, et al. Rearrangement A challenge for embodied ai. ArXiv, 2020a.
Dhruv Batra, Aaron Gokaslan, Aniruddha Kembhavi, Oleksandr Maksymets, Roozbeh Mottaghi, Manolis Savva, Alexander Toshev, and Erik Wijmans. Objectnav revisited: On evaluation of embodied agents navigating to objects. arXiv preprint arXiv:2006.13177 13171, 2020b. 1,
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253 259, 2013.
Bergen R Bugelski and Delia A Alampay. The role of frequency in developing perceptual sets. Canadian Journal of Psychology/Revue canadienne psychologie, 15(4):205, 1961.
Timothy Buschman and Earl K Miller. Top-down versus bottom-up control of attention in the prefrontal and posterior parietal cortices. science, 315(0)) 860-1862, 2007. 2
Devendra Singh Chaplot, Dhiraj Gandhi, Saurabh Gupta, Abhinav Gupta, and Ruslan Salakhutdinov. Learning to explore using active neural slam. In ICLR, 2020a.
Devendra Singh Chaplot, Dhiraj Prakashchand Gandhi, Abhinav Gupta, and Russ R Salakhutdinov. Object goal navigation using goal-oriented semantic exploration. In NeurIPS, 2020b. 3
Yuxiao Chen, Jianbo Yuan, Yu Tian, Shijie Geng, Xinyu Li, Ding Zhou, Dimitris N Metaxas, and Hongxia Yang. Revisiting multimodal representation in contrastive learning: from patch and token embeddings to finite discrete tokens. In CVPR, 2023. 3
Maxime Chevalier- ShheallierBBoisvert, Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. In arXiv, volume abs/2306. 13831, 2023. 3
Matt Deitke, Winson Han, Alvaro Herrasti, Aniruddha Kembhavi, Eric Kolve, Roozbeh Mottaghi, Jordi Salvador, Dustin Schwenk, Eli VanderBilt, Matthew Wallingford, et al. Robothor: An open simulation-to-real embodied ai platform. CVPR, 2020. 2,5
Matt Deitke, Eli VanderBilt, Alvaro Herrasti, Luca Weihs, Kiana Ehsani, Jordi Salvador, Winson Han, Erie Kolve, Aniruddha Kembhavi, and Roozbeh Mottaghi. Procthor: Large-scale embodied ai using procedural generation. In NewrIPS, 2022a.
Matt Deitke, Eli VanderBilt, Alvaro Herrasti, Luca Weihs, Jordi Salvador, Kiana Ehsani, Winson Han, Eric Kolve, Ali Farhadi, Aniruddha Kembhavi, et al. Procthor: Large-scale embodied ai using procedural generation. arXiv preprint arXiv:2206.0699 2022b. 2, 7, 15, 16
10

---

Published as a conference paper at ICLR 2024
Robert Desimone and John Duncan. Neural mechanisms of selective visual attention. Annual review of neuroscience, (8(1):193-222 1995.
Heming Du, Xin Yu, and Liang Zheng. Learning object relation graph and tentative policy for visual navigation. ECCV, 2020.
Yilun Du, Chuang Gan, and Phillip Isola. Curious representation learning for embodied intelligence. In ICCV, 2021.
Jiafei Duan, Samson Yu, Hui Li Tan, Hongyuan Zhu, and Cheston Tan. A survey of embodied ai: From simulators to research tasks. IEEE Transactions on Emerging Topics Computational Intelligence, 6(2):230-244 2022.
Kiana Ehsani, Winson Han, Alvaro Herrasti, Eli VanderBilt, Luca Weihs, Eric Kolve, Aniruddha Kembhavi, and Roozbeh Mottaghi. Manipulathor: A framework for visual object manipulation. In CVPR, 2021. 1, 5, 15
Kiana Ehsani, Ali Farhadi, Aniruddha Kembhavi, and Roozbeh Mottaghi. Object manipulation via visual target localization. In ECCV, 2022. 2, 4, 6, 7, 15
Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In CVPR, 2021.
Jiameng Fan and Wenchao Li. Dribo: Robust deep reinforcement learning via multi-view information bottleneck. In ICML, 2022. 3
Marco Federici, Anjan Dutta, Patrick Forré, Nate Kushman, and Zeynep Akata. Learning robust representations via multi-view information bottleneck. In ICLR, 2020. 3
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. JMLR, 2022. 3
Samir Yitzhak Gadre, Kiana Ehsani, Shuran Song, and Roozbeh Mottaghi. representations for embodied ai. In CVPR, 2022. 3
Continuous scene
Daniel Gordon, Abhishek Kadian, Devi Parikh, Judy Hoffman, and Dhruv Batra. Splitnet: Sim2sim and task2task transfer for embodied visual navigation. In ICCV, 2019. 3
Anirudh Goyal, Riashat Islam, Daniel Strouse, Zafarali Ahmed, Matthew Botvinick, Hugo Larochelle, Yoshua Bengio, and Sergey Levine. Infobot: Transfer and exploration via the information bottle- neck. In /CLR, 2019. 3
Karol Gregor, Danilo Jimenez Rezende, Frederic Besse, Yan Wu, Hamza Merzic, and Aaron van den Oord. Shaping belief states with generative environment models for rl. In NeurlPS, 2019. 3
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. In arXiv preprint arXiv:1811.000 1.64407, 2018. 2
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, lean-Bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent-predictive representations for multitask reinforcement learning. In /CML, 2020. 3
Tengda Han, Weidi Xie, and Andrew Zisserman. Memory-augment dense predictive coding for video representation learning. ECCV, 2020.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, 2016.
Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang. Transforming auto-encoders. In ICANN, 2011. 3
Kyle Hsu, Will Dorrell, James CR Whittington, Jiajun Wu, and Chelsea Finn. Disentanglement via latent quantization. 2023.
11

---

Published as a conference paper at ICLR 2024
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization reinforcement learning with selective noise injection and information bottleneck. In NeurlPS, 2019.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In ICLR, 2017.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In ICLR, 2017.
Apoorv Khandelwal, Luca Weihs, Roozbeh Mottaghi, and Aniruddha Kembhavi. Simple but effective: Clip embeddings for embodied ai. In CVPR, pp. 14829, 2022. 1, 3,5
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv, 2014. 15, 16
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. In ICLR, 2014. 3
Pang Wei Koh, Thao Nguyen, Yew Siang Tang. Stephen Mussmann, Emma Pierson, Been Kim, and Percy Liang. Concept bottleneck models. In ICML, 2020. 3
Alexander Kolesnikov, André Susano Pinto, Lucas Beyer, Xiaohua Zhai, Jeremiah Harmsen, and Neil Houlsby. Uvim: A unified modeling approach for vision with learned guiding codes. In NeurIPS, 2022. 3
Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Daniel Gordon, Yuke Zhu, Abhinav Gupta, and Ali Farhadi. A12-THOR: An Interactive 3D Environment for Visual AI. arXiv, 2017a. 2
Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Daniel Gordon, Yuke Zhu, Abhinav Gupta, and Ali Farhadi. Ai2-thor: An interactive environment for visual ai. In Arxiv, 2017b. 5
Klemen Kotar, Aaron Walsman, and Roozbeh Mottaghi. Entl: Embodied navigation trajectory learner. In ICCV, 2023. 3
Jacob Krantz, Erik Wijmans, Arjun Majumdar, Dhruv Batra, and Stefan Lee. Beyond the nav-graph: aiiio-and-language navigation in continuous environments. In Computer Vision-ECCCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXVIII 16, pp. 104-120. Springer, 2020. 1
Jacob Krantz, Theophile Gervet, Karmesh Yadav, Austin Wang, Chris Paxton, Roozbeh Mottaghi, Dhruv Batra, Jitendra Malik, Stefan Lee, and Devendra Singh Chaplot. Navigating to objects specified by images. In ICCV, 2023.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In ICML, 2020.
Yoseph Linde, Andres Buzo, and Robert Gray. An algorithm for vector quantizer design. In IEEE Transactions communications, 1980. 5, 18
Xingyu Lu, Kimin Lee, Pieter Abbeel, and Stas Tiomkin. Dynamics generalization via information
bottleneck in deep reinforcement learning. In arXiv, 2020. 3
Rabeeh Karimi Mahabadi, Yonatan Belinkov, and James Henderson. Variational information bottle- neck for effective low-resource fine-tuning, In ICIR, 2021. 3
Julien Mairal, Jean Ponce, Guillermo Sapiro, Andrew Zisserman, and Francis Bach. Supervised dictionary learning. In NeurIPS, 2008.
Arjun Majumdar, Gunnar A Sigurdsson, Robinson Piramuthu, Jesse Thomason, Dhruv Batra, and Gaurav Sukhatme. Ssl enables learning from sparse rewards image-got goal navigation. In ICML, 2022. 3
12

---

Published as a conference paper at ICLR 2024
Pierre Marza, Lactitia Matignon, Olivier Simonin, Dhruv Batra, Christian Wolf, and Devendra Singh Chaplot. Autonerf: Training implicit scene representations with autonomous agents. In arXiv, 2023. 3
Lina Mezghan, Sainbayar Sukhbaatar, Thibaut Lavril, Oleksandr Maksymets, Dhruv Batra, Piotr Bojanowski, and Karteek Alahari. Memory-augmented reinforcement learning for image-goal navigation. IROS, 2022.
Piotr Mirowski, Razvan Pascanu, Fabio Viola, Hubert Soyer, Andrew Ballard, Andrea Banino, Misha Denil, Ross Goroshin, Laurent Sifre, Koray Kavukcuoglu, et al. Learning to navigate in complex environments. In ICLR, 2017. 3
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.07119, 2023. 5, 8, 9
Vincent Pacelli and Anirudha Majumdar. Learning task-driven control policies via information bottlenecks. In RSS, 2020.
Deepak Pathak, Pulkit Agrawal, Alexei Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In 2017.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In ICML, 2021. 3, 17
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In NeurIPS, 2021. 3, 5, 19
Nikolay Savinov, Alexey Dosovitskiy, and Vladlen Koltun. Semi-parametric topological memory for navigation. In ICLR, 2018. 3
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra. Habitat: A platform for embodied ai research. In ICCV, 2019. 2
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy
optimization algorithms. arXiv, 2017. 4, 15, 16
Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localiza- tion. In ICCV, 2017. 8
Julian Vlad Serban, Chinnadhurai Sankar, Michael Pieper. Joelle Pineau, and Yoshua Bengio. The bottleneck simulator: model-based deep reinforcement learning approach. Journal of Artificial Intelligence Research, 2020. 3
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gater mixture-of-expere layer. In ICLR, 2017. 5,
Mohit Shridhar, Jesse Thomason, Daniel Gordon, Yonatan Bisk, Winson Han, Roozbeh Mottaghi, Luke Zettlemoyer, and Dieter Fox. Alfred: A benchmark for interpreting grounded instructions for everyday tasks. In CVPR, 2020. 1
Daniel Simons and Christopher Chabris. Gorillas in our midst: Sustained inattentional blindness for dynamic events. perception, 18(9):105--100 1074, 1999. 2
Kunal Pratap Singh, Jordi Salvador, Luca Weihs, and Aniruddha Kembhavi. A general purpose supervisory signal for embodied agents. ICCV, 2023. 2, 18, 19
Aravind Srinivas, Tsung Yi Lin, Niki Parmar, Jonathon Shlens, Pieter Abbeel, and Ashish Vaswani. Bottleneck transformers for visual recognition. In CVPR, 2021. 3
13

---

Published as a conference paper at ICLR 2024
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. In IMLR, 2014. 5
Anne M Treisman and Garry Gelade. A feature-integration theory of attention. Cognitive psychology, 12(1):97-136, 1980. 2
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. In NeurlPS, 2017. 3,5
Matthew Wallingford, Aditya Kusupati, Alex Fang. Vivek Ramanujan, Aniruddha Kembhavi, Roozbeh Mottaghi, and Ali Farhadi. Neural radiance field codebooks. In ICLR, 2023, 3, 18, 19
Luca Weihs, Matt Deitke, Aniruddha Kembhavi, and Roozbeh Mottaghi. Visual room rearrangement. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp.
1922-5931, 2021a.
Luca Weihs, Aniruddha Kembhavi, Kiana Ehsani, Sarah M Pratt, Winson Han, Alvaro Herrasti, Eric Kolve, Dustin Schwenk, Roozbeh Mottaghi, and Farhadi. Learning generalizable visual representations via interactive gameplay. In ICLR, 2021b. 3
Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near- cear-perfect pointgoal navigators from 2.5 billion frames. In ICLR, 2019. 15
Jeremy M Wolfe. Guided search 2.0 revised model of visual search. Psychonomic bulletin & review, 1:202-238, 1994. 2
Yi Wu, Yuxin Wu, Aviv Tamar, Stuart Russell, Georgia Gkioxari, and Yuandong Tian. Bayesian relational memory for semantic visual navigation. In ICCV, 2019.
Fanbo Xiang, Yuzhe Qin, Kaichun Mo, Yikuan Xia, Hao Zhu, Fangchen Liu, Minghua Liu, Hanxiao Jiang, Yifu Yuan, He Wang, al. Sapien: simulated part-based interactive environment. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11097-11107, 2020. 1
Karmesh Yadav, Arjun Majumdar, Ram Ramrakhya, Naoki Yokoyama, Alexei Baevski, Zsolt Kira. Oleksandr Maksymets, and Dhruv Batra. Ovrl-v2: A simple statc-of-art baseline for imagenav and objectnav. In arviv, 2023a.
Karmesh Yadav, Ram Ramrakhya, Arjun Majumdar, Vincent-Pierre Berges, Sachit Kuhar, Dhruv Batra, Alexei Baevski, and Oleksandr Maksymets. Offline visual representation learning for embodied navigation. In ICLR Workshop on Reincaraating Reinforcement Learning, 2023b. 1, 3
Joel Ye, Dhruv Batra, Abhishek Das, and Erik Wijmans. Auxiliary tasks and exploration enable objectgoal navigation. In ICCV, 2021a.
Joel Ye, Dhruv Batra, Erik Wijmans, and Abhishek Das. Auxiliary tasks speed up learning point goal navigation. In CoRL, 2021b. 3
Kuo-Hao Zeng, Luca Weihs, Ali Farhadi, and Roozbeh Mottaghi. Pushing it out of the way: Interactive visual navigation. CVPR, 2021.
Kuo-Hao Zeng Luca Weihs, Roozbeh Mottaghi, and Ali Farhadi. Moving forward by moving backward: Embedding action impact over action semanties. In ICLR, 2023.
14

---

Published as a conference paper at ICLR 2024
A
APPENDIX
The following items are provided in the Appendix:
Training details and code along with our pretrained models for reproducing our experiments
($ A.1).
II. A qualitative analysis of the information captured in our codebook module using nearest neighbor retrieval ($ A.2).
III. A comparison of our codehook-based architecture with other informatiion-oottleeckke baselines (8 A.3).
IV. Qualitative examples of the failure cases of our approach (S A.4).
V. An analysis of the impact of end-to-end finetuning of the CLIP visual encoder on final performance ($ A.5).
VI. An ablation study on codebook regularization and codebook collapse ($ A.6).
VII. A study of our agent's behavior compared the baseline in terms of trajectory smoothness, redundant actions, etc. (S A.7).
VIII. A comparison with other object goal navigation baselines ($ A.8).
IX. An ablation study of codebook hyperparameters including codebook size, scoring mechanism applied the codebook entries, etc. ($ A.9).
A.1 TRAINING DETAILS
Code and pretrained models are publicly available through the project page.
Object Goal Navigation
We follow (Deitke et al., 2022b) to pretrain the EmbCLIP baseline and EmbCLIP-Codebood on the the PROCTHOR-10k houses using the default action space (MoveAhead, RotateRigh RotateLeft, LockUp, LookDown, Done]. During this stage, we optmize the models by Adam (Kingma & Ba, 2014) optimizer with a fixed learning rate of 3e In addition, we fol- low (Deitke al., 2022b) to have two warm up stages by training model parameters with lower number of steps per batch for PPO training (Schulman et al., 2017). In the first stage, we set number of steps as 32, and in the second stage, we increase the number of steps 64. These two stages are trained by 1M steps, respectively. After the second stage, we increase the number of steps to 128 and keep it till the end of training (e.g., 435M steps). For reward shaping. we employ the default reward shaping defined in AlleAct: Reenalty + Ryuceess + Rdistances where Reenalty 00.11 Rsuccess = 10, and Rdistance denotes the cahnge of distances from target between two consecutive steps. In the experiments, we follow (Deitke al., 2022b) use 16 objects possible target objects, include
Apple, Baseba 1 1 Bat, Basket Ball, Bed, Bowl, Chair, Ga GarbageCan, HousePlams antt Laptop, Mug, Sofa, SprayBott le, relevision, Toi let, Vase), J. Figure 7 compares the Success Rate training curves for EmbCLIP and EmbCLIP-Codebook. As shown in the figure, adding our codebook bottleneck significantly improves the sample efficiency through faster convergence, makes the training more robust, and achieves better performance by the end of the training.
Object Displacement
Following (Ehsani et al., 2022), we optimize our agent's model parameters by Adam opti- mizer (Kingma & Ba, 2014) with a fixed learning of 34 and number of steps of 128 for DD-PPO (Wi- jmans et al., 2019) training. On APND dataset (Ehsani et al., 2021) we use 12 possible target objects, including (Apple, Bread, Tomato, Lettuce, Pot, Mug, P Potato, P Pan, Egg, Spatula, Cup, and SoapBot to During the training stage, we sample an initial location from 130 possible agent initial locations to spawn the agent and sample initial locations for source and target objects from 400 possible object-pair locations to place objects. The action space in this environment is (MoveAhead, RotateRi ght, Rotateleft, MoveArmBaseUp, MoveArmBaseDown, MoveArmåhead, MoveArmBack, MoveArmRight ght, MoveArmLeft. MoveArmÜp, MoveArmDown)
Habitat Finetuning
15

---

Published as a conference paper at ICLR 2024
0.8
Success Rate
0.8
0.4
02
EmbCLIP Codebbokk
EmbCLIP
100M 200M
300M
400M
Training Step
Figure 7: Success Rate training curve comparing EmbCLIP and EmbCLIP-Codebook
Table 5: Our proposed codebook module consistently outperforms other information bottleneck baselines across different ObjectNav benchmarks. EmbCLIP-AF AE uses an auto-encoder to bottleneck the task-conditioned embedding E and EmbCLIP-SOn -eeffftttttt applies self attention top of the goal-conditioned CLIP feature maps. (Models evaluated 210M training steps.)
Object navigation
Model
SR(S6)1
EL.]
SPL.7
EmbCLIP
133.00
45.98
EmbCLIP-AE
145.00
42.14
Pro:THOR-100 (yalidation)
EmbCLIP-SelfANe
$6.50
126.00
44.02
MubCLIP-Codcbbot (Ours)
72.49
134,00
49,37
EmbCLIP
$2.67
172.00
34,93
EmbCLIP-AE
48.67
166.00
32.58
ArchitecTHOR (0-shol)
EmbCLIP-SellAAtn
50.00
151.00
36.32
KmbCCIJP-Codebol (Ours)
55.25
160.00
35.31
EmbCLIP
64.13
108.00
52.76
EmbCLIP-AE
58.93
96.00
47.29
AI2-TTHOR (0-shot)
EmbCLLP-SelFArt
EmbCLIP-Codeben (Ours)
74.80
81,00
$7.79
As mentioned in the main paper (Sec. 4.2), follow (Deitke al., 2022b) 200M steps to perform another 200M finetuning. However, in our lightweight finetuning, we only update the model parameters inside the adaption module (shown in Fig. 4). In this finetuning stage, we use Adam (Kingma & Ba, 2014) optimizer to update the model parameters and we set the number of steps in the PPO (Schulman et al., 2017) training as 128. We try 4 learning rates, incluing (1e-4. 2e 2e-1, 3,, 4e and find 14 works best for EmbCLIP and -1 works best for the model with our codebook module. Furthermore, as mentioned in the Sec. 4.2, because the target object types in Habitat, including (Bed, Chair, House
the checkpoint
Sofa, Television. Toilet). ), are subset of the ones in PROCTHOR, we remove the parameters the goal embedder corresponding to the rest of 10 objects presented in the pretraining stage.
A.2 WHAT'S ENCODED
CODEBOOK MODULE?
We expand our analysis of nearest neighbor retrieval to further evaluate the information captured within our codebook module. Utilizing set of 10k frames from ProcTHOR, each accompanied by corresponding goal specification that may or may not visible in the frame, we input these frames and their associated goals into our pretrained LmbCLIP-Codebbod generates hidden compact representations h, which are convex combinations learnable codes {ci)/K] To visually assess the information encoded in each individual learnable code e;, perform nearest neighbor retrieval on the dataset using each code c; as query. Fig. illustrates the 8 nearest neighbors for 9 different codes. The figure reveals that certain codes primarily contribute goal visibility, determined by their proximity to the agent, while others capture various geometric and semantic details of scene layouts, such walkable areas, corners, doors, etc.
model. This process
16

---

Published as a conference paper at ICLR 2024
Table 6: We fine-tune CLIP ResNet visual encoder an end-to-end manner for an additional 30M steps. The table shows the results for the frozen backbone and after end-to-end finetuning. EmbCLIP-Collebok consistently outperforms EmbCLIP both with frozen and fine-tuned visual encoder.
Object navigation
Benchmark
Model
SR(%5)
E1.]
SPL.1
EmbCCIP
67.70
182.00
49,00
ProcTHOR- 1ok (validation)
CumbCCIP-Coddbbo
73.72
136,00
48.37
EmbCLIP CLIP ResNet tine-tuning
74.37
[38.00
52.46
EmbCLIP CmncCLiP-Codebook CLIP ResNer fine runing
79.38
120.00
52.67
EmbCLPP
55 80
222.00
38.30
ARCHITECTHOR (0-shot)
CondCLi--Coldbook
58.33
174.00
35,57
EmbCLIP + CLIP ResNet line-tuning
59.00
182.00
38.92
EmbCLIP Coddei-k CLIP ResNet fine tuning
62.38
EmbCLIP
70.00
121.00
57.10
A12-ITHOR (0-shot)
MinCCLIP-Codebod
78.40
86,00
54.39
EmbCLIP + CLIP ResNet tre-nnnnng
72.53
58.00
59.28
EmbCI IP-Codebook CLIP ResNer fine-runing
77.75
99.00
A.3 COMPARISON WITH OTHER INFORMATION- INORMATION-BOTTLEEEEKEE BOTTLENECKED BASELINES
We ablate the choice of the bottleneck architecture and conduct a comparison between our code- book module and two alternative information- against EmbCLIP-AE, E, which employs an auto-encoder on the goal-conditioned embedding E. This autoencoder comprises series of linear layers with ReLU activation functions, mapping the representation P R256 and h R before resizing back É R Additionally, we as another baseline, utilizing celf-attention on the CLIP feature maps. To achieve this, we merge the compressed CLIP feature map o with the goal embedding g, resulting in 32 X 7 x 7 goal-conditioned feature map. This map is then processed through the self-attention module, where Ixl convolutions serve as the k, q, and e functions. Due to limited computational resources and time constraints during the discussion period, the results are reported for the Object Navigation task at 210M training steps as shown in Table 5.
baselines. Specifically, evaluate
introduce EmbCLIP-SelfAL
A.4 WHAT ARE THE FAILURE CASES FOR EMBCLIP-CODEBOOKS
Fig. 9 illustrates examples of two modes of failure in our 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. The top row of Fig. 9 showcases examples of failures related the 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 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.
agent: perception and
A.5 END-TO-END FINE-TUNING OF THE VISUAL ENCODER
In this section, we investigate the impact of end-to-end finetuning the entire policy model, including the CLIP ResNet visual encoder. Initially, trained the models with the visual backbone frozen for 420 million steps, as depicted in Figure 2, with the results in Table 1. Subsequently, we fine-tuned the entire policy model. including the visual encoder, end-to-end for an additional 30 million steps. This process of fine-tuning the entire policy model incurs high computational costs. To facilitate this, we utilize multi-node training on 4 servers, each equipped with eight A-100-80GB GPUs. Moreover, we found that it is important to (1) employ larger number of samplers (128) to match the substantial batch size (32, 768) used in the CLIP's pretraining (Radford et al., 202 1),, and (2) use a small learning rate (1e-6) for training stability. As shown in Table 6, while there is a substantial improvement from fine-tuning the entire policy model, including the visual encoder, compared to the frozen one, noticeable gap still exists between fine-tuned EmbCLIP and fine-tuned Remarkably, with further end-to-end finetuning, EmbCLIP-Codes results compared to all other models.
CoLLI-CCddbbbok. achieves significantly superior
17

---

Published as a conference paper at ICLR 2024
Table 7: Adding a dropout on top the probability distribution over the latent codes acts as a codebook regularization and improves the performance. (Models are evaluated at 150M training steps.)
Object navigation
Benchmark
Model
SR(K)
SPL.
EmbCLIP
59.35
115.00
45.27
ProcTHOR-10K (validation)
EmhCLIP Codebook Dropout
126.00
EmbCLIP-Codeboot
71.39
147.00
47.46
EmbCLIP
45.67
127.00
33,68
ARCHITECTHOR (0-shor)
EmbC1.PP -Pc.e-CCo w/o Dropout
48.08
164.00
33.61
EmbCLIP-Codeie
57.17
192.00
35.65
EmhCLIPP
AI2-TIIO shot)
EnbCLIP Codebook w/o Dropout
66.93
102.00
54,13
EmhCIIIP PoCoeeokkk
73.47
108.00
54.40
A.6 CODEBOOK REGULARIZZTION
Codebook collapse is prevalent issue marked by the assignment of overconfident high probabilities to specific subset of latent codes in the codebook, causing the majority of codes to be underutilized. This restrains the models from fully leveraging the codebook's capacity. While it's important to note that our codebook module doesn't 't experience complete collapse, incorporating regularization into the codebook enhances overall performance by promoting a more balanced utilization over all the latent codes. While various solutions have been suggested, we have found 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 training curves of the two models and Table 7 presents the corresponding evaluation results. Clearly, introducing such regularization to the codebook results in improved performance. We additionally tried Linde-Buzo- Buzo-Spliitingg algorithm (Linde et al., 1980) to balance the codebook useage. More specifically, during the training process, if there is 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 depicted in Fig. 11, which shows severe imbalance and eventual collapse during the training.
more uniform average usage of the latent codes. Fig. 11 compares the
A.7 COMPARISON OF AGENT'S BEHAVIOR
To enhance our understanding of the agent's 's navigation behavior, we delved deeper than the conven- tional embodied navigation metries. We assessed the sequential actions and curvature values, offering empirical insights into the agent's 's behavior throughout the navigation episodes. Fig. 12 shows the action distribution of EmbCL.IP-Codebor compared to EmbCLIP baseline ARCHITECTHOR scenes. As shown in the plot, MabbCLIP-Codebon mostly favors MoveAhead actions while EmbCLIP baseline performs lots redundant rotations. This observation also supported by Fig. 13 which plots the frequencies of sequential MoveAhead actions taken by the agents. takes sequential MoveAhead actions much more frequently than the baseline resulting in smoother trajectories travelled by the agent. This provides further evidence for the lower average curvature for the EmbCLIP IP--Codebook trajectories as reported in Table 1. This is qualitatively visualized in the top-down maps shown in Fig. 14.
A.8 COMPARISON TO OTHER BASELINES
Table 8 compares our method with 2 other baselines SGC (Singh et al., 2023) and NRC (Wallingford et al., 2023) on 3 Object Navigation benchmarks based on the numbers reported in the papers. Scene Graph Contrastive (SGC) improves the representation by building scene graph from the gent's observations. Neural Radiance Field Codebooks (NRC) learns object-centric representations through novel view reconstruction and finetunes the representation for Object Navigation task. As shown in the table, we achieve better success rate and lower episode length compared to all methods across all 3 benchmarks.
18

---

Published as a conference paper at ICLR 2024
Table 8: Comparisons against other methods SGC (Singh et al., 2023) NRC (Wallingford et al., 2023) on Object Navigation Benchmarks. Our method outperforms all the baselines on both Success Rate and Episode Length in all 3 benchmarks.
Object navigation
Benchmark
Model
SR((5)T
EL]
SPLT
SGC (Singh et al., 2023)
51.80
204.00
34,80
ARCHITECTHOR (0-shol)
EmbCLIP
222.00
Conlebo
58.33
174.00
35.57
NRC (Wallingford et al,, 2023) (finctuned)
50.10
23.90
RoboTHOR-TE (0-sbot)
EmbCLIP (0-shat)
51.32
24.24
EmbCLIP+Codeble (0-shot)
SGC (Singh 2023)
71.40
124.00
59.30
A12-THOR (0-shot)
EmbCLIP
70.00
121.00
57.10
EmbCLIPPCodebook
78.40
86.00
54.39
Table 9: Ablations Different Codebook Hyper- Parameters Object Navigation Benchmarks.
Object Navigation
Renchmark
Codehook Indexing
Codehook Size
CoallCCnddtiooeed
SR( %))
E1.1
SPL.
1024 10
72.56
119
48.88
Softmax
73.4
49.84
256
43.62
ProcTHOR-1Ok (validation)
Top-1 gating
38.06
110
29.39
Top-8 gating
256 10
68.22
160
45.86
Top 32 gating
73.27
138
48.5
EmbC1.IP Baseline
59.81
176
43.91
1024 10
60.00
160
37.33
Softmax
34.49
48.42
31.02
ARCHITECTHOR (0-sbot)
Top-1 gating
25.83
112
20.00
Top-8 gating
256 10
52.50
222
32.07
Top 32 gating
59.33
188
36.14
EmbC1.IP Baseline
50.5
209
35.64
1024 x 10
76.4
65
57.4
Softmax
72.67
54.82
48.53
A12-TTHOR (0-shot)
Top-1 gating
46.13
100
37.5
Top-8 gating
256
73.2
129
$1.23
Top 32 gating
72.93
101
50.57
EmbCLIP Baseline
60.80
137
$2.35
A.9 CODEBOOK ABLATIONS
Table 9 compares different codebook hyperparameters including codebook size, scoring mechanism applied to the codebook entries and the choice of representation that undergoes compression by the codebook cooal-oondiioone training steps and evaluated vero-shot on ProcTHOR (validation), Architec THOR and A12-iTHOR benchmarks. We ablate different codebook sizes specifically 256 and 1024 x 10. To underscore the significance of goal-conditioning we employ the codebook to compress representations both before and after they are fused with the goal (u or F. embeddings). codebook bottleneck is only effective when the visual representation is conditioned the goal across all the benchmarks. Moreover, to further intensify the bottlenecking mechanism, we selectively retain only the top-N scores, disregarding the rest, before computing the convex combination of the codebook entries. This inspired by the gating mechanism in Mixture Experts (Riquelme al., 2021; Shazeer et al., 2017) and results in only the top N most important codes to contribute to the bottlnecked representation at any given time. As shown in the table, selecting the top 32 codes for the convex combination can sometimes work as well as using all 256 latent codes. Using a bigger codebook size (1024 instead of 256) improves the zero-shot generalization in some of the benchmarks.
not). All models are trained on PROCTHOR-
houses for 300M
shown in the table, our
19

---

Published as a conference paper at ICLR 2024
Goal::
Toilet
GarbageCan
Television
Toilet
Chair
GarbageCan
BaseballBat
BasketBal
Code 1
Goal::
HousePiant
Bowl
Chair
HousePlant
BasketBall
Apple
AlarmClock
HousePlant
Code 2
AlarmClock
HousePlant
Code 3
Goal:
HousePlant
SprayBottle
Bed
HousePlant
Television
HousePlant
Apple
SprayBottte
Code 4
Goal:
Toilet
SprayBottie
Sofa
Code 5
Apple
SprayBottle
Apple
Mug
Apple
AlammClock
Code 6
Goal:
HousePlant
HousePlant
BasketBall
Tollet
GarbageCan
Bed
Code 7
Goal:
Toilet
SprayBottle
Chair
Toilet
Toilet
GerbageCen
GarbageCan
Toilet
Code 8
Goal:
HousePlant
AlarmClock
Mug
GarbageCan
Mug
Laptop
Mug
Toilet
Code 9
Figure 8: Nearest-neighhor retrieval the hidden compact Representation Space e R00 using {c [sl queries. We show the nearest neighbors of learnable codes in the codebook module. While some codes seem to be responsible for encoding the object goal visibility depending on different proximities the agent (codes to 5), others encode other semantic and geometric information about the scene layout such as walkable surfaces (code 6,7), corners (code 8), and doors (code etc.
20

---

Published as a conference paper at ICLR 2024
Starting Pose
Fallure Case (Perception)
Goal
Fallure Case 2 (Exploration)
Figure 9: Failure cases of EmbCLLIP-Cudebook in Object Navigation. Top row shows the failure examples related perception, where the agent traverses past the goal object but fails identify it. The bottom row shows another failure mode where the agent fails to explore specific areas of the environment where the target object is located. (Agent's starting pose is shown white circle and the target object is shown in a red circle)
Codebook wio Dropout
Codebook with Dropout 0.1
0.08
0.025
Success Rele
0.000
0.015
0.04
0.6
0.010
0.2
Croccu--oaaccc ai Depautt
0.02
R155
Training Step
Figure 10: Average probability distribution P E R 206 over the latent codes comparing the codebook trained with 0.1 dropout probability and without any dropout. Using dropout clearly results in utilization of all the latent codes in codebook module. (P is reshaped 16 x for visualization.)
Figure 11: Success Rate training comparing EmbCLIP (baseline), EmbCLIP Codebook dropout regularization, and EmbCLIP- Codebook trained with Linde-Buu Buzz- Gray splitting algorithm.
uniform
CmmbLLI-Coodbbok with and without
MehCLiP-Coddeono
CmbCLLP-Cooebook
0.4
EmbCLIP
EmbCLIP
Frequency
Frequency
10a
0.2
10c
10
MoveAhead RotatoLoft RotateRight
End
LookUp LookDown
10
Actions
Sequential MoveAhead Actions
Figure 12: Action frequenciess in ArchitecTHOR validation trajectories.
Figure 13: Frequencies of sequential MoveAhead
actions in ArchiteTTHOR validation trajectories.
21

---

Published as a conference paper at ICLR 2024
EmbCLiP-Codebook
EmbCLIP
Success V
Success X
Success
Success X
Success
Success X
Success
V
Success
V
Figure 14: Top-down Maps of Agents' Trajectories. Our agent travels in smoother paths and explores the environment more effectively.
22


Reviews:
[{"review_id": "ybUMTA0odN", "reviewer": "Reviewer_X3Z2", "summary": "This paper proposes an adaptive feature modulation module for embodied AI.  The core idea is to create a task-conditioned representation bottleneck, based on a learnable codebook, to select useful features for each robotic task.  The codebook is jointly trained to optimize the task reward.  This paper demonstrates superior efficacy of the proposed method, and conduct thorough ablation study / visualization of the proposed method.", "questions": "1. 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.\n\n2. The nearest neighbor probably should be done by treating learned codes as query.  In stead of using the pooled/weighted sum features, the authors can use a one-hot vector to select each code and upsample it to 1568 dimension.  It'd be more interesting to see what are learned in the codebook and what do the retrievals look like for each task.\n\n3. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Good ablation study.  Table 3 clearly shows the selected information of the codebook. Grad-CAM visualization also makes sense.\n\n2. The paper is well written.  It's easy to follow especially that readers are primed with some high-level ideas from human vision.\n\n3. The performance gain is significant compared to the baseline.", "weaknesses": "1. 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 toe the novel task faster and better than the proposed codebook method.\n\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w5PLgEg5nQ", "reviewer": "Reviewer_YjyH", "summary": "The paper presents a parameter efficient approach to filter out task-irrelevant information encoded by visual encoders, such as CLIP, in embodied AI tasks. This approach leverages a compact, learnable codebook module to establish a task-specific filter for visual observations. The codebook is trained to optimize task performance, serving as a filter that directs the agent's attention toward task-relevant visual cues. The experimental results demonstrate performance improvement in object goal navigation and object displacement tasks across various benchmarks. Qualitative analysis illustrates that agents become more proficient in exploring their surroundings, retaining task-relevant information, and disregarding irrelevant visual details.", "questions": "I have listed my main concerns in the weaknesses section. I will be happy to increase the score if the authors can answer those with relevant experiments. Other than that, I have some other questions and suggestions that I list below:\n1) What is meant by “Samplers”, which is mentioned in the experimental details?\n2) Why is RobotTHOR missing some metrics in Table 1?\nTypos:\nPage 6: HMD semantics -> HM3D Semantics. Also missing citation for HM3D Semantics\t\nPage 5: temperture -> temperature\nPage 2: codebook better encodes -> codebook encodes better\nPage 17: gent’s -> agent’s", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1) The paper proposes a simple parameter efficient approach for adapting representations which leads to consistent performance improvement on multiple benchmarks. \n2) Because the model learns to ignore the irrelevant parts of the visual inputs, it is able to do more efficient exploration and navigation.. \n3) The proposed finetuning approach of the codebook module is clever.", "weaknesses": "1) 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. \n2) 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. \n3) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "75pE5eUMGb", "reviewer": "Reviewer_vxKH", "summary": "Visual encoders like CLIP capture general purpose scene information which includes details not relevant to the task. The paper proposes to learn a codebook module that selectively filters information from visual representations specific for a task. They demonstrate the approach on a wide range of Embodied AI tasks across several large-scale benchmarks. Additionally, they present qualitative analysis showing that the codebook module better encodes task information which lends to more robust navigational policies.", "questions": "Could the codebook be learning to localize the goal information in the scene? The input to the codebook contains both the visual observation and a language description of the goal. It could be possible that the codebook is simply acting as a bridge between the two input modalities.\n\nWhat are some of the failure cases of the EmbCLIP + codebook method? Are there instances where codebook fails to identify task objects?\n\nAre there cases where two objects of the same class are present in the scene, but the task requires the model to pick one of the objects.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "Approach is simple, easy to follow, and well motivated. \n\nCodebook representations and policy are learned end-to-end.\n\nResults are shown across a diverse set of tasks and benchmarks. Improvement over EmbCLIP seems significant. Also report a comprehensive set of metrics including a new metric Success Weighted by Episode Length which accounts for actions like changing viewing angles that also requires time and effort. \n\nInteresting experiments showing that with minimal finetuning of the Adaptation Module with a frozen codebook can help transfer to new visual domains.\n\nSaliency visualizations are informative of where the model is paying attention to and highlights relevant tasks objects.\n\nNice set of ablation studies on codebook latent size and linear probes.", "weaknesses": "Does the codebook work for other types of general purpose visual encoders? This work seems to only show the codebook module applied to EmbCLIP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DMmnUJLRUD", "reviewer": "Reviewer_NmEd", "summary": "This work starts from a very attractive motivation: redundant information will blind people to making correct actions. The authors utilize a simple yet effective trick, i.e., adding a parameter-efficient codebook, and expecting it to filter out the unnecessary information in the embeddings. The method significantly boosts the performance of the baseline EmbCLIP on several benchmarks.", "questions": "Should adding a skip connection between \\hat{E} and E help improve the current version, i.e., \\hat{E} = theta(E)+Code Module(E)? I think bottleneck representations such as ResNet etc will usually do so.\n\nFor the other questions please see weakness.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I am quite appealed by the introduction, which I believe illustrates a very interesting motivation and will inspire the community. The method is also very simple and effective. Simply incorporating a codebook module is able to improve the baseline EmbCLIP by a large margin on several benchmarks. Overall, I believe this paper makes a good contribution, especially in terms of motivations and methods.", "weaknesses": "One of the weaknesses is also related to its strengths, i.e., motivation. I was quite attracted by the example of Figure 1, in which the authors describe the situation that the keys would only lie on a flattened surface, not a corner or somewhere else. I do like the motivation and expect the authors to incorporate such a human prior into the method yet I did not find it in the paper. I would suggest the authors consider incorporating a text prompt on the CLIP text model such as \"the chair is usually on the floor\" etc. Not necessarily for all the benchmarks but I think it would be interesting to see whether several cases can be improved.\n\nThe second weakness is about codebook collapse, this is a good motivation and interesting problem that is also related to interpreting what codebook really learns. I would expect whether the current method exists codebook collapse without dropout during training. This means the authors may need to conduct ablation studies and provide illustrations about collapse or not.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)