SlowGuess/NeurIPS_2023_OCR huggingface dataset:
representations for encoding visual observations and text instructions. For our methods, we use a return-conditioned policy based on Decision Transformer (DT) [8, 43] architecture, denoted as ARP-DT (see Appendix A for details). We consider two variations: the model that uses frozen CLIP representations (denoted as ARP-DT) and the model that uses fine-tuned CLIP representations (denoted as ARP-  $\mathrm{DT + }$  ) for computing the multimodal reward. We use the same M3AE model to encode visual observations and the same transformer architecture for policy training. The main difference is that our model uses sequence with multimodal return, while the baseline uses static text representations with the concatenation of visual representations.
Comparison with language-conditioned agents Figure 3 shows that our method significantly outperforms InstructRL in all three tasks. In particular, ARP-DT outperforms InstructRL in test environments while achieving similar training performance. This result implies that our method effectively guides the agent away from pursuing unintended goals through the adaptive multimodal reward signal, thereby mitigating goal misgeneralization. Moreover, we observe that ARP-DT+, which uses the multimodal reward from the fine-tuned CLIP model, achieves superior performance to ARP-DT. Considering that the only difference between ARP-DT and ARP-DT+ is using different multimodal rewards, this result shows that improving the quality of reward can lead to better generalization performance.
4.2 RLBench Experiments
Environment We also demonstrate the effectiveness of our framework on RLBench [34], which serves as a standard benchmark for visual-based robotic manipulations. Specifically, we focus on Pick Up Cup task, where the robot arm is instructed to grasp and lift the cup. We train agents using 100 expert demonstrations collected from environments where the position of the target cup changes above the cyan-colored line in each episode (see the upper figure of Figure 4a). Then, we evaluate the agents in a test environment, where the target cup is positioned below the cyan-colored line (see the lower figure of Figure 4a). The natural language instruction x used is "grasp the red cup and lift it off the surface with the robotic arm." For evaluation, we measure the average success rate over 500 episodes where the object position is varies in each episode.
Setup As a baseline, we consider MV-MWM [62], which initially trains a multi-view autoencoder by reconstructing patches from randomly masked viewpoints and subsequently learns a world model based on the autoencoder representations. We use the same procedure for training the multi-view autoencoders for our method and a baseline. The main difference is that while MV-MWM does not use any text instruction as an input, our method trains a policy conditioned on the multimodal return as well. In our experiments, we closely follow the experimental setup and implementation of the imitation learning experiments in MV-MWM. Specifically, we adopt a single-view control setup where the representation learning is conducted using images from both the front and wrist cameras, but world model learning is performed solely using the front camera. For our methods, we train the return-conditioned policy based on the recurrent state-space model (RSSM) [25], denoted
as ARP-RSSM (see Appendix A for more details). We consider two variants of this model: the model utilizing frozen CLIP representations (referred to as ARP-RSSM) and the model that employs fine-tuned CLIP representations (referred to as ARP-RSSM+). To compute multimodal rewards using both frozen and fine-tuned CLIP, we employ the same setup as in Procgen experiments. Additional details are in Appendix D.
Results Figure 4b showcases the enhanced generalization performance of ARP-RSSM+ agents in test environments, increasing from  $20.37\%$  to  $50.93\%$ . This result implies that our method facilitates the agent in reaching target cups in unseen locations by employing adaptive rewards. Conversely, ARP-RSSM, which uses frozen CLIP representations, demonstrates similar performance to MV-MWM in both training and test environments, unlike the result in Section 4.1. We expect this is because achieving target goals for robotic manipulation tasks in RLBench requires more fine-grained controls than game-like environments.
4.3 Generalization to Unseen Instructions
(a) CoinRun-bluegem
(b) Maze III
Figure 5: Test environments used for experiments in Section 4.3.
We also evaluate our method in test environments where the agent is now required to reach a different object with an unseen shape, color, and location by following unseen language instructions associated with this new object. First, we train agents in environments with the objective of collecting a yellow coin, which is always positioned in the far right corner, and learned agents are tested on
unseen environments where the target object changes to a blue gem, and the target object's location is randomized. This new environment is referred to as CoinRun-bluegem (see Figure 5a), and we provide unseen instruction, "The goal is to collect the blue gem." to the agents. Table 1 shows that our method significantly outperforms the text-conditioned policy (InstructRL) even in CoiunRun-bluegem. This result indicates that our multimodal reward can provide adaptive signals for reaching target objects even when the color and shape change.
Table 1: Expert-normalized scores on CoinRun-bluegem test environments (see Figure 5a).
Model	Test Performance
InstructRL	63.99% ± 3.07%
ARP-DT (Ours)	77.05% ± 2.09%
ARP-DT+ (Ours)	79.06% ± 6.69%
In addition, we verify the effectiveness of our multimodal reward in distinguishing similar-looking distractors and guiding the agent to the correct goal. To this end, we train agents using demonstrations from Maze II environments, where the objective is to collect the yellow line. Trained agents are tested in an augmented version of Maze II test environments: we place a yellow gem, a red diagonal line, and a red
straight line in the random position of the map (denoted as Maze III in Figure 5b), and instruct the trained agent to reach the red diagonal line ( $\mathbf{x} =$  "Navigate a maze to collect the red diagonal line"). Table 2 shows that our method outperforms the baseline in Maze III, indicating that our multimodal reward can provide adaptive signals for achieving goals by distinguishing distractors.
Table 2: Expert-normalized scores on Maze III test environments (see Figure 5b).
Model	Test Performance
InstructRL	21.21% ± 1.52%
ARP-DT (Ours)	33.33% ± 4.01%
ARP-DT+ (Ours)	38.38% ± 3.15%
4.4 Comparison with Goal-Conditioned Agents
We compare our method with goal-conditioned methods, assuming the availability of goal images in both training and test environments. First, it is essential to note that suggested baselines rely on additional information from the test environment because they assume the presence of a goal image during the test time. In contrast, our method relies solely on natural language instruction and
Figure 6: Expert-normalized scores on training/test environments. The result shows the mean and standard variation averaged over three runs. ARP-DT shows comparable or even better generalization ability compared to goal-conditioned baselines.
does not necessitate any extra information about the test environment. As baselines, we consider a goal-conditioned version of InstructRL (denoted as GC-InstructRL), which uses visual observations concatenated with a goal image at each timestep. We also consider a variant of our algorithm that uses the distance of CLIP visual representation to the goal image for reward (denoted as GC-DT).
Figure 6 illustrates the training and test performance of goal-conditioned baselines and ARP-DT. First, we observe that GC-DT outperforms GC-InstructRL in all test environments. Note that utilizing goal image is the only distinction between GC-DT and GC-InstructRL. This result suggests that our return-conditioned policy helps enhance generalization performance. Additionally, we find that ARP-DT demonstrates comparable results to GC-DT and even surpasses GC-InstructRL in all three tasks. Importantly, it should be emphasized that while goal-conditioned baselines rely on the goal image of the test environment (which can be challenging to provide in real-world scenarios), ARP-DT solely relies on natural language instruction for the task. These findings highlight the potential of our method to be applicable in real-world scenarios where the agent cannot acquire information from the test environment.
4.5 Embedding Analysis
To support the effectiveness of our framework in generalization, we analyze whether our proposed method can induce meaningful abstractions in test environments. Our experimental design aims to address the key requirements for improved generalization in test environments: (i) the agent should consistently assign similar representations to similar behaviors even when the map configuration is changed, and (ii) the agent should effectively differentiate between goal-reaching behaviors and misleading behaviors. To this end, we measure the cycle-consistency of hidden representation from trained agents following [3, 42]. For two trajectories  $\tau^1$  and  $\tau^2$  with the same length  $N$ , we first choose  $i \leq N$  and find its nearest neighbor  $j = \arg \min_{j \leq N} ||h(o_{\leq i}^1, a_{<i}^1) - h(o_{\leq j}^2, a_{<j}^2)||_2$ , where  $h(\cdot)$  denotes the output of the causal transformer of ARP-DT (refer to Appendix A for details). In a similar manner, we find the nearest neighbor of  $j$ , which is denoted as  $k = \arg \min_{k \leq N} ||h(o_{\leq k}^1, a_{<k}^1) - h(o_{\leq j}^2, a_{<j}^2)||_2$ . We define  $i$  as cycle-consistent if  $|i - k| \leq 1$ , can return to its original point. The presence of cycle-consistency entails a precise alignment of two trajectories within the hidden space.
In our experiments, we first collect the set of success/failure trajectories from  $N$  different levels in CoinRun test environment, which is denoted as  $\tau_{\mathrm{suc}}^n$  or  $\tau_{\mathrm{fail}}^n$  where  $n \in N$ . Next, we extract hidden representations from trained agents at each timestep across all trajectories. We then measure cycle-consistency across these representations using three different pairs of trajectories (see Figure 7):
1.  $(\tau_{\mathrm{suc}}^{n_1},\tau_{\mathrm{suc}}^{n_2})$  ( $\uparrow$ ): We compute the cycle-consistency between success trajectories from different levels. This indicates whether the trained agents behave in a similar manner in success cases, regardless of different visual contexts.
2.  $(\tau_{\mathrm{fail}}^{n_1},\tau_{\mathrm{fail}}^{n_2})(\uparrow)$ : Similarly, we compute the cycle-consistency between failure trajectories from different levels.
3.  $(\tau_{\mathrm{succ}}^{n_1},\tau_{\mathrm{fail}}^{n_1})(\downarrow)$ : We measure the cycle-consistency between success trajectory and failure trajectory from the same level. This evaluates whether the agent can act differently in success and failure cases.
Figure 7: Visual observations of trajectories in CoinRun environments. We construct 3 different pairs to evaluate whether trained agents have well-aligned hidden representations.
Table 3: We investigate the cycle consistency of trained agents' hidden representations on different sets of trajectories in CoinRun environments. The results are presented as the mean and standard deviation averaged over three different seeds. Scores within one standard deviation from the highest average score are marked in bold.  $(\uparrow)$  and  $(\downarrow)$  denotes that higher/ lower values are better, respectively.
InstructRL	ARP-DT	ARP-DT+
(τl1sacc, l2sacc) (↑)	22.66% ± 4.26%	24.75% ± 1.91%	29.07% ± 1.39%
(τl1fail, τl2fail) (↑)	14.12% ± 2.18%	24.95% ± 0.46%	24.29% ± 0.90%
(τl1sacc, l1fail) (↓)	35.09% ± 2.05%	30.18% ± 1.21%	5.65% ± 0.25%
Note that  $(\uparrow) / (\downarrow)$  implies that higher/ lower value is better, respectively. For implementation, we first select ten different levels in CoinRun test environment with different coin locations. We then collect success and failure trajectories from each level and use the last ten timesteps of each trajectory for measuring cycle-consistency. Table 3 shows the percentage of timesteps that are cycle-consistent with other trajectories from different pairs. Similar to results described in Section 4.1, our proposed methods significantly improve cycle-consistency compared to InstructRL in all cases. Moreover, ARP-DT+, which utilizes the multimodal reward from the fine-tuned CLIP model, outperforms ARP-DT with the frozen CLIP.
Effect of pre-trained multimodal representations To verify the effectiveness of pre-trained multimodal representations, we compare ARP-DT+ with agents using multimodal rewards obtained from a smaller-scale multimodal transformer, which was trained from scratch using VIP and IDM objectives, denoted as ARP-DT  $^+$  (scratch). Table 4 shows a significant decrease in performance for ARP-DT  $^+$  (scratch) when compared to ARP-DT  $^+$  across all environments, par
Table 4: Ablation study of using pre-trained CLIP representations.
Env	Model	Train (%)	Test (%)
CoinRun	ARP-DT+	90.28%	± 1.59%	72.36%	± 3.48%
ARP-DT+ (scratch)	77.08%	± 1.04%	62.48%	± 5.32%
Maze I	ARP-DT+	75.47%	± 2.33%	36.13%	± 0.78%
ARP-DT+ (scratch)	18.87%	± 5.87%	32.52%	± 2.71%
Maze II	ARP-DT+	64.18%	± 3.62%	40.95%	± 2.97%
ARP-DT+ (scratch)	22.51%	± 9.78%	37.62%	± 8.73%
ticularly in the training performance within Maze environments. These findings highlight the crucial role of pre-training in improving the efficacy of our multimodal rewards.
Effect of fine-tuning objectives In Table 5, we examine the effect of fine-tuning objectives by reporting the performance of ARP-DT fine-tuned with or without the VIP loss  $\mathcal{L}_{\mathrm{VIP}}$  (Equation 3) and the IDM loss  $\mathcal{L}_{\mathrm{IDM}}$  (Equation 4). We find that the performance of ARP-DT improves with either  $\mathcal{L}_{\mathrm{VIP}}$  or  $\mathcal{L}_{\mathrm{IDM}}$ , which shows the effectiveness of the proposed losses that encourage temporal smoothness and robustness to visual distractions. We also note that the performance with
Table 5: Ablation study of fine-tuning objectives:  ${\mathcal{L}}_{\mathrm{{VIP}}}$  and  ${\mathcal{L}}_{\mathrm{{IDM}}}$  in CoinRun.
\( \mathcal{L}_{\text{VIP}} \)	\( \mathcal{L}_{\text{IDM}} \)	Train (%)	Test (%)
✘	✘	89.58% ± 2.08%	63.32 % ± 2.01%
✘	✓	89.24% ± 6.01%	67.34 % ± 2.66%
✓	✘	90.28% ± 2.17%	70.35 % ± 1.01%
✓	✓	90.28% ± 1.59%	72.36 % ± 3.48%
both objectives is the best, which implies that both losses synergistically contribute to improving the quality of the rewards.
In this paper, we present Adaptive Return-conditioned Policy, a simple but effective IL framework for improving generalization capabilities. Our approach trains return-conditioned policy using the adaptive signal computed with pre-trained multimodal representations. Extensive experimental results demonstrate that our method can mitigate goal misgeneralization and execute unseen text instructions associated with new objects compared to text-conditioned baselines. We hope our framework could facilitate future research to further explore the potential of using multimodal rewards to guide IL agents in real-world applications.
We want to thank Sihyun Yu, Junsu Kim, and anonymous reviewers for providing helpful feedback and suggestions for improving our paper. This research is supported by Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2022-0-00953, Self-directed AI Agents with Problem-solving Capability; No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST); No.2022-0-00184, Development and Study of AI Technologies to Inexpensively Conform to Evolving Policy on Ethics). This material is based upon work supported by the Google Cloud Research Credits program with the award (RDA2-TJCV-2DJT-HAHE).
[1] Abdal, Rameen, Zhu, Peihao, Femiani, John, Mitra, Niloy, and Wonka, Peter. Clip2stylegan: Unsupervised extraction of stylegan edit directions. In ACM SIGGRAPH 2022 conference proceedings, 2022.
[2] Ahn, Michael, Brohan, Anthony, Brown, Noah, Chebotar, Yevgen, Cortes, Omar, David, Byron, Finn, Chelsea, Gopalakrishnan, Keerthana, Hausman, Karol, Herzog, Alex, et al. Do as i can, not as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691, 2022.
[3] Aytar, Yusuf, Pfaff, Tobias, Budden, David, Paine, Thomas, Wang, Ziyu, and De Freitas, Nando. Playing hard exploration games by watching youtube. Advances in Neural Information Processing Systems, 31, 2018.
[4] Bacchus, Fahiem, Boutilier, Craig, and Grove, Adam. Rewarding behaviors. In Proceedings of the National Conference on Artificial Intelligence, 1996.
[5] Bacchus, Fahiem, Boutilier, Craig, and Grove, Adam. Structured solution methods for nonmarkovian decision processes. In Association for the Advancement of Artificial Intelligence, 1997.
[6] Bonardi, Alessandro, James, Stephen, and Davison, Andrew J. Learning one-shot imitation from humans without humans. IEEE Robotics and Automation Letters, 5(2), 2020.
[7] Brohan, Anthony, Brown, Noah, Carbajal, Justice, Chebotar, Yevgen, Dabis, Joseph, Finn, Chelsea, Gopalakrishnan, Keerthana, Hausman, Karol, Herzog, Alex, Hsu, Jasmine, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817, 2022.
[8] Chen, Lili, Lu, Kevin, Rajeswaran, Aravind, Lee, Kimin, Grover, Aditya, Laskin, Misha, Abbeel, Pieter, Srinivas, Aravind, and Mordatch, Igor. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, 2021.
[9] Cho, Jaemin, Yoon, Seunghyun, Kale, Ajinkya, Dernoncourt, Franck, Bui, Trung, and Bansal, Mohit. Fine-grained image captioning with CLIP reward. In Annual Conference of the North American Chapter of the Association for Computational Linguistics, 2022.
[10] Cobbe, Karl, Hesse, Chris, Hilton, Jacob, and Schulman, John. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, 2020.
[11] Cobbe, Karl W., Hilton, Jacob, Klimov, Oleg, and Schulman, John. Phasic policy gradient. In International Conference on Machine Learning, 2021.
[12] Crowson, Katherine, Biderman, Stella, Kornis, Daniel, Stander, Dashiell, Hallahan, Eric, Castricato, Louis, and Raff, Edward. Vqgan-clip: Open domain image generation and editing with natural language guidance. In Computer Vision-ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23–27, 2022, Proceedings, Part XXXVII, 2022.
[13] Cui, Yuchen, Niekum, Scott, Gupta, Abhinav, Kumar, Vikash, and Rajeswaran, Aravind. Can foundation models perform zero-shot task specification for robot manipulation? In Learning for Dynamics and Control Conference, 2022.
A ARP Architecture Details
Transformer-based policy (ARP-DT) To train agents following adaptive multimodal reward signals in Progen experiments, we introduce a return-conditioned policy  $\pi_{\theta}$  based on Decision Transformer architecture [8, 43]. Specifically, we train a decoder-only transformer to autoregressively model the following sequence:
$$
\langle o _ {1}, R _ {1}, a _ {1}, o _ {2}, R _ {2}, a _ {2}, \dots , o _ {T}, R _ {T}, a _ {T} \rangle
$$
Given the expert trajectory  $\tau$ , we first compute the target returns  $\{R_i^*\}_{i=1}^T$  of expert demonstrations by computing the multimodal reward in Equation 1, and we tokenize the sequence using embedding layer for each modality. Our model  $\pi$  parameterized by  $\theta$  comprises the following components:
Action decoder:  $\pi_{\theta}(\hat{a}_t|o_{\leq t},a_{<  t},R_{\leq t})$
Return decoder:  $\pi_{\theta}(\hat{R}_t|o_{\leq t},a_{<  t},R_{<  t})$  (5)
All tokens are fed into the causal transformer, and it produces output embeddings. The action decoder receives the embedding and predicts action  $\hat{a}_t$  and the return decoder takes the embedding for predicting multimodal return  $\hat{R}_t$  to encourage the agent to be aware of the multimodal return. Following Lee et al. [43], we train the model to predict not only the next action but also the next multimodal return by minimizing the objective below:
$$
\mathcal {L} _ {\pi} (\theta) = \mathbb {E} _ {\tau \sim \mathcal {D}} \left[ \sum_ {t \leq T} \left(\pi_ {\theta} \left(\hat {a} _ {t} \mid o _ {\leq t}, a _ {<   t}, R _ {\leq t}\right), a _ {t} ^ {*}\right) + \lambda \cdot \operatorname {M S E} \left(\pi_ {\theta} \left(\hat {R} _ {t} \mid o _ {\leq t}, a _ {<   t}, R _ {<   t}\right), R _ {t} ^ {*}\right) \right], \tag {6}
$$
where CE is the cross entropy loss, MSE is the mean squared error, and  $\lambda$  is a hyperparameter that adjusts the scale of return prediction. We call ARP methods based on this architecture ARP-DT.
RSSM-based policy (ARP-RSSM) To demonstrate the versatility of our framework across different model architectures, we introduce a return-conditioned policy based on world models [60, 62] for our RLBench experiments. Specifically, we implement the world model as a variant of the recurrent state-space model (RSSM; [25]) and condition it on multimodal return. The world model comprises the following components:
Encoder:  $s_t \sim f_\theta(s_t | s_{t-1}, a_{t-1}, o_t, R_t)$
Decoder:  $\left\{ \begin{array}{ll}\hat{o}_t\sim p_\theta (\hat{o}_t\mid s_t)\\ \hat{R}_t\sim p_\theta (\hat{R}_t\mid s_t) \end{array} \right.$  (7)
Dynamics model:  $\hat{s}_t\sim p_\theta (\hat{s}_t|s_{t - 1},a_{t - 1})$
Policy:  $\hat{a}_t\sim p_\theta (\hat{a}_t|s_t)$
The encoder extracts state  $s_t$  from previous state  $s_{t-1}$ , previous action  $a_{t-1}$ , current observation  $o_t$ , and the target multimodal return  $R_t = \sum_{i=t}^{T} R(o_i, \mathbf{x})$  with a recurrent architecture. The decoder reconstructs  $o_t$  to provide a learning signal for model states and predicts  $R_t$  to encourage the agent to be aware of the multimodal return. The policy predicts action  $a_t$  using state  $s_t$ . All model parameters  $\theta$  are optimized by minimizing the objective below:
$$
\begin{array}{l} \mathcal {L} (\theta) = \mathbb {E} _ {\tau \sim \mathcal {D}} \left[ \sum_ {t \leq T} - \ln p _ {\theta} \left(z _ {t} \mid s _ {t}\right) - \ln p _ {\theta} \left(R _ {t} ^ {*} \mid s _ {t}\right) - \ln p _ {\theta} \left(a _ {t} ^ {*} \mid s _ {t}\right) \right. \tag {8} \\ \left. + \beta \operatorname {K L} \left[ q _ {\theta} \left(s _ {t} \mid s _ {t - 1}, a _ {t - 1}, z _ {t}, R _ {t}\right) \| p _ {\theta} \left(\hat {s} _ {t} \mid s _ {t - 1}, a _ {t - 1}\right) \right] \right] \\ \end{array}
$$
We refer to ARP methods based on this architecture as ARP-RSSM.
Training instruction-following agents Humans excel at understanding and utilizing language instructions to adapt to unfamiliar situations. Consequently, there has been significant interest in training policies that incorporate natural language in IL [48, 69, 36, 7] by learning policy conditioned on both current observation and the text instruction of the task. In parallel, recent studies have leveraged large language models (LLMs) for generating inner plans over pre-defined skills from natural language instructions for solving various robotic manipulation tasks [2, 45, 30, 27, 17]. Our method can be thought of as one of language-conditioned imitation learning, which leverages natural language instructions as a reward signal by utilizing the similarity between visual observations and natural language instructions in the pre-trained multimodal embedding space.