Proleptic Temporal Ensemble for Improving the Speed of Robot Tasks Generated by Imitation Learning

2.1 Imitation Learning for Dual-Arm Manipulation

Imitation learning (behavior cloning) trains a model to mimic expert behavior by learning from demonstration data that captures actions based on observations. When a demonstrator uses a teleoperation device or other control mechanism to repetitively perform specific tasks with a robot, observations such as images and joint angles, and action data such as desired joint angles are collected during this process. During the training process, the model receives a dataset of observation-action pairs and learns a function that maps the current observation to the demonstrator’s future actions.

There are several methods for collecting demonstrator action data. One common approach involves the demonstrator teleoperating the robot using a master arms^[6], while another method uses dedicated data collection devices^[7]. Additionally, research has explored the use of wearable master arms equipped with cameras that recognize operator’s hand movements^[8], or directly collecting data by having the demonstrator wear data gloves and body camera designed for data acquisition^[9]. [Fig. 1] illustrates these data collection systems. Cameras used for collecting environmental data are typically mounted on the robot’s head or attached to the wrist to observe object manipulation. The speed at which the demonstrator performs tasks during data collection directly influences the actions generated by the learned policy.

[Fig. 1]

Various Data Collection Methods: (a) Remote control of the master device, (b) Portable gripper, (c) Wearable device and hand recognition sensor, (d) Data glove

Various studies have demonstrated successful applications of imitation learning in object manipulation tasks, such as cooking^[10] or tying shoelaces^[11]. Notable research includes works using Diffusion Policy and Action Chunking with Transformers (ACT), which utilizes Conditional Variational Autoencoders (CVAE)^[12]. Both the Diffusion Policy and ACT focus on predicting future action sequences rather than simple immediate actions, operating in a manner similar to Model Predictive Control (MPC), which is called action chunking. Additionally, introducing generative AI techniques such as diffusion models or CVAE model enables the imitation learning model multimodal behaviors of human operators.

In efforts to improve task performance and stability, C. Chi, et al.^[7] decreases the update speed of the inference by 0.5 times to slow down the execution of the generated behaviors. In these studies, data collected by humans, as shown in [Fig. 1(b)], exhibited fast speeds for the robot causing difficulty in tracking target points. However, increasing the update speed of the network inference for the faster motion of the robot comes with limitations, as it decreases the computational margin for the network inference which takes majority of the computational time due to its large model size.

2.2 Action Chunking and Ensemble

In this paper, we adopt the ACT algorithm for imitation learning. ACT fundamentally follows a CVAE structure, with the encoder and decoder using transformers. The CVAE encoder receives the robot’s action sequence and current pose as inputs and encodes the robot’s action modality into a style vector. The CVAE decoder takes the encoded style vector along with image data acquired through RGB stereo cameras and the robot’s end-effector pose sequence as inputs to output the robot’s future action sequence. The action space is also the robot’s end-effector pose. Defining the observation sequence containing the robot’s images and poses from the past time step t-k_s​ to the current time step t as $$ {\mathit{s}}_{t-k_{\mathrm{s}} : t}$$, and the action sequence from the current time step t to t+k_a​ as a $$ {\mathit{a}}_{t : t+k_{\mathrm{a}}}$$​​, the policy is expressed as $$ {\pi }_{\theta }\left({\mathit{a}}_{t : t+k_{\mathrm{a}}}|{\mathit{s}}_{t-k_{\mathrm{s}} : t}\right)$$.

[Fig. 2] illustrates the future action sequences predicted by the ACT at each step. The vector $$ {\widehat{\mathit{a}}}_{u\left|v\right.}$$ represents the estimated action at time u estimated at time v, with its magnitude defined by the degrees of freedom of the target robot. When predicting the future action sequences at each step, the action sequence predicted at instance t is as follows

[Fig. 2]

Action Chunking and ensemble

This method of predicting not only the robot’s next step action but also the action sequences up to t+k_a​ is defined as action chunking ([Fig. 1]). Action chunking was designed to address the compounding errors that arise in the robot’s intricate object manipulation tasks. Compared to traditional approaches that only infer the next step’s action, the use of the action chunking technique enhances the consistency of the robot’s behavior. In the ACT algorithm, when applying the action $$ {\widehat{\mathit{A}}}_t$$ to the robot, a weight function is used to ensemble the actions predicted at the current point from the past. First, we define the matrix B_t​, which consists of the actions inferred from the past to the present, as follows ([Fig. 1]).

Furthermore, the final input a_t​ to the robot, incorporating the weight function based on the exponential function, is defined as follows

Where the parameter m represents the slope of the weight function. A higher value of m places greater emphasis on the most recently estimated actions. A large value of m enables the robot to respond quickly to unexpected disturbances or environmental changes, resulting in high responsiveness. Conversely, a lower m generates more consistent behavior in the presence of environmental noise. Therefore, it is important to adjust the value of m to find an appropriate balance between the robot’s responsiveness and consistency. In all experiments conducted in this study, m was set to 0.05.

2.3 The speed limitations of imitation learning policies

