Lightweight Language-driven Grasp Detection using Conditional Consistency Model (Part 2)

#1. Lightweight Language-driven Grasp Detection

#Overview

Given an input RGB image and a text prompt describing the object of interest, we aim to detect the grasping pose on the image that best matches the text prompt input. We follow the popular rectangle grasp convention widely used in previous work to define the grasp.

In the diffusion model, we represent the target grasp pose as $\mathbf{x}_0$. The objective of our diffusion process of language-driven grasp detection involves denoising from a noisy state $\mathbf{x}_T$ to the original grasp pose $\mathbf{x}_0$, conditioned on the input image and grasp instruction represented by $y$. The forward process in traditional conditional diffusion model is defined as:

where the hyperparameter βₜ is the amount of noise added at diffusion step t ∈ [0,T] ⊆ ℝ.

To train a diffusion model with condition y, we use a neural network to learn the reverse process:

In our approach, we utilize the diffusion process in the continuous domain, where $\mathbf{x}_t$ is the grasp pose state at arbitrary time index $t$. Unlike popular discrete diffusion models, by using a continuous space, we can improve sample quality and reduce inference times due to the ability to traverse the diffusion process at arbitrary timesteps, allowing for more fine-grained control over the denoising process.

Figure 1: The overview of our method. First, the input RGB image and text prompt are fed into the feature encoder and ALBEF fusion. Subsequently, we concurrently train two models with the same architectures: A score network to estimate the probability flow Ordinary Differential Equation (ODE) trajectory for the diffusion process and a conditional consistency model to determine the grasp pose with a few denoising steps.

#Conditional Consistency Model for LLGD

To reduce the inference time during the denoising step of the diffusion model, we aim to estimate the original grasp pose with just a few denoising steps. Since our language-driven grasp detection task has the condition $y$, we introduce a conditional consistency model based on the consistency concept to infer the original grasp pose during the inference process directly:

where $\mathbf{f}_\theta(\mathbf{x}_\epsilon, t, y) = \mathbf{x}_\epsilon$ is the boundary condition, and $\mathbf{F}_\theta(\mathbf{x}_t,t,y)$ is a free-form deep neural network whose output has the same dimensionality as $\mathbf{x}_t$.

To train our conditional consistency model, we employ knowledge distillation from a continuous diffusion process:

where $\gamma_t$ is a non-negative function referred to as the noise schedule, and $\mathbf{w}_t$ is the standard Brownian motion. This forward process creates a trajectory of grasp poses $\{\mathbf{x}_t\}_{t=0}^T$. The grasp pose state $\mathbf{x}_t$ depends on the time index $t$ and the input image and text prompt. The grasp distribution $p(\mathbf{x}_0|y)$ from the dataset is transformed into $p(\mathbf{x}_T|y) \sim \mathcal{N}(0, \mathbf{I})$. Given the ground truth grasp pose $\mathbf{x}_0$, we can sample $\mathbf{x}_t$ at arbitrary $t$:

where

The equation (4) is a probability flow ODE. With the conditional variable $y$, it can be redefined as:

where $\nabla\log p(\mathbf{x}_t|y)$ is the score function of the conditional diffusion model.

Suppose that we have a neural network $\mathbf{s}_\phi(\mathbf{x}_t, t, y)$ that can approximate the score function $\nabla\log p(\mathbf{x}_t|y)$, i.e., $\mathbf{s}_\phi(\mathbf{x}_t, t, y) \approx \nabla\log p(\mathbf{x}_t|y)$. After training the score network, we can replace the $\nabla\log p(\mathbf{x}_t|y)$ term in the equation (6) with a neural network:

Score Function Loss. In order to approximate the score function $\nabla\log p(\mathbf{x}_t|y)$, the conditional denoising estimator minimizes the following objective:

where $\lambda(t) \in \mathbb{R}^+$ is a positive weighting function.

Proposition 1. Suppose that $\mathbf{x}_t$ is conditionally independent of $y$ given $\mathbf{x}_0$, then minimizing $\mathcal{L}_{\rm score}$ is the same as minimizing:

Proof. Because $\mathbf{x}_t$ is conditionally independent of $y$ given $\mathbf{x}_0$, we have:

where

