Robot imitation learning method based on diffusion model

Abstract

The invention discloses a robot imitation learning method based on a diffusion model, and belongs to the technical field of robot learning and body intelligence. Comprising the following steps: carrying out standardization processing on input image observation data and a robot state, and carrying out image enhancement by adopting gray scale transformation, random erasure and Gaussian blur; based on the enhanced image and the robot state, training a diffusion model to generate an action sequence, including forward diffusion, conditional feature fusion, noise prediction and joint loss optimization; in real-time control, features are extracted through an image enhancement network, sampling is accelerated through DDIM to generate an action sequence, noise scheduling coefficients are dynamically adjusted based on visual feature differences, and closed-loop optimization is achieved. According to the method, the image enhancement technology and the diffusion strategy are deeply fused, joint optimization of visual features and action sequences is achieved, the action generation accuracy and robustness of the robot under the complex visual interference and small sample conditions are remarkably improved, and the success rate is improved by 18% compared with a base line under 40 sample sizes.

Description

Technical Field

[0001] This invention belongs to the field of robot learning and embodied intelligence technology, specifically relating to a robot motion control method that deeply integrates image enhancement technology with diffusion policy, applicable to robot manipulation tasks under complex visual conditions with few samples. Background Technology

[0002] In robot vision motion control, the quality of visual input directly affects the accuracy of motion generation, but existing technologies have the following pain points: (1) Complex visual interference: There are problems such as drastic changes in lighting, local occlusion and motion blur in industrial environments, which cause the images captured by the visual sensors to have local feature loss or distortion, and traditional diffusion models have difficulty extracting robust motion control features from them; (2) Small sample learning dilemma: High-quality robot operation data collection is expensive, and the performance of existing methods drops sharply under the condition of few samples (<100 sets of demonstration data); (3) Module fragmentation problem: In traditional schemes, image enhancement algorithms and motion generation (such as imitation learning) are implemented as independent modules in series, which leads to the enhancement process ignoring the downstream motion generation requirements, which may over-enhance irrelevant features, and the diffusion model cannot back-optimize the image enhancement network parameters, resulting in a deterioration of the model effect;

[0003] While diffusion models exhibit multimodal advantages in action generation, they rely on high-quality visual input. In existing technologies, image enhancement and action generation are typically treated as independent modules, lacking end-to-end optimization, leading to a significant increase in action generation errors under low-quality visual environments. Therefore, deeply integrating image enhancement techniques with diffusion strategies to achieve joint optimization of visual features and action sequences is crucial for improving the robot's maneuverability in complex environments. Summary of the Invention

[0004] This invention proposes a diffusion strategy robot control method that integrates image enhancement technology. By constructing an "image enhancement-conditional diffusion" joint optimization framework, it achieves the co-evolution of visual feature quality improvement and action sequence generation.

[0005] This invention is achieved through the following technical solutions.

[0006] The robot imitation learning method based on a diffusion model described in this invention includes the following steps:

[0007] Step 1: Data Preprocessing and Augmentation

[0008] Load the dataset containing image observation data, robot state, and expert demonstration action sequences; standardize the image data and robot state; and process low-quality input images using image enhancement techniques, including grayscale transformation, random erasure, and Gaussian blur, as detailed below:

