SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Changyeon Kim<sup>1</sup> Younggyo Seo<sup>2</sup> Hao Liu<sup>3</sup> Lisa Lee<sup>4</sup> Jinwoo Shin<sup>1</sup> Honglak Lee<sup>5,6</sup> Kimin Lee<sup>1</sup>
$^{1}$ KAIST  $^{2}$ Dyson Robot Learning Lab  $^{3}$ UC Berkeley  $^{4}$ Google DeepMind  $^{5}$ University of Michigan  $^{6}$ LG AI Research
Developing an agent capable of adapting to unseen environments remains a difficult challenge in imitation learning. This work presents Adaptive Return-conditioned Policy (ARP), an efficient framework designed to enhance the agent's generalization ability using natural language task descriptions and pre-trained multimodal encoders. Our key idea is to calculate a similarity between visual observations and natural language instructions in the pre-trained multimodal embedding space (such as CLIP) and use it as a reward signal. We then train a return-conditioned policy using expert demonstrations labeled with multimodal rewards. Because the multimodal rewards provide adaptive signals at each timestep, our ARP effectively mitigates the goal misgeneralization. This results in superior generalization performances even when faced with unseen text instructions, compared to existing text-conditioned policies. To improve the quality of rewards, we also introduce a fine-tuning method for pre-trained multimodal encoders, further enhancing the performance. Video demonstrations and source code are available on the project website: https://sites.google.com/view/2023arp.
Imitation learning (IL) has achieved promising results in learning behaviors directly from expert demonstrations, reducing the necessity for costly and potentially dangerous interactions with environments [32, 59]. These approaches have recently been applied to learn control policies directly from pixel observations [7, 36, 57]. However, IL policies frequently struggle to generalize to new environments, often resulting in a lack of meaningful behavior [14, 58, 68, 79] due to overfitting to various aspects of training data. Several approaches have been proposed to train IL agents capable of adapting to unseen environments and tasks. These approaches include conditioning on a single expert demonstration [18, 20], utilizing a video of human demonstration [6, 76], and incorporating the goal image [16, 23]. However, such prior methods assume that information about target behaviors in test environments is available to the agent, which is impractical in many real-world problems.
One alternative approach for improving generalization performance is to guide the agent with natural language: training agents conditioned on language instructions [7, 50, 72]. Recent studies have indeed demonstrated that text-conditioned policies incorporated with large pre-trained multimodal models [22, 56] exhibit strong generalization abilities [46, 66]. However, simply relying on text representations may fail to provide helpful information to agents in challenging scenarios. For example, consider a text-conditioned policy (see Figure 1a) trained to collect a coin, which is positioned at the end of the map, following the text instruction "collect a coin". When we deploy the learned agent to test environments where the coin's location is randomized, it often fails to collect the coin. This is because, when relying solely on expert demonstrations, the agent might mistakenly think that the goal is to navigate to the end of the level (see supporting results in Section 4.1). This example shows the simple text-conditioned policy fails to fully exploit the provided text instruction and suffers from goal misgeneralization (i.e., pursuing undesired goals, even when trained with a correct specification) [15, 64].
(a) Comparison of ARP with baseline
(b) Multimodal reward curve from CoinRun
Figure 1: (a) ARP utilizes the similarity between visual observations and text instructions in the pre-trained multimodal representation space as a reward and then trains the return-conditioned policy using demonstrations with multimodal reward labels. (b) Curves of multimodal reward for fine-tuned CLIP [56] in the trajectory from CoinRun environment. Multimodal reward consistently increases as the agent approaches the goal, and this trend remains consistent regardless of the training and test environment, suggesting the potential to guide agents toward target objects in test environments.
In this paper, we introduce Adaptive Return-conditioned Policy (ARP), a novel IL method designed to enhance generalization capabilities. Our main idea is to measure the similarity between visual observations and natural language task descriptions in the pre-trained multimodal embedding space (such as CLIP [56]) and use it as a reward signal. Subsequently, we train a return-conditioned policy using demonstrations annotated with these multimodal reward labels. Unlike prior IL work that relies on static text representations [46, 51, 67], our trained policies make decisions based on multimodal reward signals computed at each timestep (see the bottom figure of Figure 1a).
We find that our multimodal reward can provide a consistent signal to the agent in both training and test environments (see Figure 1b). This consistency helps prevent agents from pursuing unintended goals (i.e., mitigating goal misgeneralization) and thus improves generalization performance when compared to text-conditioned policies. Furthermore, we introduce a fine-tuning scheme that adapts pre-trained multimodal encoders using in-domain data (i.e., expert demonstrations) to enhance the quality of the reward signal. We demonstrate that when using rewards derived from fine-tuned encoders, the agent exhibits superior generalization performance compared to the agent with frozen encoders in test environments. Notably, we also observe that ARP effectively guides agents in test environments with unseen text instructions associated with new objects of unseen colors and shapes (see supporting results in Section 4.3).
In summary, our key contributions are as follows:
- We propose Adaptive Return-conditioned Policy (ARP), a novel IL framework that trains a return-conditioned policy using adaptive multimodal rewards from pre-trained encoders.
- We introduce a fine-tuning scheme that adapts pre-trained CLIP models using in-domain expert demonstrations to improve the quality of multimodal rewards.
- We show that our framework effectively mitigates goal misgeneralization, resulting in better generalization when compared to text-conditioned baselines. We further show that ARP can execute unseen text instructions associated with new objects of unseen colors and shapes.
- We demonstrate that our method exhibits comparable generalization performance to baselines that consume goal images from test environments, even though our method solely relies on natural language instruction.
- Source code and expert demonstrations used for our experiments are available at https://github.com/csmile-1006/ARP.git
Generalization in imitation learning Addressing the challenge of generalization in imitation learning is crucial for deploying trained agents in real-world scenarios. Previous approaches have shown improvements in generalization to test environments by conditioning agents on a robot demonstration [18, 20], a video of a human performing the desired task [76, 6], or a goal image [16, 44, 23]. However, these approaches have a disadvantage: they can be impractical to adopt in real-world scenarios where the information about target behaviors in test environments is not guaranteed. In this work, we propose an efficient yet effective method for achieving generalization even in the absence of specific information about test environments. We accomplish this by leveraging multimodal reward computed with current visual observations and task instructions in the pre-trained multimodal embedding space.
Pre-trained representation for reinforcement learning and imitation learning Recently, there has been growing interest in leveraging pre-trained representations for robot learning algorithms that benefit from large-scale data [54, 73, 61, 52]. In particular, language-conditioned agents have seen significant advancements by leveraging pre-trained vision-language models [46, 66, 78, 38], drawing inspiration from the effectiveness of multimodal representation learning techniques like CLIP [56]. For example, InstructRL [46] utilizes a pre-trained multimodal encoder [22] to encode the alignment between multiple camera observations and text instructions and trains a transformer-based behavior cloning policy using encoded representations. In contrast, our work utilizes the similarity between visual observations and text instructions in the pre-trained multimodal embedding space in the form of a reward to guide agents in the test environment adaptively.
We provide more discussions on related work in more detail in Appendix B.
In this section, we introduce Adaptive Return-conditioned Policy (ARP), a novel IL framework for enhancing generalization ability using multimodal rewards from pre-trained encoders. We first describe the problem setup in Section 3.1. Section 3.2 introduces how we define the multimodal reward in the pre-trained CLIP embedding spaces and use it for training return-conditioned policies. Additionally, we propose a new fine-tuning scheme that adapts pre-trained multimodal encoders with in-domain data to enhance the quality of rewards in Section 3.3.
We consider the visual imitation learning (IL) framework, where an agent learns to solve a target task from expert demonstrations containing visual observations. We assume access to a dataset  $\mathcal{D} = \{\tau_i\}_{i=1}^N$  consisting of  $N$  expert trajectories  $\tau = (o_0, a_0^*, \dots, o_T, a_T^*)$  where  $o$  represents the visual observation,  $a$  means the action, and  $T$  denotes the maximum timestep. These expert demonstrations are utilized to train the policy via behavior cloning. As a single visual observation is not sufficient for fully describing the underlying state of the task, we approximate the current state by stacking consecutive past observations following common practice [53, 74].
We also assume that a text instruction  $\mathbf{x} \in \mathcal{X}$  that describes how to achieve the goal for solving tasks is given in addition to expert demonstrations. The standard approach to utilize this text instruction is to train a text-conditioned policy  $\pi(a_t | o_{\leq t}, \mathbf{x})$ . It has been observed that utilizing pre-trained multimodal encoders (like CLIP [56] and M3AE [22]) is very effective in modeling this text-conditioned policy [46, 49, 66, 67]. However, as shown in the upper figure of Figure 1a, these approaches provide the same text representations regardless of changes in visual observations. Consequently, they would not provide the agent with adaptive signals when encountering previously unseen environments. To address this limitation, we propose an alternative framework that leverages  $\mathbf{x}$  to compute similarity with the current visual observation within the pre-trained multimodal embedding space. We then employ this similarity as a reward signal. This approach allows the reward value to be adjusted as the visual observation changes, providing the agent with an adaptive signal (see Figure 1b).
3.2 Adaptive Return-Conditioned Policy
Multimodal reward To provide more detailed task information to the agent that adapts over timesteps, we propose to use the visual-text alignment score from pre-trained multimodal encoders.
Specifically, we compute the alignment score between visual observation at current timestep  $t$  and text instruction  $\mathbf{x}$  in the pre-trained multimodal embedding space as follows:
$$
r _ {\phi , \psi} \left(o _ {t}, \mathbf {x}\right) = s \left(f _ {\phi} ^ {\text {v i s}} \left(o _ {t}\right), f _ {\psi} ^ {\text {t x t}} (\mathbf {x})\right). \tag {1}
$$
Here,  $s$  represents a similarity metric in the representation space of pre-trained encoders: a visual encoder  $f_{\phi}^{\mathrm{vis}}$  parameterized by  $\phi$  and a text encoder  $f_{\psi}^{\mathrm{txt}}$  parameterized by  $\psi$ . While our method is compatible with any multimodal encoders and metric, we adopt the cosine similarity between CLIP [56] text and visual embeddings in this work. We label each expert state-action trajectory as  $\tau^{*} = (R_{0},o_{0},a_{0}^{*},\dots,R_{T},o_{T},a_{T}^{*})$  where  $R_{t} = \sum_{i=t}^{T}r_{\phi,\psi}(o_{i},\mathbf{x})$  denotes the multimodal return for the rest of the trajectory at timestep  $t^{1}$ . The set of return-labeled demonstrations is denoted as  $\mathcal{D}^{*} = \{\tau_{i}^{*}\}_{i=1}^{N}$ .
Return-conditioned policy Using return-labeled demonstrations  $\mathcal{D}^*$ , we train return-conditioned policy  $\pi_{\theta}(a_t|o_{\leq t},R_t)$  parameterized by  $\theta$  using the dataset  $\mathcal{D}^*$  and minimize the following objective:
$$
\mathcal {L} _ {\pi} (\theta) = \mathbb {E} _ {\tau^ {*} \sim \mathcal {D} ^ {*}} \left[ \sum_ {t \leq T} l \left(\pi_ {\theta} \left(a _ {t} \mid o _ {\leq t}, R _ {t}\right), a _ {t} ^ {*}\right) \right] \tag {2}
$$
Here,  $l$  represents the loss function, which is either the cross entropy loss when the action space is defined in discrete space or the mean squared error when defined in continuous space.
The main advantage of our method lies in its adaptability in the deployment by adjusting to multimodal rewards computed in test environments (see Figure 1a). At the test time, our trained policy predicts the action  $a_{t}$  based on the target multimodal return  $R_{t}$  and the observation  $o_{t}$ . Since the target return  $R_{t}$  is recursively updated based on the multimodal reward  $r_{t}$ , it can provide a timestep-wise signal to the agent, enabling it to adapt its behavior accordingly. We find that this signal effectively guides the agent to prevent pursuing undesired goals (see Section 4.1 and Section 4.2), and it also enhances generalization performance in environments with unseen text instructions associated with objects having previously unseen configurations (as discussed in Section 4.3).
In our experiments, we implement two different types of ARP using Decision Transformer (DT) [8], referred to as ARP-DT, and using Recurrent State Space Model (RSSM) [25], referred to as ARP-RSSM. Further details of the proposed architectures are provided in Appendix A.
3.3 Fine-Tuning Pre-trained Multimodal Encoders
Despite the effectiveness of our method with pre-trained CLIP multimodal representations, there may be a domain gap between the images used for pre-training and the visual observations available from the environment. This domain gap can sometimes lead to the generation of unreliable, misleading reward signals. To address this issue, we propose fine-tuning schemes for pre-trained multimodal encoders  $(f_{\phi}^{\mathrm{vis}}, f_{\psi}^{\mathrm{xt}})$  using in-domain dataset (expert demonstrations)  $\mathcal{D}$  in order to improve the quality of multimodal rewards. Specifically, we propose fine-tuning objectives based on the following two desiderata: reward should (i) remain consistent within similar timesteps and (ii) be robust to visual distractions.
Temporal smoothness To encourage the consistency of the multimodal reward over timesteps, we adopt the objective of value implicit pre-training (VIP) [52] that aims to learn smooth reward functions from action-free videos. The main idea of VIP is to (i) capture long-range dependency by attracting the representations of the first and goal frames and (ii) inject local smoothness by encouraging the distance between intermediate frames to represent progress toward the goal. We extend this idea to our multimodal setup by replacing the goal frame with the text instruction  $\mathbf{x}$  describing the task objective and using our multimodal reward  $R$  as below:
$$
\begin{array}{l} \mathcal {L} _ {\mathrm {V I P}} (\phi , \psi) = \underbrace {(1 - \gamma) \cdot \mathbb {E} _ {o _ {1} \sim \mathcal {O} _ {1}} [ r _ {\phi , \psi} (o _ {1} , \mathbf {x}) ]} _ {\text {l o n g - r a n g e d e p e n d e n c y l o s s}} \\ + \underbrace {\log \mathbb {E} _ {\left(o _ {t} , o _ {t + 1}\right) \sim \mathcal {D}} \left[ \exp \left(r _ {\phi , \psi} \left(o _ {t} , \mathbf {x}\right) + 1 - \gamma \cdot r _ {\phi , \psi} \left(o _ {t + 1} , \mathbf {x}\right)\right) \right] .} _ {\text {l o c a l s m o o t h n e s s l o s s}} \tag {3} \\ \end{array}
$$
<sup>1</sup>We assume the discount factor  $\gamma$  as 1 in our experiments, and our method can also be applied in the setup with the discount factor less than 1.
(a) CoinRun
(b) Maze I
Figure 2: Environments from OpenAI Procgen benchmarks [10] used in our experiments. We train our agents using expert demonstrations collected in environments with multiple visual variations. We then perform evaluations on environments from unseen levels with target objects in unseen locations. See Section 4.1 for more details.
(c) Maze II
Here  $\mathcal{O}_1$  denotes a set of initial visual observations in  $\mathcal{D}$ . One can see that the local smoothness loss is the one-step temporal difference loss, which recursively trains the  $r_{\phi,\psi}(o_t, \mathbf{x})$  to regress  $-1 + \gamma \cdot r_{\phi,\psi}(o_{t+1}, \mathbf{x})$ . This then induces the reward to represent the remaining steps to achieve the text-specified goal  $\mathbf{x}$  [31], making rewards from consecutive observations smooth.
Robustness to visual distractions To further encourage our multimodal reward to be robust to visual distractions that should not affect the agent (e.g., changing textures or backgrounds), we introduce the inverse dynamics model (IDM) objective [55, 33, 41]:
$$
\mathcal {L} _ {\mathrm {I D M}} (\phi , \psi) = \mathbb {E} _ {(o _ {t}, o _ {t + 1}, a _ {t}) \sim \mathcal {D}} [ l (g (f _ {\phi} ^ {\text {v i s}} (o _ {t}), f _ {\phi} ^ {\text {v i s}} (o _ {t + 1}), f _ {\psi} ^ {\text {t x t}} (\mathbf {x})), a _ {t} ^ {*}) ], \tag {4}
$$
where  $g(\cdot)$  denotes the prediction layer which outputs  $\hat{a}_t$ , predicted estimate of  $a_t$ , and  $l$  represents the loss function which is either the cross entropy loss when the action space is defined in discrete space or the mean squared error when it's defined in continuous space. By learning to predict actions taken by the agent using the observations from consecutive timesteps, fine-tuned encoders learn to ignore aspects within the observations that should not affect the agent.
Fine-tuning objective We combine both VIP loss and IDM loss as the training objective to fine-tune pre-trained multimodal encoders in our model:
$$
\mathcal {L} _ {\mathrm {F T}} (\phi , \psi) = \mathcal {L} _ {\mathrm {V I P}} (\phi , \psi) + \beta \cdot \mathcal {L} _ {\mathrm {I D M}} (\phi , \psi),
$$
where  $\beta$  is a scale hyperparameter. We find that both objectives synergistically contribute to improving the performance (see Table 7 for supporting experiments).
We design our experiments to investigate the following questions:
1. Can our method prevent agents from pursuing undesired goals in test environments? (see Section 4.1 and Section 4.2)
2. Can ARP follow unseen text instructions? (see Section 4.3)
3. Is ARP comparable to goal image-conditioned policy? (see Section 4.4)
4. Can ARP induce well-aligned representation in test environments? (see Section 4.5)
5. What is the effect of each component in our framework? (see Section 4.6)
4.1 Procgen Experiments
**Environments** We evaluate our method on three different environments proposed in Di Langosco et al. [15], which are variants derived from OpenAI Progen benchmarks [10]. We assess the generalization ability of trained agents when faced with test environments that cannot be solved without following true task success conditions.
- CoinRun: The training dataset consists of expert demonstrations where the agent collects a coin that is consistently positioned on the far right of the map, and the text instruction is "The goal
Figure 3: Expert-normalized scores on training/test environments. The result shows the mean and standard variation averaged over three runs. ARP-DT denotes the model that uses pre-trained CLIP representations, and ARP-DT+ denotes the model that uses fine-tuned CLIP representations (see Section 3.3) for computing the multimodal reward.
is to collect the coin." Note the agent may mistakenly interpret that the goal is to proceed to the end of the level, as this also leads to reaching the coin when relying solely on the expert demonstrations. We evaluate the agent in environments where the coin's location is randomized (see Figure 2a) to verify that the trained agent truly follows the intended objective.
- Maze I: The training dataset consists of expert demonstrations where the agent reaches a yellow cheese that is always located at the top right corner, and the text instruction is "Navigate a maze to collect the yellow cheese". The agent may misinterpret that the goal is to proceed to the far right corner, as it also results in reaching the yellow cheese when relying only on expert demonstrations. To verify that the trained agent follows the intended objective, we assess the trained agents in the test environment where the cheese is placed at a random position (see Figure 2b).
- Maze II: The training dataset consists of expert demonstrations where the agent approaches a yellow diagonal line located at a random position, and the text instruction is "Navigate a maze to collect the line". The agent might misinterpret the goal as reaching an object with a yellow color because it also leads to collecting the object with a line shape when relying only on expert demonstrations. For evaluation, we consider a modified environment with two objects: a yellow gem and a red diagonal line. The goal of the agent is to reach the diagonal line, regardless of its color, to verify that the agent truly follows the intended objective (see Figure 2c).
Implementation For all experiments, we utilize the open-sourced pre-trained CLIP model<sup>2</sup> with ViT-B/16 architecture to generate multimodal rewards. Our return-conditioned policy is implemented based on the official implementation of InstructRL [46], and implementation details are the same unless otherwise specified. To collect expert demonstrations used for training data, we first train PPG [11] agents on 500 training levels that exhibit ample visual variations for 200M timesteps per task. We then gather 500 rollouts for CoinRun and 1000 rollouts for Maze in training environments. All models are trained for 50 epochs on two GPUs with a batch size 64 and a context length of 4. Our code and datasets are available at https://github.com/csmile-1006/ARP.git. Further training details, including hyperparameter settings, can be found in Appendix C.
Evaluation We evaluate the zero-shot performance of trained agents in test environments from different levels (i.e., different map layouts and backgrounds) where the target object is either placed in unseen locations or with unseen shapes. To quantify the performance of trained agents, we report the expert-normalized scores on both training and test environments. To report training performance, we measure the average success rate of trained agents over 100 rollouts in training environments and divide it by the average success rate from the expert PPG agent used to collect demonstrations. For test performance, we train a separate expert PPG agent in test environments and compute expert-normalized scores in the same manner.
Baseline and our method As a baseline, we consider InstructRL [46], one of the state-of-the-art text-conditioned policies. InstructRL utilizes a transformer-based policy and pre-trained M3AE [22]
$^{2}$ https://github.com/openai/CLIP
(a) Pick Up Cup
(b) Training/test performance
Figure 4: (a) Image observation of training and test environments for Pick Up Cup task in RLBench benchmarks [34]. (b) Success rates on both training and test environments. The result represents the mean and standard deviation over four different seeds. ARP-RSSM denotes the model that uses frozen CLIP representations for computing the multimodal reward, and ARP-RSSM+ denotes the model that incorporates fine-tuning scheme in Section 3.3.