Data-Driven Stochastic Motion Evaluation and Optimization with Image
by Spatially-Aligned Temporal Encoding

It is not easy to represent a complex distribution in a high-dimensional space by EBM using sampling such as MCMC. To tackle this problem, in Score matching [8, 9] where no sampling is needed, given the gradient of the probability density of a real sample [8], EBM is trained so that the gradient estimated by EBM gets close to the given real gradient. In Noise Contrastive Estimation [10, 11], given a simple distribution where sampling is easy, EBM is trained to discriminate between this simple distribution and the distribution of real samples. Even with these training methods [8, 9, 10, 11], however, EBM inference still need to estimate the complex high-dimensional distribution by sampling.

Another solution is to reduce the sampling difficulty both in the training and inference stages by using other generative models such as VAE [7, 12] and NF [13, 14]. In [7] and [12], sampling is achieved in the low-dimensional space induced by VAE. In [14], initial inaccurate samples are produced fast by NF, and then these initial samples are given to EBM for more accurate sampling. Even with these methods, however, accurate sampling is still difficult to predict the accurate high-dimensional long-term motions.

Difficulties in our motion prediction problem include (i) long-term error accumulation, (ii) stochasticity of motions, and (iii) multi-domain data processing.

Even with recent LSTM [15, 16, 17, 18] and GRU [19, 20], it is difficult to predict a long-time complex motion due to gradient vanishing. This problem is alleviated by Transformer [21, 22, 23] using an attention mechanism [24], which keeps the gradient stable [25].

Stochastic motions can be predicted by probabilistic generative models such as VAE [26, 27, 2] and GAN [28, 29, 30]. However, these models do not explicitly estimate the probability density of a motion. This disadvantage prevents us from optimizing the motion using maximum likelihood estimation using the probability density as the likelihood. This problem is resolved by enabling probabilistic models to explicitly estimate the probability density, for example, NF [31, 32] and EBM [7, 33]. However, NF and EBM still have several problems as described in Sec. II-A.

For multi-domain data processing such as our method using image and motion data, one of the simplest ways is to concatenate these data. The expressiveness of the concatenated data can be improved by feature extraction such as convolution and pooling for the image data [34, 35, 36]. However, spatial pooling makes is difficult to precisely localize the motion in the image [37]. Inconsistency between the coordinate systems of the image and motion data is also problematic. For the former problem, we propose efficient high-dimensional feature extraction that focuses on the motion trajectory on the image. Since the latter problem can be resolved by coordinate rearrangement (e.g., motion heatmapping in the image [2]), our method also employs it.

An image and a motion sequence are denoted by $I$ and $P$, respectively. $P$ is a set of $(T+1)$ temporal pose vectors $\bm{p}_{t}$ where $t\in\{0,\cdots,T\}$. In our experiments, $\bm{p}_{t}$ is a 11-D vector consisting of the 3D location of the robot hand ($\bm{x}_{t}$), its orientation ($\bm{r}_{t}$), its status ($s_{t}$), and a time stamp $t$. $\bm{x}_{t}$, $\bm{r}_{t}$, and $s_{t}$ are a 3D positional vector in the camera coordinate system, a 6D vector representing the 3D orientation [38], and a scalar value between 0 and 1 representing the open (1) and close (0) status of the hand gripper, respectively. $I^{i}$ and $P^{i}$ denote the $i$-th data, where $i\in\{1,\cdots,N\}$, in our training dataset. In each of all pairs in the training dataset, $P^{i}$ can accomplish the task from $I^{i}$.

Motion prediction from the initial state represented by the image is illustrated in Fig. 2. In our method, $I$ and $P$ are fed into EBM (denoted by $E_{\theta}$) to estimate the probability density $p(P|I)$ of the consistency between $I$ and $P$:

$$p(P|I)=\frac{\exp(-E_{\theta}(I,P))}{Z(I)},$$

(1)

where $\theta$ denotes the parameters in EBM. $Z$, which depends only on $I$, is defined as follows for normalization:

$$Z(I)=\int\exp(-E_{\theta}(I,P))dP$$

