Prim-LAfD: A Framework to Learn and Adapt Primitive-Based Skills from Demonstrations for Insertion Tasks

Impedance control is used to render the robot as a mass-spring-damping system following the dynamics below,

where $\mathbf{M}$, $\mathbf{D}$, $\mathbf{K}$ are the desired mass, damping, and stiffness matrices, and $\mathbf{F}_{ext}$ denotes the external wrench. $\mathbf{\ddot{x}}_{d}$, $\mathbf{\dot{x}}_{d}$, $\mathbf{x}_{d}$ are the desired Cartesian acceleration, velocity, and pose of the end-effector, and $\mathbf{\ddot{x}}$, $\mathbf{\dot{x}}$, $\mathbf{x}$ are the current values correspondingly. We assume a small velocity in our tasks and set $\mathbf{\ddot{x}}$, $\mathbf{\dot{x}}$ to 0, thus arriving at this control law,

where $\boldsymbol{\tau}$ is the control torque, $\mathbf{F}$ is the control wrench, $\mathbf{J}(\mathbf{q})$ is the Jacobian, and $\mathbf{g}(\mathbf{q})$ is the gravity compensation force.

Throughout an insertion task, we would like to design a desired trajectory and a variable impedance to guide the robot movement. In favor of stability and ease of learning, we use a diagonal stiffness matrix $\mathbf{K}=\textbf{Diag}[K_{x},K_{y},K_{z},K_{roll},K_{pitch},K_{yaw}]$, and, for simplicity, the damping matrix $\mathbf{D}$ is scaled such that the system is critically damped.

In summary, our insertion policy output, $\mathbf{a}_{t}\in\mathcal{A}$, fed to the impedance controller defined above, is composed of a desired end-effector pose $\mathbf{x}_{d}$ and the diagonal elements of the stiffness matrix $\mathbf{k}=\{K_{x},K_{y},K_{z},K_{roll},K_{pitch},K_{yaw}\}$. The input to the policy, $\mathbf{s}_{t}\in\mathcal{S}$, consists of end-effector pose $\mathbf{x}_{t}$ and the sensed wrench $\mathbf{f}_{t}$, and is extensible to more modalities such as RGB and depth images.

In this section, we provide a detailed design on our insertion policy, which entails a state machine with state-dependent motion primitives. Each motion primitive $\mathcal{P}_{m}$ associated with the $m$-th state defines a desired trajectory $f_{\boldsymbol{\theta}_{m}}(\mathbf{x}_{enter},\mathcal{T})$, an exit condition checker $h_{\boldsymbol{\theta}_{m}}(\cdot):\mathcal{S}\rightarrow\{1,0\}$, and a 6-dimensional stiffness vector $\mathbf{k}_{m}$. $\boldsymbol{\theta}_{m}$ contains all the learnable parameters in the primitive $\mathcal{P}_{m}$. $\mathbf{x}_{enter}$ denotes the end-effector pose upon entering the $m$-th state. $\mathcal{T}$ contains the task information such as the 6 DoF poses of the peg and the hole; often, the hole pose defines the task frame of the motion primitives.

In the following, we formally describe the $4$ motion primitives used in the peg-in-hole tasks, as shown in Figure 2.

In this section, we illustrate how to learn the primitive parameters $\boldsymbol{\Theta}=\{\boldsymbol{\theta}_{1},\boldsymbol{\theta}_{2},\boldsymbol{\theta}_{3},\boldsymbol{\theta}_{4}\}$. The core idea is using a black-box optimizer to optimize a task-relevant objective function $J(\cdot)$. While a similar idea has been explored in Johannsmeier et al. (2019), the objective function used is simply the measured task execution time. The major drawback of this objective function is that the objective signal is sparse and can only be triggered when the task is successfully executed. This makes the optimizer challenging to find a feasible region initially, especially when the primitive parameter space is large. Motivated by this, we propose a dense objective function that measures the likelihood of the induced execution trajectory being sampled from the distribution of successful task demonstrations $\mathcal{E}_{D}=\{\boldsymbol{\xi_{i}}\}(i=1,2,...,M)$. Assuming the trajectroies are Markovian, a trajectory rollout $\boldsymbol{\xi}=[\boldsymbol{x}_{0},\boldsymbol{x}_{1},...,\boldsymbol{x}_{n-1}]$ is modeled as:

In order to learn $p(\boldsymbol{x}_{i}|\boldsymbol{x}_{i-1})$ from demonstrations, we first use a Guassian Mixture Model (GMM) to model the joint probability as $p(\left[\begin{matrix}\boldsymbol{x}_{i}\\ \boldsymbol{x}_{i-1}\end{matrix}\right])=\sum_{j=1}^{K}{\phi_{j}\mathcal{N}(\boldsymbol{\mu}_{j},\boldsymbol{\Sigma}_{j})}$ , where $\sum_{j=1}^{K}{\phi_{j}}=1$, and $K$ is the number of GMM clusters.