[0009] (1) Grayscale transformation; Image processing that enhances data diversity by randomly adjusting image contrast, hue, and saturation, allowing the model to adapt to different lighting conditions (sunny day, cloudy day, artificial light), and avoiding overfitting caused by the uniform brightness of the training data. The original image is transformed using the transformation function T. f(x,y) (in f(x,y) Indicates the original image in coordinates (x,y) Pixel value at that location, x, y (These are the horizontal and vertical coordinates of a pixel in the image, used to uniquely identify the pixel's location) are converted into an enhanced image (representing the enhanced image after grayscale transformation in coordinates). (x,y) (Pixel value at location), the formula is:

[0010]

[0011] in This is a grayscale transformation function used to adjust the contrast, hue, or saturation of the original pixel values.

[0012] (2) Random Erasure: Randomly select a rectangular area in the image and replace its pixel values ​​with fixed values ​​(such as black, gray, or the image mean) or delete it directly (leaving blanks) to simulate a scenario of partial occlusion or damage to the image. In the original image I (Its coordinates) (x,y) The pixel value at that location is I(x,y) Randomly generate rectangular areas on ) Ω (Based on the coordinates that satisfy the region boundary) (x,y) (compose), and fill the pixel values ​​in that area with v (This represents the fixed pixel value filled into the rectangular area Ω, which can be black, gray, or the pixel mean of the original image I, etc.), to obtain the enhanced image. The formula is:

[0013]

[0014] in This indicates the enhanced image after random erasure in coordinates. The pixel value at that location.

[0015] (3) Gaussian Blur: Gaussian blur uses a Gaussian kernel to convolve an image, smoothing pixel values ​​to reduce high-frequency details and simulate the effect of blurry focus or motion blur. It uses the standard deviation. (Its value determines the degree of fuzziness,) The larger the Gaussian kernel, the stronger the blurring effect. Convolution operation with the original image I (representing a two-dimensional image matrix composed of pixel values) This represents a convolution operation, used to apply a Gaussian kernel to the original image I. (Perform pixel-by-pixel convolution calculations) to obtain a blurred image. The formula is:

[0016]

[0017] Among them, Gaussian kernel The mathematical expression is:

[0018]

[0019] here x, y This represents the relative coordinates of the pixels within the Gaussian kernel (indicating the position of the pixels within the kernel relative to the kernel center). This is the blurred image obtained after convolution.

[0020] Step 2: Diffusion Strategy Training

[0021] Based on the enhanced image and robot state, a diffusion model is trained to generate action sequences, specifically including:

[0022] (1) Forward diffusion process (noise addition): in the action sequence a 0 Gaussian noise is gradually added to obtain a noisy action sequence. a k The formula is:

[0023]

[0024] illustrate: a 0 This is the initial action sequence, i.e., the original action sequence without noise; a k For the first k The noisy action sequence of the step is a 0 Add to the base k The sequence obtained after step Gaussian noise; a k For the first k The noise figure of the step is used to control the proportion of noise added, and its value ranges from [0,1]. To follow a standard normal distribution ( I Gaussian noise (with the identity matrix) is the real noise added during the forward diffusion process.

[0025] (2) Conditional feature fusion: Enhanced image features and robot state features are fused by cross-attention to obtain conditional features. f cond The formula is:

[0026]

[0027] illustrate: f cond The fused conditional features are obtained by fusing enhanced image features and robot state features through cross-attention; CrossAttn is the cross-attention mechanism, which is used to achieve cross-modal feature fusion. The enhanced image at time t is obtained by image enhancement processing from the original observation image; To enhance the image Visual features extracted using a ResNet network; q t Let be the robot's quaternion state at time t, used to describe the robot's posture; SE3-Encoder( q t ) represents the quaternion state of the robot q t The state features extracted by the SE3 encoder are used; SE3 is an encoding network for processing 3D pose.

[0028] (3) Noise prediction; predicting noise residuals using the UNet network. The input is a noisy action sequence. a k Time step k and conditional features f cond The formula is:

[0029]

[0030] illustrate: The predicted values ​​are the noise residuals, generated by the UNet network from the noise action sequence. Time step k and conditional feature f cond The prediction was obtained; This is a network model used to predict noise residuals; a k For the first k The noise action sequence of the step (same as the forward diffusion process) a k ); k This is the time step, i.e., the index of the current step in the diffusion process; f cond For the fused conditional features (in the fusion of conditional features) f cond ).

[0031] (4) Training loss: Calculate the mean square error between predicted noise and actual noise, and combine it with the visual feature consistency loss. The formula is as follows:

[0032]

[0033] Note: L is the total training loss, which consists of the mean square error between predicted noise and real noise and the visual feature consistency loss; True Gaussian noise (same as in the forward diffusion process) ); For the predicted noise residual (same as in noise prediction) ); It is the L2 norm, used to calculate the mean square error between the predicted noise and the actual noise; γ The weighting coefficients for visual feature consistency loss are used to balance the proportions of the two losses. ô t The enhanced image at time t; o t The image is the original observation image acquired at time t (without image enhancement processing); To enhance the image ô t Visual features extracted using the VGG network; Original observation image o t Visual features extracted using the VGG network; It is the L1 norm, used to calculate the visual feature difference between the enhanced image and the original image (i.e., visual consistency loss).

[0034] Step 3: Real-time control process

[0035] In real-time robot control, the following steps are performed:

[0036] (1) Observation phase: Acquiring current images o t and robot state s t And perform image enhancement to obtain the enhanced image. ô t :

[0037]

[0038] Note: o t The image is the original image captured at time t; U-Net is the network model used for image enhancement. ô t The image at time t is the enhanced image obtained after image enhancement.

[0039] (2) Feature encoding: extracting visual featuresf v and state characteristics f s :

[0040]

[0041]

[0042] illustrate: f v Visual features, derived from enhanced images ô t Extracted using a ResNet-18 network; ô t The enhanced image at time t; f s State characteristics are determined by the robot's quaternion states. q t Extracted using the SE3 encoder; q t Let t be the robot's quaternion state at time t, used to describe the robot's posture.