(2)

In the inference stage, given $I^{test}$ and any motion $P$, EBM is required to estimate the consistency between $I^{test}$ and $P$. While $Z(I)$ in Eq. (1) is computationally intractable, $Z(I)$ is not changed because $I$ is fixed to $I^{test}$. Instead of the probability density, therefore, the energy term $-E_{\theta}(I^{test},P)$ is used as the output of EBM.

To increase $-E_{\theta}(I^{test},P)$, DMO updates $P$ in order to bring $P$ to its ideal motion (denoted by $P^{test}$ paired with $I^{test}$), which is not available in the inference stage, as follows:

$$\hat{P}_{(1)}=DMO(I,P),$$

(3)

where $I^{test}$ is substituted to $I$ in the inference stage. Since it is not guaranteed that $\hat{P}_{(1)}$ is sufficiently close to $P^{test}$, DMO is recurrently used by rewriting Eq. (3) to Eq. (4):

$$\hat{P}_{(r+1)}=DMO(I,\hat{P}_{(r)}),$$

(4)

where the initial motion $P$ in Eq. (3) is denoted by $\hat{P}_{0}$, and $r$ is an integer greater than or equal to zero.

As initial motions each of which is fed into DMO in Eq. (4), $\{P^{1},\cdots,P^{N}\}$ in the training dataset are reused in our method so that $\hat{P}^{i}_{(r+1)}=DMO(I^{test},\hat{P}^{i}_{(r)})$. Among all $\{P^{1},\cdots,P^{N}\}$, several of them are close to $P^{test}$, while the remaining ones are far from $P^{test}$. For efficiency, only the former motions are used. These motions are determined by selecting the top $n$ energy values computed by EBM, and fed into DMO expressed by Eq. (4).

In our method, both EBM and DMO are implemented as neural networks trained as described in Secs. III-C and III-D, respectively.

If previous EBM training methods [33] are used for estimating $E_{\theta}$, $\theta$ is optimized by maximizing the log-likelihood of the following loss function $\mathcal{L}_{EBM}(\theta)$:

$$\mathcal{L}_{EBM}(\theta)=\frac{1}{N}\sum_{i=1}^{N}\left(-E_{\theta}(I^{i},P^{i})-\log Z_{\theta}(I^{i})\right)$$

(5)

The gradient of $\mathcal{L}_{EBM}(\theta)$ is expressed as follows:

$$\begin{split}\nabla\mathcal{L}_{EBM}(\theta)=&\frac{1}{N}\sum_{i=1}^{N}\left(-\nabla E_{\theta}(I^{i},P^{i})\right.\\ &\left.+{\mathbb{E}}_{\widecheck{P}^{i}\sim p(P|I^{i})}\nabla E_{\theta}(I^{i},\widecheck{P}^{i})\right),\end{split}$$

(6)

where $\widecheck{P}^{i}$ is a motion synthesized by sampling based on $p(P|I^{i})$ given by EBM. This sampling process is not easy in high dimension, as described in Sec. II-A.

To consider how to avoid the difficulty in this sampling, we focus on Eq. (6). With $-\nabla E_{\theta}(I^{i},P^{i})$ in Eq. (6), $p(P^{i}|I^{i})$ as the probability density of a positive sample where $P^{i}$ accomplishes the task from $I^{i}$ is increased. On the other hand, $\nabla E_{\theta}(I^{i},\widecheck{P}^{i})$ decreases $p(\widecheck{P}^{i}|I^{i})$ as the probability density of a negative sample where $\widecheck{P}^{i}$ cannot achieve the task from $I^{i}$. In our method, these positive and negative samples are selected from the training dataset as follows. As described in Sec. III-B, $\{P^{1},\cdots,P^{N}\}$ in the training dataset are reused as initial motions fed into DMO in the inference stage. To make this inference succeed, $P^{i}$ is expected to be updated towards $P^{test}$, if $I^{i}$ is the closest to $I^{test}$. On the other hand, $P^{j}$ where $j\neq i$ is expected not to be updated to $P^{i}$. These relationships are represented in EBM so that $(I^{i},P^{i})$ and $(I^{i},P^{j})$ are trained as the positive and negative samples, respectively.

