Text-Aware Diffusion for Policy Learning

We thank Reviewer KsL9 for their thoughtful comments and helpful feedback on our work. Below, we seek to both address the reviewer’s listed weaknesses and answer the posed questions:

On hyperparameter tuning: we concur with the reviewer that our approach benefits from initial tuning of the noise level, and reward computation terms such as $w_1$ and $w_2$. However, we demonstrate that such hyperparameter settings generally transfer across environments and robotic states with minimal modification. For example, the Humanoid, Dog, and MetaWorld experiments were conducted with the same TADPoLe noise level and $w_1, w_2$ terms. The same $w_1, w_2$ terms were used for Video-TADPoLe, and were also kept consistent between Humanoid and Dog environments - only the noise level range changed (with justification and discussion in Appendix A.1). Therefore, although there are more terms to tune, tuning does not generally need to occur ad hoc but can be preserved flexibly across environments with decent effect.

On relying on external diffusion models for world knowledge: we agree with the reviewer’s analysis that the diffusion model is utilized to provide world knowledge to a policy during optimization. As demonstrated in the paper, this offers a natural way for priors summarized from large-scale pretraining to benefit policy learning, enabling the learning of policies that act and behave according to “natural” poses and motions summarized from natural image and video datasets, as well as enabling flexible text-conditioning. As diffusion models improve in encoding world knowledge, stemming from efforts on understanding better physics, dynamics, etc. we anticipate the quality of policies learned through TADPoLe to directly improve as well.

On user study details: In Table A7 of our rebuttal Appendix page, we provide the raw user study results on text-alignment, which correspond to the checkmark and x-marks in Table 1. Among 25 anonymous and randomly selected online participants, we deem a prompt to be correctly learned when more than 50% of the participants vote for text-behavior alignment (denoted by a checkmark). We observe that TADPoLe generally has the highest positive vote ratios.

Clarifying $\mathbf{\hat\epsilon}_{\boldsymbol{\phi}}$: we thank the reviewer for highlighting the notational discrepancy; indeed, we overload $\mathbf{\hat\epsilon}_{\boldsymbol{\phi}}$: as a standalone it represents a noise prediction, and when utilized as a function it represents the neural network that produces the prediction. We will clarify this overloaded notation in the caption of Figure 3 in the final draft of the paper.

On replacing the diffusion model with CLIP-style: indeed, the reviewer’s suggestion is what VLM-RM is doing, which we compare against in our experiments. Utilizing a diffusion model over a CLIP-based approach has conceptual benefits aligned with the intuition of the reviewer. As a CLIP-based alignment model is trained on clean pairs of image and text, during inference it therefore assumes that the input for the image is always clean. During ordinary retrieval, from a set of natural images, this is not an issue. During policy optimization, however, it essentially has the flexibility to search for the most text-aligned image amongst potentially unnatural-looking poses. Having a base assumption that every query image is “natural”, it may therefore cause the policy to learn an unnatural-looking pose but the model deems highly-aligned with the text prompt. On the other hand, a diffusion model is trained to not only respect text-alignment, but also understand what is a natural-looking image at all; as a generative model, it seeks to generate a natural-looking image and therefore must learn a notion of "naturalness" from its pretraining data. Therefore, an approach like TADPoLe can result in more natural-looking policies learned (as evidenced by the qualitative results reported in Tables 1, 2, and 3 in comparison to VLM-RM).

On VLM-RM standing performance: another distinguishing factor between using a CLIP-style model and a diffusion model in our method is that CLIP provides a deterministic reward signal. In other words, for the same query image and query text prompt pair, the CLIP model will always provide the same alignment reward score. When utilized for policy optimization, for a fixed text prompt, the policy can essentially be seen as taking actions to search for the singular frame that will achieve the highest deterministic alignment score and then maintaining it for perpetuity. Therefore, goal-achievement tasks like Humanoid Stand are very stable for CLIP-based approaches to learn; it learns to achieve a standing position and maintain it. However, there are two drawbacks: having a singular deterministic score which encourages goal-achievement does not translate well to continuous locomotion, where there is no canonical pose to achieve, and the pose that it does achieve may not be aligned with a notion of naturalness as mentioned above.

On failure modes: In our updated submission website, we provide Video-TADPoLe results for the prompts “a person walking to the left” and “a person walking to the right”. Whereas there are seeds of the Humanoid successfully walking with respect to the provided prompt, there are also seeds where the person walks in the opposite direction. How to provide fine-grained control over the text-conditioning to focus on key words such as direction (beyond simply walking) is interesting to explore for future work; furthermore, a potential way to remedy this limitation is with improvements to the underlying text-conditioned diffusion model itself, such as with motion adapters.

