SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jialu Gao $^{1*}$ , Kaizhe Hu $^{1,2,3*}$ , Guowei Xu $^{1}$ , Huazhe Xu $^{1,2,3}$
$^{1}$  Tsinghua University  $^{2}$  Shanghai Qi Zhi Institute  $^{3}$  Shanghai AI Lab
gaojialululu@gmail.com, huazhe_xu@mail.tsinghua.edu.cn
Pre-trained text-to-image generative models can produce diverse, semantically rich, and realistic images from natural language descriptions. Compared with language, images usually convey information with more details and less ambiguity. In this study, we propose Learning from the Void (LfVoid), a method that leverages the power of pre-trained text-to-image models and advanced image editing techniques to guide robot learning. Given natural language instructions, LfVoid can edit the original observations to obtain goal images, such as "wiping" a stain off a table. Subsequently, LfVoid trains an ensembled goal discriminator on the generated image to provide reward signals for a reinforcement learning agent, guiding it to achieve the goal. The ability of LfVoid to learn with zero in-domain training on expert demonstrations or true goal observations (the void) is attributed to the utilization of knowledge from web-scale generative models. We evaluate LfVoid across three simulated tasks and validate its feasibility in the corresponding real-world scenarios. In addition, we offer insights into the key considerations for the effective integration of visual generative models into robot learning workflows. We posit that our work represents an initial step towards the broader application of pre-trained visual generative models in the robotics field. Our project page: LfVoid.github.io.
What is the simplest way to provide guidance or goals to a robot? This question is answered multiple times: a set of expert demonstrations [1, 2, 3, 4], goal images [5, 6, 7], or natural language instructions [8, 9]. However, these answers either require a laborious and sometimes prohibitive effort to collect data, or contain ambiguity across modalities. The desire to alleviate the challenges begs the question: can we generate goal images from natural language instructions directly, without physically achieving the goals for robots?
Large text-to-image generative models [10, 11, 12, 13] have achieved exciting breakthroughs, demonstrating an unprecedented ability to generate plausible images that are semantically aligned with given text prompts. Trained on extensive datasets, these models are thought to possess a basic understanding of the world [14]. Therefore, we aim to harness the power of these large generative models to provide unambiguous visual goals for robotic tasks without any in-domain training.
In the past, a common approach to leverage off-the-shelf large generative models in robot learning involves using these models for data augmentation on expert datasets to improve policy general-
ization [15, 16, 17]. Despite the success, these methods still rely on human demonstrations, not generative models, as the main source of guidance.
DALL-E-Bot [18], an early attempt to utilize web-scale generative models in a zero-shot manner for guiding manipulation tasks, utilizes DALL-E 2 [10] to generate goal images for object rearrangement tasks. However, the generated images are often too diverse and largely disagree with real-world scenarios, requiring segmentation masks for object matching and a rule-based transformation planner to close the visual discrepancy.
The recent uprising of language-based image editing methods [19, 20, 21] provides a new paradigm: editing the generated images to align with language descriptions, such as "wiping a stain off a table", while keeping the visual appearance of the other objects and the background roughly unchanged. However, the edited images are usually examined by visual appearance rather than embodied tasks.
In this work, we introduce Learning from the Void (LfVoid), a method that uses image editing techniques on pre-trained diffusion models to generate visual goals for Reinforcement Learning (RL). LfVoid contains a unique image editing pipeline capable of performing appearance-based and structure-based editing to generate goal images according to language instructions. With these goal images, LfVoid improves upon example-based RL methods [5, 6] to solve several robot control tasks without the need for any reward function or demonstrations.
We evaluate LfVoid on three simulated environment tasks and the corresponding real-world scenarios. Empirical results show that LfVoid can generate goal images with higher fidelity to both the editing instructions and the source images when compared to existing image editing techniques, leading to better downstream control performance in comparison to language-guided and other image-editing-based methods. We also validate the feasibility of LfVoid in real-world settings: LfVoid can retrieve meaningful reward signals from generated images comparable to those from true goal images. Based on these observations, we discuss key considerations for the current generative model research when aiming for real-world robotic applications.
Our contributions are threefold: First, we propose an effective approach to leverage the knowledge encapsulated in large pre-trained generative models and apply it in a zero-shot manner to guide robot learning tasks. Second, we provide empirical evidence that suggests LfVoid-edited images provide guidance more effectively than other methods including the direct usage of text prompts. Lastly, we identify the existing gap between the capabilities of current text-conditioned image generation models and the demands in robotic applications, thus shedding light on a potential direction for future research in this domain.
Image editing with diffusion models. Recent work in text-to-image diffusion models has demonstrated a strong ability to generate images that are semantically aligned with given text prompts [22, 23, 10, 11, 12, 24, 25]. Based on these generative models, Prompt-to-Prompt proposes to perform text-conditioned editing through the injection of the cross-attention maps. Imagic [20] proposes to perform editing by using interpolation between source and target embeddings to generate images with edited effects. InstructPix2Pix [26] demonstrates another approach for training a diffusion model on the paired source and target images and the corresponding editing instructions. Apart from editing techniques, several methods focusing on controlled editing have been proposed. Textual-Inversion [27] and DreamBooth [28] both aim to learn a special token corresponding to a user-specified object and at inference time can generate images containing that object with a diverse background. Directed Diffusion [21] focuses on object placement generation, which can control the location of a specified object to reside in a given area.
Visual RL with implicit rewards. Standard reinforcement learning methods require an explicit reward function to guide the agent towards desired goal states, yet these reward functions may not always be available or may not capture the essence of the task at hand. One possible way of learning without access to explicit reward functions is using self-supervised learning to learn a representation where the distances can be used as reward signals [29, 30]. Another line of work uses example-based RL methods that aim to utilize the goal states to guide the learning process. VICE, proposed by Fu et al. [5], establishes a general framework for learning from goal observations: it trains a discriminator with the observations in the replay buffer as negative samples and goal images
as positive samples. The positive logits for new observations are used as a reward to guide the agent toward the goal observation. Building upon this, Singh et al. [6] introduce VICE-RAQ, integrating the VICE algorithm with active querying methods through periodically asking a human to provide labels for ambiguous observations. In addition, it proposes a label-smoothing technique for a better reward shaping. Other works [31, 32] also explore ways to improve the reward shaping of VICE as well. In a parallel direction, Eysenbach et al. [7] develop Recursive Classification of Examples (RCE), a method that learns the Q function of actor-critic methods directly from the goal observations.
Large generative models for robot control. Current application of large generative models in robotic tasks mainly focuses on using Large Language Models for planning [8, 9, 33, 34, 35] or learning a language conditioned policy [36, 37], while the application of pre-trained image generative models for robotics tasks are limited. A common way of leveraging large-scale diffusion models for robot learning is by using diffusion models to perform data augmentation on training data, as suggested in CACTI [15], GenAug [16] and ROSIE [17]. Another line of work suggests using diffusion models to generate plans to solve robotics tasks [38, 39, 40]. More recently, video generation and editing techniques have also been used with imitation learning to guide robot policies [41, 42]. More specifically, UniPi[42] trains a text-to-video generation model on web-scale datasets and expert demonstrations to generate image sequences for planning and inverse modeling. The most related work is DALL-E-Bot [18], which uses the DALL-E 2[10] model to generate goal plans and performs object rearrangement according to these goals. However, unlike our work, DALL-E-Bot does not apply image editing techniques for goal generation and requires rule-based matching and predefined pick-and-place actions to bridge the visual gap between generated images and true observations.
In this study, we aim to provide zero-shot visual goals to reinforcement learning agents for manipulation tasks using only text prompts. Our approach utilizes the information grounded in large-scale pre-trained visual generative models to create semantically meaningful visual goals. We first generate a synthetic goal image dataset from raw observations using image editing techniques, then employ example-based visual reinforcement learning that is optimized for our task with the generated dataset.
The structure of this section unfolds as follows: Section 3.1 elucidates the image editing techniques and adjustments applied to create goal images from text prompts and initial image observations, as depicted in Figure 1(a). Section 3.2 discusses the execution of example-based visual reinforcement learning using the generated goal images, as outlined in Figure 1(b).
3.1 Visual goal generation
In the first phase of our methodology, we edit and generate goal images from raw observations based on natural language guidance provided by humans. Given a source prompt  $\mathcal{P}$ , a source image  $x_{src}$ , and an editing instruction, LfVoid synthesizes a target image  $x_{tgt}$  using a pre-trained Latent Diffusion Model (LDM) [12]. We consider two types of editing instructions: either a target prompt  $\mathcal{P}^*$  describing the appearance changes in the target image, or a bounding-box region  $\mathcal{B}$  and a set of tokens  $\mathcal{I}$  corresponding to the object to be relocated when structural changes are needed.
To provide sufficient visual guidance for downstream reinforcement learning tasks, the generated goal images should highlight the visual changes while preserving the irrelevant scene as much as possible, which is a challenging requirement even for state-of-the-art image editing techniques. To this end, we integrate a number of different techniques in the visual goal generation pipeline of LfVoid, consisting of a feature extracting module, an inversion module, and an editing module, as outlined in Figure 1(a). Following conventional notations, we denote  $x_{t}$  as the synthesized LDM image at time step  $t \in \{T, T - 1, \dots, 0\}$ , where  $T$  is the total diffusion time steps. Specifically,  $x_{T}$  denotes the Gaussian noise, and  $x_{0}$  denotes the generated image. The techniques used in each module and their purpose are described as follows.
3.1.1 Feature extracting module
To ensure high fidelity to the source image  $x_{src}$ , we learn a unique token  $sks$  that encapsulates the visual features of objects within  $x_{src}$  as in DreamBooth [28]. This process involves optimizing the diffusion model parameters along with the special token  $sks$ , using a set of images that contain the
Figure 1: An overview of LfVoid. The goal of LfVoid is to learn robot policies requiring only language descriptions from humans. LfVoid consists of two parts: (a) Goal image generation, where we apply image editing on the initial observations according to different editing instructions to obtain a visual goal dataset; (b) Example-Based Visual RL, where we perform reinforcement learning on the generated dataset to achieve the desired goal image in various environments.
target object. As a result, we are able to derive a specialized model that can accurately retain key details of  $x_{src}$ , such as the color and texture of a cube to be positioned, or the shape of a Franka robot arm for manipulation. As later demonstrated in our experiments, this module substantially boosts the resemblance in details of our edited images to the corresponding source images.
3.1.2 Inversion module
After employing the special token  $sks$  to capture the essential features of the source scene, we invert the provided source image  $x_{src}$  to a diffusion process. A simple inversion technique using the Denoising Diffusion Implicit Model (DDIM) [23] sampling scheme calculates the samples  $x_{T}, x_{T-1}, \ldots, x_{0}$  by reversing the ODE process. However, the inverted image  $x_{0}$  obtained this way often deviates from  $x_{src}$  due to cumulative errors. Null-text inversion [43] aims to mitigate this discrepancy by fine-tuning the "null-text" embedding used in the classifier-free guidance [44]  $\Phi_{t}$  for each  $t \in \{T, T-1, \ldots, 1\}$  to control the generation process. Thus, with the initial noise sample  $x_{T}$  obtained through DDIM inversion and the optimized "null-text" embedding  $\Phi_{T}, \Phi_{T-1}, \ldots, \Phi_{1}$ , the diffusion process is able to generate the image  $x_{src}'$ , an approximation of the source image  $x_{src}$ . Based on the diffusion process obtained with the inversion module, LfVoid can perform precise and detailed editing control on the source image, which will be discussed in Section 3.1.3.
3.1.3 Editing module
Depending on the specific requirements, we divide the editing tasks into two distinct categories: appearance-based and structure-based image editing. Appearance-based editing involves maintaining the structural layout of the image while altering certain visual aspects, such as cleaning the surface of a table or lighting up an LED bulb. In contrast, structure-based editing involves changes in the layout of the image, such as relocating objects within the image from one area to another.
Appearance-based editing. In tasks involving appearance-based editing, we employ the Prompt-to-Prompt editing control technique. When the Latent Diffusion Model (LDM) generates an image  $x_0$  conditioned on a text prompt  $\mathcal{P}$ , the information encapsulated in  $\mathcal{P}$  influences the diffusion process via the cross-attention layers [12, 11, 19]. Prompt-to-Prompt reveals that the spatial configuration
Figure 2: Image editing with structural changes. The details of P2P-DD technique used in the editing module of LfVoid to perform structural changes. LfVoid performs a combination of P2P attention map injection with Directed Diffusion attention editing on the diffusion process generating the target image  $x_{tgt}$ , based on the Gaussian noise  $x_{T}$  and the optimized null-text embeddings  $\Phi_T, \Phi_{T-1}, \ldots, \Phi_1$  obtained through the Inversion Module of LfVoid.
of the image  $x_0$  largely depends on the cross-attention maps  $M_t$  within these cross-attention layers, especially during the initial diffusion steps. Consequently, it suggests replacing the attention maps  $M_t^*$  of the target image diffusion process with the attention maps  $M_t$  from the source image diffusion process for the first  $N$  time steps, beginning at time step  $t = T$ :
$$
\operatorname {P 2 P E d i t} \left(M, M ^ {*}, t\right) = \left\{ \begin{array}{l l} M _ {t} & \text {i f} t > T - N \\ M _ {t} ^ {*} & \text {o t h e r w i s e} \end{array} \right. \tag {1}
$$
This method ensures that the structure of the source image encapsulated in the attention maps  $M_{t}$  is conserved in the target image in a more controlled manner, which is crucial in downstream RL tasks.
Structure-based editing. A notable limitation of Prompt-to-Prompt editing is the incapacity to spatially move existing objects across the image, i.e., implement structural changes. Therefore, we introduce a novel approach termed P2P-DD for structure-based editing tasks. P2P-DD combines Prompt-to-Prompt control with the idea of Directed Diffusion [21], extending its capability to perform object replacement.
Directed Diffusion [21] illustrates the possibility of achieving object replacement through direct editing of attention maps during the first  $N$  time steps of the generation process. The method performs attention strengthening on the cross-attention maps corresponding to the tokens in  $\mathcal{I}$  (recall that  $\mathcal{I}$  is the token sets representing the object of interest), through calculating a Gaussian strengthening mask (SM) of the bounding-box region  $\mathcal{B}$ . It also performs attention annealing to the remaining area  $\bar{\mathcal{B}}$  by applying a constant weakening mask (WM), and the two masks are weighted by a scalar  $c$ :
$$
\operatorname {D D E d i t} \left(M _ {t}, \mathcal {B}, \mathcal {I}\right) = \left\{ \begin{array}{l l} M _ {t} \odot \mathbf {W M} (\bar {\mathcal {B}}, \mathcal {I}) + c \cdot \mathbf {S M} (\mathcal {B}, \mathcal {I}) & \text {i f} t > T - N \\ M _ {t} & \text {o t h e r w i s e} \end{array} \right. \tag {2}
$$
Our proposed P2P-DD method aims to encourage the background and other details of the image generated by Directed Diffusion with a higher resemblance to the source image. P2P-DD first injects cross-attention maps  $M_t$  into  $M_t^*$  as suggested by Prompt-to-Prompt. This will preserve the structure and background information from the source image, providing a decent starting point for attention map editing. Next, P2P-DD performs pixel value strengthening and weakening on the attention maps  $M_t$  according to the bounding-box  $\mathcal{B}$  and the tokens in  $\mathcal{I}$ , as suggested in the Directed Diffusion, to achieve the desired object placement:
$$
\mathrm {P 2 P - D D E d i t} \left(M _ {t}, M _ {t} ^ {*}, \mathcal {B}, \mathcal {I}\right) = \mathrm {D D E d i t} \left(\mathrm {P 2 P E d i t} \left(M, M ^ {*}, t\right), \mathcal {B}, \mathcal {I}\right) \tag {3}
$$
This combined editing control is only applied in the first  $N$  diffusion time steps, i.e.  $\{T,\dots,T - N + 1\}$ . After the first  $N$  time steps, we cease any control on attention maps and allow the target diffusion
Wipe (Sim)
LED (Sim)
Push (Sim)
Figure 3: Visualization of robot manipulation tasks. We evaluate LfVoid on three simulation tasks: Wipe, Push, and LED, as well as three real-robot environments correspondingly.
Wipe (Real)
LED (Real)
Push (Real)
process to complete in a conventional denoising manner. The details of P2P-DD are shown in Figure 2. Attention visualization of the generation process can be found in Appendix A.1, and we provide a detailed description of the full algorithm in Appendix A.2.
3.2 Example-based visual reinforcement learning
We modify and extend the approach of VICE [5] to devise our method for example-based visual reinforcement learning. At the core of our process, we employ the image editing methods described in Section 3.1 on observations gathered from randomly initialized environments, producing a dataset of 1024 target images for each task. During training, these target images serve as positive samples, while the images sampled from the agent's replay buffer after a number of random exploration steps are treated as negative samples. We train a discriminator using these instances with the binary cross-entropy loss and use the output of the discriminator for new observations for the positive class as the agent reward.
The discriminator shares the CNN encoder with the reinforcement learning agent and performs classification over the output latent representations. To improve the reward shaping, we utilize the label mixup method [6], which performs a random linear interpolation of the 0-1 labels and their corresponding hidden vectors to obtain continuous labels between 0 and 1. We discover that restricting the negative instances from the recent portion of the replay buffer (the last  $5\%$ ) promotes the discriminator's ability to discern subtle differences between the target and current observations. Additionally, we find that our method enjoys an ensemble of discriminators for classification results (i.e., RL rewards), which gives the agent more representative rewards.
For the reinforcement learning backbone algorithm, we use DrQ-v2 [45] for visual RL training. Based on Twin Delayed DDPG (TD3) [46], DrQ-v2 enhances image representation by applying random crop augmentation to the input images. Collectively, these refinements constitute our approach to example-based visual reinforcement learning, making the "learning" from an initial observation and language prompts possible.
In this section, we present a comprehensive evaluation of LfVoid across both simulated environments and real-world robotic tasks. An illustration of each task can be found in Figure 3.
The environments we use are: 1) LED, where the robot reaches for a switch or touches the LED light directly to turn the light from red to green; 2) Wipe, where the robot needs to wipe out stains from the table; 3) Push, where the robot needs to push a red cube to a goal position indicated by a green dot. The simulated tasks are developed based on the Robosuite benchmark, while we provide corresponding real-world tasks for each environment. A full description of the environments is provided in Appendix B.1.
4.1 Goal generation
In this section, we evaluate the ability of LfVoid and other baseline methods to generate goal images according to two types of editing instructions: appearance-based editing and structure-based editing.
Figure 4: Appearance-based goal image editing results. We compare LfVoid against three recent image editing methods: Imagic, InstructPix2Pix, and DALL E-2, as well as three ablations of LfVoid. Note that InstructPix2Pix has different prompts because it requires an editing instruction rather than a description of the image. In both simulation and real-world settings, LfVoid can generate images that are better aligned with given text prompts while preserving the remaining source scene. See Appendix C for more examples.
Baselines and ablations. We select three recent image editing methods as baselines: Imagc [20], InstructPix2Pix [26], and mask-based DALL E-2 editing used in DALL-E-Bot [18](with masks provided by humans). For ablations, DDIM P2P removes both the feature extracting module and the inversion module, and DDIM DD additionally removes P2P in the editing module for structure-based editing. On top of DDIM DD, Null-text DD adds the inversion module back. For both appearance-based and structure-based editing tasks, DreamBooth P2P/P2P-DD ablates the inversion module, and Null-text P2P/P2P-DD ablates the feature extracting module.