If $y$ and $t$ are fixed, we can define a transition probability that does not depend on these variables, $q(\mathbf{x}_0) = p(\mathbf{x}_0|y)$, $\kappa(\mathbf{x}_t)=\mathbf{s}_\phi(\mathbf{x}_t,t,y)$. According to Vincent P., 2011, we have:

From the equations (9) and (10), we can prove the equivalence of the two objective functions.

Discretization. Consider discretizing the time horizon $[\epsilon,T]$ into $N-1$ with boundary $t_1=\epsilon<t_2<t_3<\ldots<t_{N}=T$. If $N$ is sufficiently large, we can use an ODE-solver to estimate the next discretization step:

Conditional Consistency Model Loss. To enable fast sampling, we expect that the predicted point $\hat{\mathbf{x}}_{t_i}$ and $\mathbf{x}_{t_{i+1}}$ to lie on the same probability flow ODE trajectory. We propose conditional consistency loss to enforce this constraint:

where $\hat{\mathbf{x}}_{t_i}$ is calculated in Equation 12, $\mathbf{x}_{t_{i+1}}$ is sampling from Gaussian distribution in Equation 5, and $\theta$ is the parameters of neural network $\mathbf{f}$.

Additionally, we need to minimize the discrepancy between the predicted and ground truth grasp poses with the detection loss:

The overall training objective for our method is:

#Network Details

The input of our network is the image and a corresponding grasping text prompt represented as $e$ (for example, "grasp the fork at its handle"). We first extract the image feature using a 12-layer vision transformer ViT image encoder. The input text prompt is encoded by a text encoder using BERT or CLIP. We then combine and learn the features of the input text prompt and input image using the ALBEF fusion network. The output of the fusion features is fed into a score network, and our conditional consistency model is used to learn the grasp pose. Figure 1 shows the detail of our network.

Score Network. In practice, we utilize a score network composed of several MLP layers to extract three components: the noisy grasp pose $\mathbf{x}_t$, the time index $t$, and the conditional vision-language embedding $y$. Subsequently, these features are concatenated, and the score function is extracted through a final MLP layer. It is crucial to ensure that the output dimension of the scoring network is identical to the dimension of the input $\mathbf{x}_t$ because, fundamentally, the score function is the gradient of the grasp pose distribution given the condition $y$. Our conditional consistency model's network has an architecture similar to the scoring network; however, its output is the predicted grasp pose.

Algorithm 1: Inference Process

Input: Image and text prompt, conditional consistency model $\mathbf{f}_\theta(\mathbf{x},t,y)$, number of inference steps $P$, sequence of time points $t_1 = \epsilon < t_2 < t_3 < \dots < t_{P} = T$, noise scheduler $\alpha_t = e^{\rho_t}$.

$y \gets \text{ALBEF (image, prompt)}$

Initial grasp noise $\mathbf{x}_T \sim \mathcal{N}(0,\mathbf{I})$

$\mathbf{x}_0 \gets \mathbf{f}_\theta(\mathbf{x}_T,T,y)$

For $i = P - 1$ to $2$:

• Sample $\mathbf{z} \sim \mathcal{N}(0,\mathbf{I})$
• $\mathbf{x}_{t_i} \gets \sqrt{\alpha_{t_i}}\mathbf{x}_0 + \sqrt{1 - \alpha_{t_i}}\mathbf{z}$
• $\mathbf{x}_0 \gets \mathbf{f}_\theta(\mathbf{x}_{t_i},t_i,y)$

Output: Final grasp pose $\mathbf{x}_0$

#Training and Inference

During training, we freeze the text and image encoder, then train the ALBEF fusion, the scoring network, and the consistency model end-to-end. The score network and the conditional consistency model share the same architecture. We trained both models simultaneously for 1000 epochs with a batch size of 8 using the Adam optimizer. The training time takes approximately three days on an NVIDIA A100 GPU. Regarding the parameters of the conditional consistency model, we empirically set $T = 1000$, $\epsilon = 1$, and $N = 2000$. After training the scoring network and the conditional consistency model $\mathbf{f}_\theta(\mathbf{x}_t,t,y)$, we can sample the grasp pose given the input image and language instruction prompt in a few denoising steps using our algorithm 1.

Figure 2: Robot Hands with different ultilities.

#Next

In the next post, we will evaluate the effectiveness of our proposal.