We further represent the Gaussian mean $\boldsymbol{\mu}_{j}$ and variance $\boldsymbol{\Sigma}_{j}$ as: $\boldsymbol{\mu}_{{j}}=\begin{bmatrix}\boldsymbol{\mu}^{{1}}_{{j}}\\ \boldsymbol{\mu}^{{2}}_{{j}}\end{bmatrix}$, $\boldsymbol{\Sigma}_{{j}}=\begin{bmatrix}\boldsymbol{\Sigma}^{{11}}_{{j}}&\boldsymbol{\Sigma}^{{12}}_{{j}}\\ \boldsymbol{\Sigma}^{{21}}_{{j}}&\boldsymbol{\Sigma}^{{22}}_{{j}}\end{bmatrix}$. We can then derive the conditional probability $p(\boldsymbol{x}_{i}|\boldsymbol{x}_{i-1})=\sum_{j=1}^{K}{\phi_{j}\mathcal{N}(\overline{\boldsymbol{\mu}}_{j},\overline{\boldsymbol{\Sigma}}_{j})}$, where

Then, the objective function is designed as $J(\boldsymbol{\xi})=\log{p(\boldsymbol{\xi};\boldsymbol{\Theta})}+B$, where the first term encourages exploring parameters inducing similar trajectories to the successful demonstration traces, and the second term $B$ denotes a sparse bonus reward if the task succeeds. We use black-box optimizers to solve $\boldsymbol{\Theta}^{*}=\underset{\boldsymbol{\Theta}}{\mathrm{argmax}}J(\boldsymbol{\Theta})$, and Bayesian Optimization (BO) is selected in our work. Expected Improvement (EI) is used as the acquisition function, and we run BO for $N$ iterations. The learned parameter $\boldsymbol{\Theta}^{*}$ that achieves maximum $J(\boldsymbol{\Theta})$ during $N$ training iterations is selected as the optimal primitive configuration. Note that BO can be seamlessly replaced by other black-box optimization methods and the optimizer choice is not the focus of this work.

In this section, we detail our method on how to leverage prior experience when adapting to a novel insertion task, in particular, how to adapt previously learned peg-in-hole policies to different tasks with unseen hole shapes. Our adaptation procedure is composed of two core steps: measuring task similarities and transferring similar task policies to the unseen shape.

As shown in Figure 1, our hardware setup consists of a 6-DoF FANUC Mate 200iD robot and an ATI Mini45 Force/Torque sensor. The clearances for all 3D printed peg-hole pairs are 1mm; the waterproof and RJ45 are unaltered off-the-shelf connectors. To mimic the pose estimation error during industrial deployments, a uniform perturbation error of +/- 5mm in translation and +/- 6 degree in orientation is applied along each dimension. The controller takes the policy output at 10 Hz and computes the torque command streamed to the robot at 1000Hz. All the learnable parameters of the policy and their initial range are listed in Table 1. Two metrics are used to evaluate the effectiveness and efficiency of the approaches: 1) the number of iterations the robot takes to accomplish the first successful insertion during training (denoted as number of iterations), and 2) the success rate of the optimal policy after a fixed number of iterations (denoted as success rate).

In this experiment, we applied our learned objective function from demonstration (LfD) for primitive parameter optimization as described in Section 3.3, and compare the results against primitive optimization by minimizing the measured task execution time (Time) Johannsmeier et al. (2019). When learning the objective function, we collected $10$ demonstrations for each insertion task through kinesthetic teaching. The demonstrations have $43.5$ time steps on average across different insertion tasks. The number of GMM clusters is set as $K=25$.

For each of the 8 peg-hole pairs, we run BO for 40 iterations. Within each iteration, the current policy is executed twice with independently sampled hole poses. The average objective of the two trials is used as the final objective value consumed by the BO step. The optimal policy is selected as the policy at the BO step achieving the optimal objective value, and evaluated by being executed 20 trials with independently sampled hole poses. As shown in Figure 4, LfD outperforms Time on almost all the insertion tasks in both number of iterations and success rate. In some tasks, e.g., parallelogram and RJ45, Time cannot find proper parameters to achieve a single successful trial within 40 iterations, while LfD successfully navigates through the parameter space and accomplishes successful trials for all of the tasks. This validates the learned dense objective function, provides richer information for primitive learning than sparse signals alone like task success and completion time.

We now examine how our generalization method described in Section 3.4 performs when transferring to unseen shapes. We consider the leave-one-out cross-validation setting, i.e., when presented 1 of the 8 tasks as the task of interest, the learning algorithm has access to all interaction data during policy learning on the other 7 tasks. Two sets of experiments are conducted. First, we apply the full method described in Section 3.4 (Full) to learn a reduced search space using the $L=3$ most similar tasks, within which the primitive parameters are optimized.

Compared with LfD where parameters are optimized over the full space, Full reached a comparable or better success rate, meanwhile achieved the first task completion with a lower number of iterations. Second, we consider learning the search space without measuring the task similarities (NoSim), i.e., the new search space is obtained using all the other tasks instead of only the similar ones. As seen in Figure 4, Full outperforms NoSim consistently on all insertion tasks, indicating the significance of finding similar tasks using the task geometry information before learning the new search space.