The movements of the robot derived from imitation learning policies are entirely dependent on the collected data used during the training process. In the data collection process, the operator manipulates the robot using a teleoperation device, and various task speeds are recorded depending on the operator’s skill level. Even if only fast task data is collected from skilled operators to improve the robot’s task speed, the robot’s task speed based on the imitation learning policy is still constrained by data dependency. The next chapter introduces a proleptic temporal ensemble technique that improves the robot’s task speed during the inference process without requiring additional training.

3.1 Proleptic Temporal Ensemble

To accelerate the speed of robot movements generated by imitation learning algorithms, we propose the PTE technique. [Fig. 3] serves as an example to illustrate the proposed method, assuming that five actions are predicted at each step, and displays the actions predicted over the last six steps in blocks. The blocks marked as ‘α’ in [Fig. 3] represent the actions used at the current time, predicted from the past, and are the components defined in Equation (2) as ​B_t.

[Fig. 3]

Structure of action chunking

The proposed PTE uses B_t+f(f<k_a) instead of B_t in the calculation of input a_t as described in Equation (3). B_t+f is defined as follows

It consists of f fewer $$ {\widehat{\mathit{a}}}_{u\left|v\right.}$$ compared to B_t​. For example, if f is 2, B_t+2 is defined by the blocks marked as ‘γ’ in [Fig. 3].

As a result, the PTE expects to enhance the robot’s task speed by using a weighted average of future control inputs through the ensemble method, instead of relying on the inferred current control input as the target value at the present time.

3.2 Proleptic Variables

The proposed PTE technique sets future action points as targets for the robot, allowing it to move more quickly to follow these targets. Research^[7] has shown that reducing the action update cycle to increase task speed is not suitable for object contact tasks that require precision at high speeds. Therefore, the proleptic variable f needs to be set to an appropriate value depending on the task. Precise tasks require careful execution, while non-contact free motions are better suited for faster speeds. The next chapter validates the changes in task success rates and task speeds based on the proleptic variable f through repeated experiments.

4.1 Block Color Classification Task

To analyze the performance of the algorithm presented in this paper, we selected a block color classification task as the robot’s operation. [Fig. 4] illustrates the classification task environment. Two six-degree-of-freedom robotic arms are each equipped with a one-degree-of-freedom gripper to perform the task of picking up yellow and blue 2-stud blocks and placing them into boxes of the corresponding colors. The blocks are identical in size and shape, differing only in color, and their orientations and positions are randomly arranged. The boxes are placed at either end of the table without any fixation.

[Fig. 4]

Experimental environment

4.2 Data Collection for Learning

[Fig. 5] shows the data collection environment. To remotely control the robots, we developed a master control device with a total of seven degrees of freedom, including six degrees for the robot joints and one for the gripper. The master device was designed to have a scale size of 2:1 relative to the robotic arm, and the data structure transmitted from the master device to the robot is as follows

[Fig. 5]

Teleoperation environment

Where, $$ {\mathit{p}}_{\mathrm{r}}\in {\mathbb{R}}^{3\times 1}$$, $$ {\mathit{o}}_{\mathrm{r}}\in {\mathbb{R}}^{3\times 1}$$, $$ {\theta }_{\mathrm{r}}\in \left[\mathrm{0,\; 1}\right]$$ represent the position of the endpoint of the right arm master device, the Euler angles corresponding to its orientation, and the gripping degree of the gripper, respectively. The subscripts r and l denote the right and left arms, respectively.

Surrounding environmental data was collected using three stereo cameras. As shown in [Fig. 5], two cameras were attached to the wrist, while the remaining one was fixed at a third position. The images captured by the stereo cameras are displayed in [Fig. 6], resulting in a total of six images with a resolution of 640x480.

[Fig. 6]

Stereo camera images

The demonstrator operated the robot using a teleoperation device, repeatedly performing the task of picking up a total of six blocks and placing them into boxes of the corresponding colors. The blocks were randomly arranged in the center of the table, as illustrated in [Fig. 4], [Fig. 5], and [Fig. 7]. If the box position to which a block needed to be moved during the classification process was outside the robotic arm’s workspace, the task was carried out by transferring the block between the robots. The right arm overcomes workspace limitations by passing the blue block to the left arm, which then places it in the blue box on the left side.

[Fig. 7]

Data collecting environment: (a) Basic, (b) boxes swapped, (c) with lights off, (d) with lights off and boxes swapped

Since images constitute a significant portion of the data, various lighting conditions were used to collect data, as seen in [Fig. 7]. The colors of the boxes were also alternated to obtain diverse data. For each environment, data was collected through a total of 1,000 demonstrations, with 250 repetitions per environment, resulting in a total dataset of 1.1 million samples collected at a rate of 20 samples per second.

4.3 Policy Learning

[Fig. 8] illustrates the policy architecture utilizing a transformer-based encoder and decoder. Stereo images with a resolution of 640 x 480 are combined in pairs to create three images of size 1280 x 480, which are then input into the encoder. The term ‘pose’ is defined as the robot’s information as follows:

[Fig. 8]

Architecture of action chunking with transformer