[0043] (3) Action generation: Generate action sequences for the next H steps by accelerating sampling using DDIM:

[0044]

[0045] illustrate: a t+H The sequence of actions for the next H steps is the sequence of actions the robot needs to execute; DDIM-Sampler is the DDIM accelerator sampler used to generate the action sequence. f v For visual features (in the same feature encoding) f v ); f s For state features (same as f in feature encoding) s ); z init The initial noise is the random initialization noise of the sampling process, which follows a standard normal distribution.

[0046] (4) Closed-loop optimization: Calculate the feature difference after every two steps and dynamically adjust the noise scheduling coefficient:

[0047]

[0048]

[0049] illustrate: Δf Visual feature differences are used to predict visual features. Visual features observed in actual observation The L2 norm; The noise scheduling coefficient is determined by differences in visual features. Δf After passing through the Softmax function (including weights) W and bias b This is obtained through dynamic adjustment; W and b These are the weights and bias parameters of the Softmax layer; Softmax is a normalization function used to map feature differences to noise scheduling coefficients.

[0050] The technical advantages of this invention are as follows: through end-to-end joint optimization of image enhancement and diffusion strategies, the accuracy and robustness of robot action generation under complex visual interference and small sample conditions are significantly improved. Experiments show that, with a sample size of 40, the success rate of the method combined with random erasure enhancement is 18% higher than the baseline. Attached Figure Description

[0051] Figure 1 The loss of the EDP model was validated for both standard and small sample sizes.

[0052] Figure 2 The training loss for the EDP model is given for both standard and few samples.

[0053] Figure 3 shows the validation loss of the improved EDP model under standard samples.

[0054] Figure 4 To improve the validation loss of the EDP model with a small number of samples.

[0055] Figure 5 This is a flowchart of the present invention. Detailed Implementation

[0056] This invention will be further explained with reference to the accompanying drawings and the model training process.

[0057] The overall process of this invention is as follows: Figure 3 As shown, the specific steps are as follows:

[0058] 1.1 Data Preprocessing and Augmentation

[0059] • Load the dataset: The input includes image observation data (such as RGB or depth maps), robot states (joint angles, end-effector pose), and expert-demonstrated action sequences. (See Table 1)

[0060] Table 1 Dataset Format

[0061]

[0062] • Calculate normalization parameters: Standardize the image data (pixel values) and robot states (such as joint angle ranges).

[0063] • Image enhancement: An enhancement network with a U-Net structure is used to denoise and improve the contrast of low-quality input images, enhance the diversity of image data distribution, and output high-quality image features.

[0064] 2.2 Feature Extraction and Diffusion Strategy Training

[0065] • Observation feature extraction:

[0066] o Image data → Extract multi-scale visual features using CNN (ResNet).

[0067] The robot's state is converted into a 6D rotational representation using an isotropic encoder, enhancing its kinematic generalization ability.

[0068] • Forward diffusion (noise addition): Gaussian noise is gradually added to the action sequence to construct the diffusion process.

[0069] • UNet prediction noise residuals:

[0070] Input: Noisy action sequence + enhanced visual features + robot state features.

[0071] Output: Predicted noise residuals, used for inverse denoising.

[0072] • Loss calculation and backpropagation:

[0073] The difference between the predicted noise and the actual noise is calculated using the mean square error (MSE).

[0074] o Jointly optimize the image enhancement network and diffusion model to ensure consistency between enhanced features and action generation.

[0075] • EMA (Exponential Moving Average) model parameter updates: stabilize the training process and avoid model oscillations.

[0076] 3.3 Closed-loop control and real-time optimization

[0077] • Observation data: The robot collects environmental images and its own status in real time.

[0078] • Dynamic noise scheduling: Based on the feature differences between the enhanced image and the original image, the noise intensity of the diffusion process is adjusted to adapt to dynamic environmental changes.

[0079] • Action sequence generation: Accelerate sampling through DDIM (Denoising Diffusion Implicit Models) to generate robot actions for the next N steps.

[0080] • Execution and feedback: After each K steps, the environment is re-observed and subsequent actions are optimized to form a closed-loop control.

[0081] 3.3 Data Analysis (as shown in Table 2)

[0082] Table 2 Success rates of different image enhancement algorithms

[0083]

[0084] 1. Performance benchmark of the base model (EDP)

[0085] The current algorithm performs well with a success rate of 74% on 100 samples, but the success rate drops sharply to 34% when the sample size decreases to 40. This indicates that:

[0086] (1) The model is severely overfitted under small sample size: the limited training data cannot support effective generalization ability.