We thank Reviewer TWnR for their detailed comments and thorough questions. We seek to address their concerns below:

On pre-existing well-rendered videos: we would like to clarify a potential misunderstanding - our method does not require any pre-existing well-rendered videos for environments TADPoLe is applied to, and therefore there is no source of ground-truth image/video demonstrations for arbitrary environments (no offline videos or demonstrations are collected, and pre-existing expert policies are not needed). Rather, we utilize the rendering function at each timestep on-the-fly to generate a dense text-conditioned reward, as computed by a large-scale generally-pretrained diffusion model. We do not update the diffusion model with in-domain examples whatsoever. We therefore showcase how large-scale pretraining over natural images, videos, and text-alignment can directly transfer to the zero-shot optimization of policies that behave in alignment with text-conditioning, as well as natural-looking goal poses or motions.

On alignment with motivation: we appreciate that the reviewer believes our motivation on using text-conditioned diffusion models for reward generation is clear; in light of our clarification that TADPoLe can be applied to arbitrary environments without demanding well-rendered video examples a priori (and therefore no pre-trained policies are needed either), we believe that our final proposed approach is well-aligned with the goal of making reinforcement learning more practical and scalable. Indeed, TADPoLe offers an avenue for supervising policies that behave not only in alignment with natural language, but also in alignment with natural-looking poses or motions captured from large-scale pretraining on natural images and videos.

On ablations: the reviewer has highlighted three design decisions of interest: the symlog operation, the noise level range, and the selection of the two weights $w_1$ and $w_2$. We utilize the symlog operation to normalize the scales of the computed reward (in conjunction with $w_1$ and $w_2$ terms); indeed, we observe in Figure A1 where over all noise levels, the final computed reward stays roughly between 0 and 1, despite changes to the visual input or text prompt. Furthermore, this consistent normalization enables the reuse of hyperparameter settings (such as noise level, weights $w_1$ and $w_2$) across tasks and environments and even other diffusion models (e.g. Video Diffusion models) for consistent policy learning. We demonstrate this in a quantitative manner in Table A6 of the updated rebuttal Appendix page, by showcasing how removing the symlog normalization reduces consistent policy learning across environments and diffusion models with the same hyperparameter settings. For the noise level range used, we have provided much justification in the Appendix of our original paper submission through a discussion (Section A, A.1) as well as graphical Figures (Figure A.1). Similarly, in our original submission we have also reported that the values of the two weights $w_1$ and $w_2$ were selected through a hyperparameter sweep on Humanoid standing and walking performance (Section 4.1, Implementation), which we provide in Table A4, along with a discussion on their effects in Appendix C.3.

On experimentation: the focus of our work is evaluating zero-shot optimization of policies conditioned on natural language. We highlight that text-conditioned policy learning through foundation models is a recent and active area of exploration. We compare against other recent work in VLM-RM (ICLR 24), LIV (ICML 23), and Text2Reward (ICLR 24), while also proposing novel baselines (ViCLIP-RM). We further highlight the diversity of our baselines to position TADPoLe in a broader context: Text2Reward uses a language-only approach to create a text-conditioned reward function, whereas VLM-RM, LIV, and ViCLIP-RM compute rewards on-the-fly in a multimodal manner. Furthermore, our range of tasks includes not only Dog and Humanoid (Table 1, 3), but also robotic manipulation tasks from MetaWorld (Table 4). We extend our range of tasks to include multiple novel prompts, as well as those with subtle details (Figure 4).

On comparison with Diffusion-Reward: our work focuses on learning text-conditioned policies in a zero-shot manner through a large-scale generally-pretrained diffusion model. Diffusion-Reward (concurrent work, to be published at ECCV 24) is not a natural comparison for our approach, as it does not enable text-conditioned policy learning (their diffusion model is conditioned only on a history of frames), and their method explicitly requires expert video demonstrations from the environment to train their in-domain diffusion model. On the other hand, TADPoLe (along with listed baselines like VLM-RM, LIV, Text2Reward) requires no in-domain or expert demonstrations, and is text-conditioned. However, we have adapted a version of Diffusion-Reward for comparison purposes, where rewards are conditioned not on history but on a natural language input, and reported it in Table A8 of the rebuttal Appendix page. We demonstrate that TADPoLe outperforms the adapted Diffusion-Reward implementation in terms of overall success rate aggregated over all tasks in the suite. We further note that Diffusion-Reward requires multiple denoising steps for each dense reward computation (in practice they use 10 steps), whereas TADPoLe only requires one denoising step per reward computation, thus substantially improving in computation speed and complexity.

