Text-Aware Diffusion  for Policy Learning

In this work, we propose Text-Aware Diffusion for Policy Learning (TADPoLe), a framework that utilizes a large-scale pretrained, frozen text-conditioned diffusion model to generate dense reward signals for policy learning. We demonstrate that TADPoLe enables the learning of novel, zero-shot policies that are flexibly and accurately conditioned on natural language inputs, across multiple robot configurations and environments, for both goal-achievement and continuous locomotion tasks.  Furthermore, we observe that behaviors learned by TADPoLe are qualitatively appealing due to alignment with natural priors summarized from large-scale pretraining.

Experiments

We perform comprehensive experiments to demonstrate the effectiveness of TADPoLe on goal achievement, continuous locomotion, and robotic manipulation tasks. For evaluation benchmarks, we select the Dog and Humanoid environments from the DeepMind Control Suite, both of which are known to be advanced in difficulty due to their large action space and complex transition dynamics. Additionally, we select 8 robotic manipulation tasks from Meta-World, which are balanced for diversity and complexity.

Goal-Achievement

We showcase text-conditioned goal-reaching tasks learned via TADPoLe, for both Dog and Humanoid environments, and compare it against other text-to-reward approaches.  TADPoLe is able to successfully learn a variety of behaviors, from standing upright to doing splits to kneeling.  Noticeably, the results on the Dog experiments appear more aligned with how dogs naturally "stand".

Prompt Sensitivity

We further investigate whether or not TADPoLe is sensitive to subtle variations of the input prompt.  When we change the conditioning phrase from "a person standing" to "a person standing with hands above head" and train an agent as before, we find that the resulting policy indeed respects the additional detail in the text specification and keeps its hands above its head.  We demonstrate that a similar behavior for the prompt of "a person standing with hands on hips", and show that the agent learns to keep both its hands on its hips.   We take this as further evidence that TADPoLe is a powerful approach to learning learning text-conditioned policies, capable of respecting fine-grained details and subtleties of the input prompts.  We also include an additional example of a Dog learning to chase its tail. 

Continuous Locomotion

We further explore the ability of Video-TADPoLe to learn continuous locomotion behaviors conditioned on natural language specifications.  This is challenging for approaches that statically select a canonical goal-frame to achieve, such as CLIP or LIV, and we propose Video-TADPoLe, which leverages large-scale pretrained text-to-video generative models, as a promising direction forward.  Indeed, Video-TADPoLe outperforms ViCLIP-RM, our proposed extension of VLM-RM into text-video alignment, in both quantitative and qualitative evaluations across both Dog and Humanoid.

Robotic Manipulation

We investigate how well TADPoLe can be applied to learn robotic manipulation tasks through dense text-conditioned feedback.  Whereas Meta-World normally provides manually designed dense rewards guiding the agent to complete the task, we demonstrate TADPoLe's promising ability by replacing the ground-truth dense reward signal with TADPoLe's text-conditioned reward. We perform thorough comparisons between TADPoLe and VLM-RM by evaluating them on a diverse set of selected Meta-World tasks, in which we observe that TADPoLe provides meaningful zero-shot dense supervision that enables success across a variety of robotic manipulation tasks through text-conditioning. 

User Study

We perform an additional larger paid user study, with 25 anonymous participants sourced from the general public, through the Prolific platform.  In determining which motion was believed to be more natural between a diffusion-based supervisor against a CLIP-based supervisor, the majority of users preferred the policies learned by TADPoLe or Video-TADPoLe in all cases.  We take this as further evidence that our framework's usage of a pre-trained generative model as a reward signal also supervises the policy to behave in a visually natural way.

Failure Modes

We report failure modes of Video-TADPole and TADPoLe.  We discover that for certain seeds, Video-TADPoLe is able to learn a Humanoid agent that walks according to the direction specified by the text prompt; however, for others, Video-TADPoLe may learn to walk in the opposite direction (visualized below).  This demonstrates that although the agent consistently learns to walk, it may not always respect details such as direction.  How to provide fine-grained control over the text-conditioning to focus on key words such as direction is interesting to explore for future work.

Visualizing TADPoLe Rewards for Natural Videos

We visualize Video-TADPoLe dense text-conditioned rewards for natural videos, to demonstrate how it can assign aligned prompts a higher reward than less-aligned prompts, thus supporting its potential as a text-conditioned reward signal for policy learning in real-world settings.  We utilize natural videos of robot arms performing tasks from the Bridge Dataset (version 2), and demonstrate that our reward computation assigns higher dense rewards per frame for a more-aligned text prompt describing the action.  We also use natural videos of humans performing repetitive exercises from the RepNet dataset, which we do not visualize explicitly due to potential anonymization concerns (we can provide exact videos and details at the reviewers' request, with permission from the AC/PCs).  We similarly find that the Video-TADPoLe dense reward exhibits preference for more-aligned text prompts relating to the natural video.

Additional Robotic Experiments

We implement TADPoLe for additional robotic tasks from Adroit, and FrankaKitchen.  For Adroit, we implement TADPoLe with a DrM RL Backbone, but otherwise keep all hyperparameters the same as what was reported in the MetaWorld experiments.  For FrankaKitchen we keep the exact same implementation as what was used for MetaWorld.

We observe that TADPoLe is able to replace dense ground-truth rewards with a purely text-conditioned reward, to still accomplish  complex dextrous manipulation tasks.  We reiterate that TADPoLe accomplishes this using a general-purpose, large-scale pretrained text-to-image diffusion model that has not (to our knowledge) observed any in-domain examples from Adroit.

We also report an example of a task that TADPoLe is able to learn in FrankaKitchen - FrankaKitchen has been known to struggle in online RL settings without observing prior demonstrations.  Furthermore, LIV, a prior work, has found that fine-tuning for Franka-Kitchen is needed for language-grounding, and a text-conditioned reward is virtually meaningless without fine-tuning (Appendix G.4, and Figure 20 vs. Figure 21) .  As TADPoLe does not use fine-tuning, but still attempts language-conditioned policy learning, we expect performance on FrankaKitchen to be challenging.