SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wonje Choi, Woo Kyung Kim, SeungHyun Kim, Honguk Woo*
{wjchoi1995, kwk2696, kimsh571, hwoo}@skku.edu
For embodied reinforcement learning (RL) agents interacting with the environment, it is desirable to have rapid policy adaptation to unseen visual observations, but achieving zero-shot adaptation capability is considered as a challenging problem in the RL context. To address the problem, we present a novel contrastive prompt ensemble (CONPE) framework which utilizes a pretrained vision-language model and a set of visual prompts, thus enabling efficient policy learning and adaptation upon a wide range of environmental and physical changes encountered by embodied agents. Specifically, we devise a guided-attention-based ensemble approach with multiple visual prompts on the vision-language model to construct robust state representations. Each prompt is contrastively learned in terms of an individual domain factor that significantly affects the agent's egocentric perception and observation. For a given task, the attention-based ensemble and policy are jointly learned so that the resulting state representations not only generalize to various domains but are also optimized for learning the task. Through experiments, we show that CONPE outperforms other state-of-the-art algorithms for several embodied agent tasks including navigation in AI2THOR, manipulation in egocentric-Metaworld, and autonomous driving in CARLA, while also improving the sample efficiency of policy learning and adaptation.
In the literature of vision-based reinforcement learning (RL), with the advance of unsupervised techniques and large-scale pretrained models for computer vision, the decoupled structure, in which visual encoders are separately trained and used later for policy learning, has gained popularity [1, 2, 3]. This decoupling demonstrates high efficiency in low data regimes with sparse reward signals, compared to end-to-end RL. In this regard, several works on adopting the decoupled structure to embodied agents interacting with the environment were introduced [4, 5], and specifically, pretrained vision models (e.g., ResNet in [6]) or vision-language models (e.g., CLIP in [7, 8]) were exploited for visual state representation encoders. Yet, it is non-trivial to achieve zero-shot adaptation to visual domain changes in the environment with high diversity and non-stationarity, which are inherent for embodied agents. It was rarely investigated how to optimize those popular large-scale pretrained models to ensure the zero-shot capability of embodied agents.
Embodied agents have several environmental and physical properties, such as egocentric camera position, stride length, and illumination, which are domain factors making significant changes in agents' perception and observation. In the target (deployment) environment with uncalibrated settings on those domain factors, RL policies relying on pretrained visual encoders remain vulnerable to domain changes.
*Honguk Woo is the corresponding author.
Figure 1 provides an example of egocentric visual domain changes experienced by embodied agents due to different camera positions. When policies learned in the source environment are applied to the target environment, zero-shot performance can be significantly degraded, unless the visual encoder could adapt not only to environmental differences but also to the physical diversity
of agents. In this paper, we investigate RL policy adaptation techniques for embodied agents to enable zero-shot adaptation to domain changes, by leveraging prompt-based learning for pretrained models in the decoupled RL structure. To this end, we present CONPE, a novel contrastive prompt ensemble framework that uses the CLIP vision-language model as the visual encoder, and facilitates dynamic adjustments of visual state representations against domain changes through an ensemble of contrastively learned visual prompts. In CONPE, the ensemble employs attention-based state composition on multiple visual embeddings from the same input observation, where each embedding corresponds to a state representation individually prompted for a specific domain factor. Specifically, the cosine similarity between an input observation and its respective prompted embeddings is used to calculate attention weights effectively.
Through experiments, we demonstrate the benefits of our approach. First, RL policies learned via CONPE achieve competitive zero-shot performance upon a wide variety of egocentric visual domain variations for several embodied agent tasks, such as navigation tasks in AI2THOR [9], vision-based robot manipulation tasks in egocentric-Metaworld, and autonomous driving tasks in CARLA [10]. For instance, the policy via CONPE outperforms EmbCLIP [7] in zero-shot performance by  $20.7\%$  for unseen target domains in the AI2THOR object navigation. Second, our approach achieves high sample-efficiency in the decoupled RL structure. For instance, CONPE requires less than  $50.0\%$  and  $16.7\%$  of the samples compared to ATC [1] and  $60\%$  and  $50\%$  of the samples compared to EmbCLIP to achieve comparable performance in seen and unseen target domains in the AI2THOR object navigation.
In the context of RL, our work is the first to explore policy adaptation using visual prompts for embodied agents, achieving superior zero-shot performance and high sample-efficiency. The main contributions of our work are as follows.
- We present a novel CONPE framework with an ensemble of visual contrastive prompts, which enables zero-shot adaptation for vision-based embodied RL agents.
- We devise visual prompt-based contrastive learning and guided-attention-based prompt ensemble algorithms to represent task-specific information in the CLIP embedding space.
- We experimentally show that policies via CONPE achieve comparable or superior zero-shot performance, compared to other state-of-the-art baselines, for several tasks. We also demonstrate high sample-efficiency in policy learning and adaptation.
- We create the datasets with various visual domains in AI2THOR, egocentric-Metaworld and CARLA, and make them publicly accessible for further research on RL policy adaptation.
In RL formulation, a learning environment is defined as a Markov decision process (MDP) of  $(S,A,\mathcal{P},R)$  with state space  $s\in S$ , action space  $a\in A$ , transition probability  $\mathcal{P}:S\times A\to S$  and reward function  $R:S\times A\to \mathbb{R}$ . The objective of RL is to find an optimal policy  $\pi^{*}:S\rightarrow A$  maximizing the sum of discounted rewards. For embodied agents, states might not be fully observable, and the environment is represented by a partially observable MDP (POMDP) of a tuple  $(S,A,\mathcal{P},R,\Omega ,\mathcal{O})$  with an observation space  $o\in \Omega$  and a conditional observation probability [11]  $\mathcal{O}:S\times A\to \Omega$ .
Given visual domains in the dynamic environment, we consider policy adaptation to find the optimal policy that remains invariant across the domains or is transferable to some target domain, where each domain is represented by a POMDP and domain changes are formulated by different  $\mathcal{O}$ . We denote domains as  $D = (\Omega, \mathcal{O})$ . Aiming to enable zero-shot adaptation to various domains, we formulate
Figure 1: Visual Domain Changes of Embodied Agents
(i) Prompt-based Contrastive Learning
(ii) Guided-Attention-based Prompt Ensemble (iii) Zero-shot Policy Deployment
Figure 2: CONPE Framework. The CLIP visual encoder is enhanced offline via (i) prompt-based contrastive learning that generates the visual prompt pool, and a policy is learned online by (ii) guided-attention-based prompt ensemble that uses the prompt pool. In (iii) zero-shot deployment, the policy is immediately evaluated upon domain changes.
the policy adaptation problem as finding the optimal policy  $\pi^{*}$  such that
$$
\pi^ {*} = \operatorname {a r g m a x} _ {\pi} \left[ \mathbb {E} _ {D \sim p (D)} \left[ \sum_ {t = 1} ^ {\infty} \gamma^ {t} R \left(s _ {t}, \pi \left(o _ {t}\right)\right) \right] \right] \tag {1}
$$
where  $p(D)$  is a given domain distribution and  $\gamma$  is a discount factor of the environment.
For embodied agents, the same state can be differently observed depending on the configuration of properties such as egocentric camera position, stride length, illumination, and object style. We refer to such a property causing domain changes in the environment as a domain factor. Practical scenarios often involve the interplay of multiple domain factors in the environment.
3 Our Approach
3.1 Framework Structure
To enable zero-shot policy adaptation to unseen domains, we develop the CONPE framework consisting of (i) prompt-based contrastive learning with the CLIP visual encoder, (ii) guided-attention-based prompt ensemble, and (iii) zero-shot policy deployment, as illustrated in Figure 2. The capability of the CLIP visual encoder is enhanced using multiple visual prompts that are contrastively learned on expert demonstrations for several domain factors. This establishes the visual prompt pool in (i). Then, the prompts are used to train the guided-attention-based ensemble with the environment in (ii). To enhance learning efficiency and interpretability of attention weights, we use the cosine similarity of embeddings. The attention module and policy are jointly learned for a specific task so that resulting state representations tend to generalize across various domains and be optimized for task learning. In deployment, a non-stationary environment where its visual domain varies according to the environment conditions and agent physical properties is considered, and the zero-shot performance is evaluated in (iii).
3.2 Prompt-based Contrastive Learning
To construct domain-invariant representations with respect to a specific domain factor for egocentric perception data, we adopt several contrastive tasks for visual prompt learning, which can be learned on a few expert demonstrations. For this, we use a visual prompt
$$
p ^ {v} = \left[ e _ {1} ^ {v}, e _ {2} ^ {v}, \dots , e _ {u} ^ {v} \right], e _ {i} ^ {v} \in \mathbb {R} ^ {d} \tag {2}
$$
where  $e_i^v$  is a continuous learnable vector with the image patch embedding dimension  $d$  (e.g., 768 for CLIP visual encoder) and  $u$  is the length of a visual prompt. Let a pretrained model  $T_{\phi}$  parameterized by  $\phi$  maps observations  $o \in \Omega$  to the embedding space  $\mathcal{Z}$ . With a contrast function  $P: \Omega \times \Omega \to \{0,1\}$  [1, 2, 12] to discriminate whether an observation pair is positive or not,
Figure 3: Guided-Attention-based Prompt Ensemble. The cosine similarity-guided attention module  $\mathcal{G}$  yields task-specific state representations from multiple prompted embeddings and is learned with a policy network  $\pi$ .
consider an  $m$ -sized batch of observation pairs  $\mathcal{B}_P = \{(o_i,o_i')\}_{i\leq m}$  containing one positive pair  $\{(o_k,o_k')|P(o_k,o_k') = 1\}$  for some  $k\leq m$ . Then, we enhance the capability of  $\mathcal{T}_{\phi}$  by learning a visual prompt  $p^v$  through contrastive learning, where the contrastive loss function [13] is defined as
$$
\mathcal {L} _ {\mathrm {C O N}} \left(p ^ {v}, \mathcal {B} _ {P}\right) = - \log \left(\frac {S \left(\mathcal {T} _ {\phi} \left(o _ {k} , p ^ {v}\right) , \mathcal {T} _ {\phi} \left(o _ {k} ^ {\prime} , p ^ {v}\right)\right)}{\sum_ {i \neq k} S \left(\mathcal {T} _ {\phi} \left(o _ {i} , p ^ {v}\right) , \mathcal {T} _ {\phi} \left(o _ {i} ^ {\prime} , p ^ {v}\right)\right)}\right), S (x, y) = \frac {1}{\lambda} \exp \left(\frac {\langle x , y \rangle}{\| x \| \| y \|}\right). \tag {3}
$$
As in [14], for latent vectors  $x, y \in \mathcal{Z}$ , their similarity in the embedding space  $\mathcal{Z}$  is calculated by  $S(x,y)$ , where  $\lambda$  is a hyperparameter. By conducting the prompt-based contrastive learning on  $n$  different domain factors, we obtain a visual prompt pool
$$
\mathbf {p} ^ {v} = \left[ p _ {1} ^ {v}, p _ {2} ^ {v}, \dots , p _ {n} ^ {v} \right]. \tag {4}
$$
Through this process, each visual prompt in  $\mathbf{p}^v$  encapsulates domain-invariant knowledge pertinent to its respective domain factor.
3.3 Guided-Attention-based Prompt Ensemble
To effectively integrate individual prompted embeddings from multiple visual prompts into a task-specific state representation, we devise a guided-attention-based prompt ensemble structure, as shown in Figure 3 where the attention weights on the embeddings are dynamically computed via the attention module  $\mathcal{G}$  for each observation.
Given observation  $o$  and the learned visual prompt pool  $\mathbf{p}^v$ , an image embedding  $z_0 = \mathcal{T}_{\phi}(o)$  and prompted embeddings  $\mathbf{z} = [z_1 = \mathcal{T}_{\phi}(o,p_1^v),\dots,z_n]$  are calculated. Then,  $z_o$  and  $\mathbf{z}$  are fed to the attention module  $\mathcal{G}$ , where attention weights  $\omega_{i}$  for each prompted embedding  $z_{i}$  are optimized. Since directly computing the attention weights using  $z_0$  and  $\mathbf{z}$  is prone to have an uninterpretable local optima, we introduce a guidance score  $g_{i}$  based on the cosine similarity between the input image and visual prompted image embeddings in  $\mathcal{Z}$ , i.e.,  $g_{i} = \frac{\langle z_{0},z_{i}\rangle}{\|z_{0}\|\|z_{i}\|}$ . Given that larger  $g_{i}$  signifies a stronger conformity of an observation to the domain factor relevant to the prompted embedding  $z_{i}$ , we use  $g_{i}$  to steer the attention weights, aiming to not only improve learning efficiency but also provide interpretability. With guidance  $g_{i}$ , we compute the attention weights  $\omega_{i}$  by
$$
\omega_ {i} = \frac {\exp \left(u _ {i} / \tau\right)}{\sum_ {k} \exp \left(u _ {k} / \tau\right)}, \quad u _ {i} = \frac {\langle z _ {0} , k _ {i} \rangle}{\sqrt {d}} g _ {i} \tag {5}
$$
Algorithm 1 Procedure of CONPE Framework
Dataset  $\mathcal{D} = \{(o_1,o_1'),\ldots \}$  , replay buffer  $Z_{D}\gets \emptyset$  , pretrained vision-language model  $\mathcal{T}_{\phi}$  Visual prompt pool  $\mathbf{p}^v = [p_1^v,\dots,p_n^v ]$  , attention module  $\mathcal{G}$  , policy  $\pi$  1: /* Prompt-based Contrastive Learning */   
2: for  $i = 1,\dots,n$  do   
3: while not converge do   
4: Sample a batch  $\mathcal{B}_{P_i} = \{(o_j,o_j')\}_{j\leq m}\sim \mathcal{D}$    
5: Update prompt  $p_i^v\gets p_i^v -\nabla \mathcal{L}_{\mathrm{CON}}(p_i^v,\mathcal{B}_{P_i})$  using (3)   
6: end while   
7: end for   
8: /* Prompt Ensemble-based Policy Learning */   
9: for each environment step do   
10: Sample action  $a = \pi (\mathcal{G}(\mathcal{T}_{\phi}(o),\mathbf{z}))$  using (5), (6)   
11:  $Z_{D}\gets Z_{D}\cup \{(\mathbf{z},a,r)\}$    
12: Jointly optimize policy  $\pi$  and module  $\mathcal{G}$  on  $\{(z_j,a_j,r_j)\}_{j\leq m}\sim Z_D$    
13: end for
where  $k_{i}$  is the projection of  $z_{i}$ ,  $d$  is dimension of  $z$ , and  $\tau$  is a softmax temperature. Then, state embedding  $Z$  is obtained by
$$
Z = \mathcal {G} \left(z _ {0}, \mathbf {z}\right) = z _ {0} + \sum_ {i = 1} ^ {n} \omega_ {i} z _ {i}. \tag {6}
$$
Algorithm 1 shows the procedures in CONPE, where the first half corresponds to prompt-based contrastive learning (in Section 3.2) and the other half corresponds to joint learning of a policy  $\pi(Z)$  and the attention module  $\mathcal{G}$ . As  $\mathcal{G}$  is optimized by a given RL task objective in the source domains (in line 12), the resulting  $Z$  tends to be task-specific, while  $Z$  is also domain-invariant by the ensemble of contrastively learned visual prompts based on  $\mathcal{G}$  with respect to the combinations of multiple domain factors. The entire algorithm can be found in Appendix.
Experiments. We use AI2THOR [9], Metaworld [15], and CARLA [10] environments, specifically configured for embodied agent tasks with dynamic domain changes. These environments allow us to explore various domain factors such as camera settings, stride length, rotation degree, gravity, illuminations, wind speeds, and others. For prompt-based contrastive learning (in Section 3.2), we use a small dataset of expert demonstrations for each domain factor (i.e., 10 episodes per domain factor). For prompt ensemble-based policy learning (in Section 3.3), we use a few source domains randomly generated through combinatorial variations of the seen domain factors (i.e., 4 source domains). In our zero-shot evaluations, we use target domains that can be categorized as either seen or unseen. The seen target domains are those encountered during the phase of prompt-based contrastive learning, while these domains are not present during the phase of prompt ensemble-based policy learning. On the other hand, the unseen target domains refer to those that are entirely new, implying that they are not encountered during either learning phases.
Baselines. We implement several baselines for comparison. LUSR [16] is a reconstruction-based domain adaptation method in RL, which uses the variational autoencoder structure for robust representations. CURL [2] and ACT [1] employ contrastive learning in RL frameworks for high sample-efficiency and generalization to visual domains. ACO [12] utilizes augmentation-driven and behavior-driven contrastive tasks in the context of RL. EmbCLIP [7] is a state-of-the-art embodied AI model, which exploits the pretrained CLIP visual encoder for visual state representations.
Implementation. We implement CONPE using the CLIP model with ViT-B/32, similar to VPT [17] and CoOp [18]. In prompt-based contrastive learning, we adopt various contrastive learning schemes including augmentation-driven [2, 1, 19] and behavior-driven [12, 20, 21, 22] contrastive learning, where the prompt length sets to be 8. In policy learning, we exploit online learning (i.e., PPO [23]) for AI2THOR and imitation learning (i.e., DAGGER [24]) for egocentric-Metaworld and CARLA.
Table 1: Zero-shot Performance. The policies of each method (CONPE and the baselines) are learned on 4 source domains. The Source column presents the performance for those source domains. In all evaluations, we use 30 seen target domains and 10 unseen target domains. The Seen Target column presents the performance for the seen target domains, and the Unseen Target column presents the performance for the unseen target domains. The unseen target domains are not used for representation learning.
(a) Zero-shot Performance in AI2THOR with Object and Point Goal Navigation Tasks
Method	ObjectNav.	PointNav.
Source	Seen Target	Unseen Target	Source	Seen Target	Unseen Target
LUSR	53.3±1.1	21.3±1.9	15.1±1.8	85.6±4.6	71.8±3.8	62.4±5.8
CURL	51.3±1.0	8.0±0.1	6.9±1.3	70.8±7.4	55.2±2.7	54.8±3.0
ATC	82.2±9.7	72.3±3.3	51.3±8.6	95.0±3.3	89.1±1.9	81.9±3.6
ACO	55.0±23.8	39.6±21.5	35.8±5.8	91.1±6.3	73.4±2.0	67.5±2.8
EmbCLIP	89.3±3.0	77.6±1.3	59.0±6.4	95.3±4.6	84.5±1.9	77.4±1.4
CONPE	96.3±1.0	83.3±0.3	79.7±6.4	97.8±1.0	89.7±1.6	84.3±2.0
(b) Zero-shot Performance in egocentric-Metaworld with Reach and Reach-wall Tasks
Method	Reach	Reach-Wall
Source	Seen Target	Unseen Target	Source	Seen Target	Unseen Target
LUSR	100.0±0.0	46.0±15.1	44.7±2.3	50.0±10.0	33.3±6.1	30.7±6.4
CURL	100.0±0.0	53.3±5.0	46.7±3.1	43.3±15.3	2.0±0.0	0.7±1.2
ATC	100.0±0.0	71.3±8.1	72.0±2.0	66.7±5.8	5.3±1.2	4.0±0.0
ACO	100.0±0.0	52.0±2.0	44.0±3.5	63.3±15.3	8.7±2.3	4.7±1.2
EmbCLIP	100.0±0.0	64.7±6.1	66.7±4.2	100.0±0.0	58.0±7.2	49.3±5.0
CONPE	100.0±0.0	88.7±3.1	86.7±3.1	100.0±0.0	75.3±3.1	67.3±2.3
(c) Zero-shot Performance in CARLA with Different Maps
Method	Map 1	Map 2
Source	Seen Target	Unseen Target	Source	Seen Target	Unseen Target
LUSR	2141.9	635.1±606.2	1073.9±212.6	2279.6	1173.7±914.3	2159.4±146.5
CURL	945.4	864.2±638.0	1256.0±61.6	1050.1	1089.9±824.0	2190.3±10.2
ATC	2280.5	1684.4±368.2	1073.7±618.8	2272.2	2253.9±218.7	2200.1±307.8
ACO	2265.8	1545.6±596.1	1330.0±144.5	2270.6	2360.9±88.0	2415.5±53.0
EmbCLIP	2235.7	1732.2±588.6	1415.1±669.9	2262.7	2139.1±655.9	2401.3±12.3
CONPE	2237.5	1738.0±163.5	1933.4±29.7	2277.2	2422.5±79.6	2512.9±15.7
4.1 Zero-shot Performance
Table 1 shows zero-shot performance of CONPE and the baselines across source, seen and unseen target domains. We evaluate with 3 different seeds and report the average performance (i.e., task success rate in AI2THOR and egocentric-Metaworld, the sum of rewards in CARLA). As shown in Table 1(a), CONPE outperforms the baselines in the AI2THOR tasks. It particularly surpasses the most competitive baseline, EmbCLIP, by achieving  $5.2 \sim 5.7\%$  higher success rate for seen target domains, and  $6.9 \sim 20.7\%$  for unseen target domains. For egocentric-Metaworld, as shown in Table 1(b), CONPE demonstrates superior performance with a significant success rate for both seen and unseen target domains, which is  $17.3 \sim 24.0\%$  and  $18.0 \sim 20.0\%$  higher than EmbCLIP, respectively. For autonomous driving in CARLA, we take into account external environment factors, such as weather conditions and times of day, as domain factors that can influence the driving task. In Table 1(c), CONPE consistently maintains competitive zero-shot performance across all conditions, outperforming the baselines.
In these experiments, LUSR shows relatively low success rates, as the reconstruction-based representation model can abate some task-specific information from observations, which is critical to conduct vision-based complex RL tasks. EmbCLIP shows the most comparative performance among the baselines, but its zero-shot performance for target domains is not comparable to CONPE. In contrast,
CONPE effectively estimates the domain shifts pertaining to each domain factor through the use of guided attention weights, leading to robust performance in both seen and unseen target domains.
(a) Success Rate in Seen Target Domain
(b) Success Rate in Unseen Target Domain
Figure 4: Sample-efficiency of Prompt Ensemble-based Policy Learning for Object Navigation in AI2THOR. The x-axis represents the number of samples (timesteps) used for policy learning, while the y-axis represents the task success rate for zero-shot evaluation.
Sample Efficiency. Figure 4 presents performance with respect to samples (timesteps) that are used by CONPE and baselines for policy learning. Compared to the most competitive baseline EmbCLIP, CONPE requires less than  $60.0\%$  timesteps (online samples) for seen target domains and  $50.0\%$  for unseen target domains to have comparable success rates.
(a) Prompted Embeddings
(b) Attention Weights
Figure 5: Prompt Ensemble Interpretability. In (a), the embeddings in the big circle are intra prompted embeddings obtained by varying domains within a domain factor, and the embeddings in the rectangle are inter prompted embeddings obtained by changing the visual prompts with aligned observation. The closely located intra prompted embeddings indicate the domain-invariant knowledge, while the inter prompted embeddings clustered by different visual prompts indicate the alignment between the visual prompts and the domain factors. In (b), each cell represents attention weight  $\omega_{i}$  applied for prompted embedding  $z_{i}$ .
Prompt Ensemble Interpretability. Figure 5(a) visualizes the prompted embeddings using the prompt pool obtained through CONPE. For intra prompted embeddings, we use observation pairs, where each pair is generated by varying domains within a domain factor. We observe that the embeddings are paired to form the domain-invariant knowledge because the visual prompt is learned through prompt-based contrastive learning. The inter prompted embeddings specify that each prompt distinctly clusters prompted embeddings that correspond to the domains associated to its domain factor. Figure 5(b) shows examples of the attention weight matrix of CONPE for four different domains. The x-axis denotes the visual prompts and the y-axis denotes timesteps. This shows the consistency of the attention weights on the prompts across the timesteps in the same domain.
4.2 Prompt Ensemble with a Pretrained Policy