We thank Reviewer iiYV for their helpful feedback and suggestions. We try to address their concerns below:

Qualitative evaluation: We report quantitative comparisons whenever available (Tables 1, 3, 4), and we agree that human evaluation may inevitably introduce subjectivity. However, they are necessary when mere quantitative metrics are not sufficient to capture the performance difference, for example when evaluating the novel text-conditioned behaviors. We attempt to make such evaluation as impartial as possible; we use the videos of policy behavior at the last timestep of training without cherry-picking, and query 25 random participants without prior exposure to the task through an anonymized platform to estimate a general response from the human population. We provide details of our qualitative evaluation procedure in Section 4.1, and also additional fine-grained user study results in our updated rebuttal Appendix page (Table A7).

On text prompt scalability and complexity: we agree that exploring text prompt scalability and sensitivity is important. We refer the reviewer to Figure (4), where we report preliminary investigations on how sensitive the method is when text prompts are extended with details (“a person standing” -> “a person standing with hands above head”), and when subtle details are modified in the text conditioning (“a person standing with hands above head” -> “a person standing with hands on hips”). However, despite showing that TADPoLe indeed demonstrates a level of sensitivity and scalability to long, detailed, text prompts, we agree that extremely complex prompts, such as multiple sequential instructions over time, may still be a challenge - developing this is a worthwhile future direction. We also note that our method is generally applicable to both image and video generative models, and therefore provides a way to handle text prompts based on both visual appearance as well as motion. As the modeling power of base text-to-video diffusion models improves, so too do we expect our approach to scale to respect more complex text prompts accurately.

On comparative baselines: we highlight that text-conditioned policy learning through foundation models is a recent and active area of exploration. We compare against other recent work on text-aware rewards in VLM-RM (ICLR 24), LIV (ICML 23), and Text2Reward (ICLR 24), while also proposing novel baselines (ViCLIP-RM). We further highlight the diversity of our baselines: Text2Reward uses a language-only approach to create a text-conditioned reward function, whereas VLM-RM, LIV, and ViCLIP-RM compute rewards on-the-fly in a multimodal manner. In addition, in the new Table A8 of our rebuttal Appendix page, we further compare with LIV and Diffusion-Reward (a concurrent work, to be published at ECCV 24) on Meta-World, where TADPoLe enjoys the best performance.

On computational overhead: Our reward computation only uses one denoising step to generate text-conditioned rewards, which avoids the overhead of multiple iterations of denoising steps usually performed with generating data through diffusion models in vanilla inference. Furthermore, our reward computation is general across diffusion model implementations, and the complexity of the diffusion model can be adjusted (such as through distillation) to fit desired resource constraints while still utilizing our proposed reward computation.

On reward scaling: we discover that the symlog operation helps to normalize the raw computed reward signals across tasks and environments to be roughly on the same scale. We are therefore able to reuse hyperparameter settings (such as noise level, weights $w_1$ and $w_2$) across tasks and environments and even other diffusion models (e.g. Video Diffusion models) for consistent policy learning. We demonstrate this in a quantitative manner in Table A6 of the updated rebuttal Appendix page, by showcasing how removing the symlog normalization reduces consistent policy learning across environments and diffusion models with the same hyperparameter settings.

On long-term dependencies: exploring how TADPoLe performs in tasks requiring long-term planning and dependencies is an interesting and worthwhile direction for future work. Due to the flexibility of text-conditioning, describing a task with long-term dependencies can be done in a way that is more like providing a sparse reward (e.g. describing just the desired goal) versus a manner that is more akin to providing a dense reward through multiple subtasks (e.g. providing detailed instructions). We anticipate the following considerations to potentially help solve long-term dependencies through TADPoLe: choosing appropriate underlying learning techniques (e.g. one that performs exploration, which can discover how to solve tasks that requires long-term dependencies), performing prompt tuning to discover the best prompts for a task, and learning to attend on different portions of a detailed text prompt (e.g. instructions) based off of progress of the policy’s performance.

On task generalization: we would like to clarify that from the perspective of TADPoLe, all environments are unseen and novel; we utilize generally-pretrained diffusion model checkpoints, with no tasks or even examples of the downstream visual environment observed during pretraining. Policy supervision is computed purely from priors captured within diffusion models with large-scale pretraining and no in-domain examples are used to update the diffusion model. We believe that applying TADPoLe across Humanoid, Dog, and MetaWorld already demonstrates strong task generalization capability, as the agents and visual characteristics of each environment differ greatly from each other. Furthermore, we demonstrate how text-conditioned policies can be learned across these distinct environments without minimal modification to its hyperparameters, even generalizing across which pretrained diffusion model is used.