The unit vectors i, j, k convey the orientation information of each robotic arm’s endpoint, representing the columns of the rotation matrices for the right and left arms as follows:

The training was conducted using eight GTX 4090 servers with NVIDIA’s Apex library for distributed learning, utilizing a total of 1.1 million data samples over approximately one week and 100,000 training steps. The batch size was set to 256, and the AdamW optimizer was employed as the parameter update algorithm. The policy was structured as $$ {\pi }_{\theta }\left({\mathit{a}}_{t : t+24}|{\mathit{s}}_{t-1 : t}\right)$$ where it receives the latest two environmental inputs to predict the future sequence of 24 actions. The update frequency for the robot’s target point was set to 20 Hz.

4.4 Evaluation of PTE Performance

In this paper, we conducted four types of experiments. The first experiment is a four-block sorting task to validate the effectiveness of the ACT-based imitation learning policy. The second experiment is a repetitive single-block moving task to evaluate the speed performance of PTE. The third experiment is an unseen object sorting task to assess the scalability of the policy. The fourth is a four-block sorting task to validate the feasibility of PTE-ACT.

First, to evaluate the effectiveness of the ACT-based imitation learning policy, we conducted a four-block sorting experiment. The experiment was implemented with an f value of 0 and was executed a total of 60 times. The experimental environment was set up by randomly scattering the positions and orientations of two blue blocks and two yellow blocks, as well as the location of the box, under varying lighting conditions.

The success rate for the classification task was recorded at 65%. The failure cases included 12% due to color classification errors, 18% attributed to covariate shift problems, and 5% where blocks were ejected during grasping. The covariate shift problem refers to errors occurring in state spaces that were not addressed by the collected data. For example, there were cases where two robotic arms attempted to grasp a block positioned in the center simultaneously, leading to a halt to avoid a collision, or situations where they failed to recover their posture under specific conditions.

The second experiment is a single-block sorting task to evaluate the speed performance of PTE. In case of multi-block sorting tasks, the variance in elapsed time is high due to the number of block handovers between robots, which depends on the initial block arrangement. Therefore, to objectively evaluate the speed performance of PTE, we performed a repetitive single-block sorting task. Experiments were conducted to pick up one block and place it into the nearest box for various values of f. [Fig. 9] illustrates the experimental process. For each value of f, 20 trials were conducted, and the average time taken and success rates were listed in [Table 1]. Instances where the block was dropped during the picking process or not placed inside the box were considered failures. The experimental results demonstrated that PTE improved the task speed by approximately 2 times while maintaining a success rate of 100%. Additionally, although the success rate decreased to 75%, it demonstrated a threefold improvement in task speed.

[Fig. 9]

Performance Evaluation Experimental Procedure: (a) Initial, (b) Grasp, (c) Carry, (d) Release

[Table 1]

Experiment results with various f

4.5 Discussion and Future Work

The third experiment is an unseen object sorting task to evaluate the scalability of the policy. In the data collection process, only yellow and blue 2-stud blocks were utilized. To investigate the scalability of the policy, classification experiments were conducted on previously uncollected data for blue 4-stud block, green 3-stud block, purple 2-stud block, and differently shaped yellow-green blocks [Fig. 10]. The experimental results indicated that the blue and purple blocks were classified into the blue box, while the green and yellow-green blocks were classified into the yellow box.

[Fig. 10]

Snapshots of the object classification task for unseen objects during the training process with f = 0

The final experiment is a four-block sorting task to comprehensively evaluate the performance of PTE as shwon in [Fig. 11]. PTE-ACT took only 13 seconds to sort four blocks, including block handovers. This contrasts with the 37 seconds required by ACT alone to sort four blocks, as shown in [Fig. 10].

[Fig. 11]

Snapshots of the object classification task with f = 19

For environmental perception, this study employed three stereo cameras. The wrist camera was used to observe the gripping state of the objects, while the head camera was utilized to monitor the color information of the boxes placed on both sides. Let us assume a block classification task conducted using the policy $$ {\pi }_{\theta }\left({\mathit{a}}_{t : t+k_{\mathrm{a}}}|{\mathit{s}}_{t-k_{\mathrm{s}} : t}\right)$$ with a low k_s. [Fig. 11(f)] depicts a scene where a block is handed over to the left, but when viewed in isolation, it could be interpreted as the block being received from the right. If only the wrist camera were used for the classification task, there would be no color information regarding both boxes, resulting in observations of the two robots repeatedly handing blocks to each other. In future work, we plan to conduct research on generating appropriate actions when the acquired environmental information is the same, but the situations differ.

This paper proposes the PTE algorithm for improving the task speed of robot actions generated by imitation learning policies. Upon examining the collected data, the four-block sorting task performed by the teleoperator took an average of 34 seconds. Successful cases of the autonomous block sorting task using only ACT elapsed an average of 41 seconds. The PTE-ACT reduced the time to 12 seconds, demonstrating a performance that is approximately 3 times faster than human teleoperation tasks. However, PTE showed a tendency to decrease stability in dexterous object manipulation tasks, which led to a lower success rate. Future work will focus on research to improve both speed and success rate simultaneously.