[0087] (2) Insufficient diversity of data distribution: Even with an increase in the number of samples, the model struggles to learn new knowledge patterns.

[0088] Therefore, data augmentation techniques are used to create diverse datasets and simulate a wider range of data distributions, thereby effectively improving the robustness and generalization ability of the model under small sample conditions.

[0089] 2.2 Impact of Data Augmentation Methods on Model Performance

[0090] (1) EDP + Gaussian blur

[0091] Enhanced parameters: A combined strategy of "color adjustment + Gaussian blur" is adopted. Color adjustment simulates a cloudy day effect by reducing brightness (0.6-0.8), contrast (0.7-0.9), saturation (0.5-0.7), and slightly shifting the cool tone (hue±0.05). Gaussian blur uses a 5×5 kernel, with the blur intensity varying in the range of sigma=0.1-2.0 to simulate the inaccurate focus effect of motion blur.

[0092] Performance: With a sample size of 40, the success rate was 0.42, which is 0.08 higher than the EDP model with the same sample size (0.34), indicating that the enhancement strategy (simulating cloudy weather + motion blur interference) helps to improve the stability of the model in small sample scenarios.

[0093] (2) EDP + Random Erasure

[0094] Enhanced parameters: A combination strategy of "color adjustment + random erasure" is adopted. The color adjustment parameters are the same as above (simulating a cloudy day); the probability of random erasure is 0.5, the erased area occupies 0.01-0.02 of the original image (small area), the aspect ratio is 0.8-1.2 (close to a square), and the erased area is filled with black (value=0) to simulate the scene of partial occlusion or damage to the image.

[0095] Performance: With a sample size of 40, the success rate was 0.52, an improvement of 0.18 compared to the EDP model with the same sample size. This result indicates that, under the current parameter settings (simulating occlusion + cloudy weather effect), the performance improvement effect of the Gaussian blur strategy differs from that of the Gaussian blur strategy, and is even 0.1 higher than that of the Gaussian blur strategy.

[0096] In summary, reducing the sample size significantly reduces the success rate of the EDP model; and different data augmentation methods have different impacts on performance, with random erasure (simulating occlusion + cloudy weather) showing a more obvious positive effect with small samples.

Claims

1. A robot imitation learning method based on a diffusion model, characterized in that, Includes the following steps: Step 1: Data Preprocessing and Enhancement: Standardize the input image observation data and robot state, and use image enhancement techniques to process low-quality images; Step 2: Diffusion strategy training: Based on the enhanced image and robot state, train the diffusion model to generate action sequences; Step 3: Real-time control process: During the robot's operation, images and status are acquired in real time, action sequences are generated and executed, and closed-loop optimization is performed.

2. The method according to claim 1, characterized in that, The image enhancement technology includes at least one of the following: Grayscale transformation adjusts image contrast, hue, or saturation. Random erasure involves randomly selecting areas in the image and replacing them with fixed pixel values. Gaussian blur uses a Gaussian kernel to convolve an image to simulate a blurring effect.

3. The method according to claim 1, characterized in that, The image enhancement network adopts a U-Net structure to improve the diversity of image data distribution.

4. The method according to claim 1, characterized in that, The diffusion strategy training includes: The forward diffusion process involves gradually adding Gaussian noise to the action sequence; Conditional feature fusion integrates enhanced image features with robot state features through a cross-attention mechanism; Noise prediction: Predicting noise residuals using the UNet network; The training loss is calculated, including the mean square error between predicted noise and real noise, and the visual feature consistency loss between the enhanced image and the original image.

5. The method according to claim 4, characterized in that, The formula for conditional feature fusion is: in, The enhanced image at time t, To enhance the visual features extracted from images via the ResNet network, q t The state is the robot's quaternion.

6. The method according to claim 4, characterized in that, The formula for the training loss is: in, This is real noise. For predicting noise, γ is the weighting coefficient, ô t To enhance the image, o t The original observation image, To extract visual features using a VGG network, Visual features extracted using the VGG network.

7. The method according to claim 1, characterized in that, The real-time control process includes: (1) Observation phase: Acquire the current image and status, and obtain the enhanced image through the image enhancement network; (2) Feature encoding: Extracting visual features and state features; (3) Action generation: Generate future multi-step action sequences by accelerating sampling through DDIM; (4) Closed-loop optimization: The noise scheduling coefficient is dynamically adjusted based on the difference between the predicted and observed visual features.

8. The method according to claim 7, characterized in that, The noise dispatch coefficient is calculated using the following formula: Where Δf represents the visual feature difference. To predict visual features, For the actual observed visual features, W and b are the weights and bias parameters of the Softmax layer.

Images