We thank Reviewer z3di for their comments and feedback on our work; we are happy to hear that the reviewer appreciates the novelty of our approach, and we seek to address their concerns below.

On comparisons: we performed apple-to-apple comparisons with three recent text-aware rewards, namely VLM-RM (ICLR 24), LIV (ICML 23), and Text2Reward (ICLR 24). We utilize the exact same underlying model architecture, hyperparameters and optimization procedure, where only the text-aware reward is changed across methods, to make it as comparable as possible (Table 1, 3). We highlight that using diffusion models has superior or comparable quantitative performances against other text-aware rewards under an apple-to-apple comparison, while having superior text-alignment and naturalness benefits (Table 2). To the best of our knowledge, there is not a diffusion reward baseline which we can compare with in an apple-to-apple setup. A concurrent work (to be published at ECCV 24), Diffusion-Reward, is the only work that leverages diffusion models for reward modeling. However, it does not condition on natural language, and thus does not enable text-conditioned policy learning, which is a central focus of our work. Furthermore, Diffusion-Reward relies on in-domain expert video demonstrations, while TADPoLe does not.

On the use of TD-MPC: we utilize TD-MPC as a standardized backbone to evaluate different text-aware reward methods in a comparable setting, to focus on text-conditioned reward quality. We choose TD-MPC due to it being the first model able to solve the complex Dog environment (given ground-truth rewards for walking and standing), with strong performance on Humanoid as well. Given its powerful modeling capability, we select it as a testbed for text-aware reward comparison. On Figure 5: we place the baseline reward provided by the ground truth function and the TADPoLe reward on separate plots because they have different scales and are not directly comparable. However, we show that side-by-side, the trends across the two rewards are positively correlated, which is desirable. For TADPoLe reward plots across novel prompts, which do not have associated ground truth reward functions, additional examples can be found in Figure A11 of the Appendix.

On Metaworld comparisons: Most prior methods which report performance on MetaWorld rely on (often expert-produced) video demonstrations from a similar domain or the target environment directly. We therefore focus our comparison on other methods that utilize a large-scale pretrained model for zero-shot (no in-domain demonstrations) text-conditioned policy learning, namely VLM-RM. Other work (namely Diffusion-Reward) utilize in-domain diffusion models trained explicitly on expert data and do not explore text-conditioned policy learning, which is the focus of our work. Furthermore, most performant state-of-the-art methods on MetaWorld were trained directly on dense ground-truth reward functions, whereas we evaluate the ability to recreate such behavior from dense text-conditioned rewards in a zero-shot manner. For further comparison, however, we include additional LIV results, as well as results from an adapted version of the Diffusion-Reward reward computation that is conditioned on natural language rather than historical frames, and uses a general text-to-video diffusion model in Table A8 of the rebuttal Appendix page. TADPoLe maintains superior performance in aggregate across the task suite.

On real-world experiments: we agree that real-world experimentation would be nice to have. However, our focus is on exploring text-to-policy capabilities through diffusion models, which is a novel approach. We demonstrate our approach not only on the difficult Humanoid and Dog environments, which have high action spaces and complex transition dynamics, but also on Metaworld. In reference, prior work (VLM-RM) that used vision-language model supervision also did not apply it to real-world environments. However, in terms of future work, we do agree with the reviewer that deploying on real-world robotics is an exciting direction for further exploration. On Computational Cost and Speed: the computational cost of TADPoLe is simply one denoising forward step of a pretrained diffusion model to generate each dense reward. In comparison, Diffusion-Reward requires multiple denoising steps to generate one singular reward (in practice, 10 steps are used). TADPoLe is therefore computationally cheaper and faster, when the same diffusion model is utilized, in terms of reward computation.

On naturalness via user surveys: Evaluations on naturalness were performed through user surveys (25 random participants), using videos achieved by policies at the last step of training. These are provided on the submission website, for reviewers to visualize.

On the effect of prompts on rewards: We refer the reviewer to Figure (4), where we find that different lengths and slight adjustments to the prompt during training indeed impact the final performance of the policy. To begin with, we extend a prompt in length from “a person standing” to “a person standing with hands above head”, and observe that the longer description can be respected. Furthermore, when we perform slight adjustments to the prompt, with “a person standing with hands on hips”, we confirm that the policy learned is able to respect subtle details such as placement of the arms. We anticipate extensive prompt tuning to produce further benefits; but for the scope of our work (demonstrating a rich text-conditioned reward from large-scale pretrained diffusion models) we leave tuning the most performant prompt for a specific desired task to the enthusiast.