In addition to the positive and negative samples directly borrowed from the training dataset, our method augments samples for training EBM so that $P$ close to $P^{j}$ should be also a negative sample. This data augmentation is done by simple low-dimensional sampling via VAE. VAE trains an embedding space from a set of all $N$ motion sequences (i.e., $P^{1},\cdots,P^{N}$) in the training dataset. The trained VAE is denoted by $V$. $P^{j}$ is fed into $V$ to get $\mathcal{N}(\bm{\mu}^{j},\bm{\sigma}^{j})$, as shown in the upper part of Fig. 3. Let $\bm{z}^{j}$ be sampled from $\mathcal{N}(\bm{\mu}^{j},\bm{\sigma}^{j})$. The motion decoded from $\bm{z}^{j}$ (denoted by $\tilde{P}^{j}$) is close to its original motion $P^{j}$ Since it is easy to represent this simple Gaussian distribution in a low-dimensional embedding space, sampling from this distribution in $V$ allows us to stably synthesize $\tilde{P}^{j}$ close to $P^{j}$.

As shown in the lower part of Fig. 3, with the positive and negative samples prepared by the aforementioned manner, EBM is trained with the following loss function:

$$\begin{split}\nabla\mathcal{L}_{EBM}(\theta)=&\frac{1}{N}\sum_{i=1}^{N}\left(-\nabla E_{\theta}(I^{i},P^{i})+\sum_{j\neq i}\nabla E_{\theta}(I^{i},P^{j})\right.\\ &\left.+\sum_{j=1}^{N}{\mathbb{E}}_{\tilde{P}^{j}\sim\mathcal{N}(\bm{\mu}^{j},\bm{\sigma}^{j})}\nabla E_{\theta}(I^{i},\tilde{P}^{j})\right)\end{split}$$

(7)

Different from general motion optimizers such as MCMC sampling [39], DMO alleviates the difficulty of hyper-parameter tuning by being trained in a data-driven manner as follows. To train DMO, $(I^{i},P^{i})$ in which the task can be accomplished is given as the ground-truth positive sample in the training dataset. DMO is trained so that any $P$ close to $P^{i}$ is updated toward $P^{i}$. In reality, however, $P$ is not randomly distributed but should be limited within realistic robot motions (e.g., within range of articulated motion).

In our method, a set of such realistic $P$ around $P^{i}$ is synthesized via the trained VAE, $V$, as done in the training of EBM. This training procedure is illustrated in Fig. 4. $P^{i}$ is fed into $V$ to get the motion sampled from $V$. This sampled motion (denoted by $\tilde{P}^{i}$) is close to its original motion $P^{i}$, as described in Sec. III-C.

By substituting $(I^{i},\tilde{P}^{i})$ to $(I,\hat{P}_{(r)})$ where $r=0$ in Eq. (4), DMO gets the updated motion $\hat{P}^{i}_{(1)}$. As with the inference stage, DMO expressed by Eq. (4) is recurrently used to get $\{\hat{P}^{i}_{(2)},\cdots,\hat{P}^{i}_{(R+1)}\}$.

With $\{\hat{P}^{i}_{(1)},\cdots,\hat{P}^{i}_{(R+1)}\}$, the parameters of DMO denoted as $\omega$ are optimized with the following loss in which the difference between $\hat{P}^{i}_{(r)}$ and its ideal motion $P^{i}$ is minimized:

$$\mathcal{L}_{DMO}(\omega)=\frac{1}{N}\sum_{i=1}^{N}\sum_{r=1}^{R+1}||P^{i}-\hat{P}^{i}_{(r)}||_{2},$$

(8)

where $R$ denotes the number of recurrently-updated motions. In $\mathcal{L}_{DMO}(\omega)$, motion data augmentation is achieved by training all $(R+1)$ motions (i.e., $\hat{P}^{i}_{(1)},\cdots,\hat{P}^{i}_{(R+1)}$) instead of training only $\hat{P}^{i}_{(R+1)}$.

We evaluate our method using the RLBench dataset [42] in which a robot performs various tasks on the simulator, Pyrep [43]. The nine tasks used in our experiments are shown in Fig. 6. For each task, 1,000 and 100 training and test sequences were generated in the simulator. In all the sequences, the Franka Panda robot was used. The initial state was captured as an image in each sequence. Each image consists of RGB channels and a depth channel. The image size is $256\times 256$ pixels. As VAE, ACTOR [44] was used.

Our proposed motion optimizer (DMO) is evaluated with other motion optimizers so that each of the following optimizers is substituted for “Motion Optimizer” in Fig. 2:

The success rate of each optimizer is shown in Table I. In all tasks except RT, our methods (a) and (b) outperform the others. This verifies the superiority of the proposed combination of probabilistic and deterministic methods over the probabilistic methods (c), (d), (e), and (f). The examples of motions predicted by (a) Ours-1 and (e) VAEBM-L, which is the best among the others, are shown in Fig. 7. We can see that Ours-1 can finely rectify the predicted motions for accomplishing the tasks. In comparison between (a) Ours-5 and (b) Ours-1, Ours-5 is a bit better. This might be because of motion data augmentation by iterative motion updates.

More examples are shown in the supplementary video.

The following four variants of the feature extractor in EBM and DMO are evaluated.

Their task success rates are shown in Table II. Our spatially-aligned temporal coding (a) and (b) outperform (c) and (d), while the effectiveness of the feature concatenation done in (a) is not validated.

Table III shows the cost (i.e., the number of parameters and the computational cost) of the four variants. Table III shows only the results of EBM because the dominant part is shared by EBM and DMO. Furthermore, when multiple motions are evaluated and updated with the same image by EBM and DMO, respectively, we can reuse $F$ for reducing the cost. The results in the cases where $F$ is reused and not reused are shown in Reuse and All columns, respectively.

The costs are almost same between (a), (b), (c), and (d) because the cost in CNN is larger than the costs in Transformer and MLP. On the other hand, it can be confirmed that the costs in Reuse are much smaller than those in All. This property is important because the performance is gained as the number of input motions (i.e., $n$) increases, as demonstrated later in Table VI. If the number of input motions becomes $n$-fold, the total cost in Reuse also becomes $n$-fold.

The following two variants of the EBM training methods are evaluated, while their prediction method is fixed to be that described in Sec. III-B.

The task success rates of these two variants are shown in Table IV. The scores in (b) are significantly worse. We interpret this performance drop in what follows.

In (b), the loss is expressed by $\sum_{j\neq i}\nabla E_{\theta}(I^{i},P^{j})$ is removed from Eq. (7). This loss can be minimized not only if any $(I,P)$ can be correctly discriminated between positive and negative data but also if $\tilde{P}$ synthesized by $V$ is discriminated from $P^{i}$ independently of $I^{i}$. This latter case might be caused because the loss is decreased to almost zero, while the task success rate is low. This problem can be resolved by synthesizing training data that cannot be discriminated from $P^{i}$. While a possible way to synthesize such realistic training data is gradient-based adversarial training such as GAN [48], it is unstable [49, 50]. We also confirmed that EBM training methods using Langevin MCMC [39] and VAEBM [12], both of which exploit gradient information, are unstable. Instead of synthesizing realistic data, on the other hand, real training data is used as negative samples, as expressed in Eq. (7).

In Table VI and  VI, the performance changes with the changes in $n$ and $R$ are shown, respectively. By increasing $n$, the performance is gained in many tasks such as PC, PK, PR, PU, SW, TP, while $n$ has less impact in CB and RT. Since $n$ is the number of motions evaluated in EBM and rectified in DMO, it is natural that the task success rate can be increased. By increasing $R$, on the other hand, the performance change is fluctuating. This might be caused because a limited number of training data might be insufficient, if the task is achievable by a variety of motions. In such a case, since the training samples are sparsely distributed, overfitting is caused by increasing the number of iterative updates (i.e., $R$).